Update app.py

app.py

@@ -1,3 +1,86 @@
+print("[1/7] Importing streamlit...", end=" ", flush=True)
+
+import streamlit as st
+print("done")
+print("[2/7] Importing something...", end=" ", flush=True)
+import os
+import pandas as pd
+from datetime import datetime
+from PIL import Image
+from PIL import ImageFilter
+from io import BytesIO
+import zipfile
+import base64
+print("done")
+print("[3/7] Importing deepdanbooru...", end=" ", flush=True)
+
+import deepdanbooru as dd
+print("done")
+print("[4/7] Importing huggingface_hub...", end=" ", flush=True)
+
+import huggingface_hub
+print("done")
+print("[5/7] Importing tensorflow...", end=" ", flush=True)
+
+import tensorflow as tf
+print("done")
+print("[6/7] Importing numpy...", end=" ", flush=True)
+
+import numpy as np
+print("done")
+print("[6/7] Importing transformers...", end=" ", flush=True)
+
+from transformers import pipeline
+print("done")
+
+# ページごとの表示数
+PAGE_SIZE = 20
+
+# ファイルの保存先フォルダー
+photos_folder = "photos"
+
+# インデックスファイルのパス
+index_file_path = "index.csv"
+
+# タグ付け関数
+def predict_tags(image: Image.Image, score_threshold: float) -> tuple[dict[str, float], dict[str, float], str]:
+    _, height, width, _ = model.input_shape
+    image = np.asarray(image)
+    image = tf.image.resize(image, size=(height, width), method=tf.image.ResizeMethod.AREA, preserve_aspect_ratio=True)
+    image = image.numpy()
+    image = dd.image.transform_and_pad_image(image, width, height)
+    image = image / 255.
+    probs = model.predict(image[None, ...])[0]
+    probs = probs.astype(float)
+
+    indices = np.argsort(probs)[::-1]
+    result_all = dict()
+    result_threshold = dict()
+    for index in indices:
+        label = labels[index]
+        prob = probs[index]
+        result_all[label] = prob
+        if prob < score_threshold:
+            break
+        result_threshold[label] = prob
+    result_text = ', '.join(result_all.keys())
+    return result_threshold, result_all, result_text
+
+# NSFW 判定関数
+def predict_nsfw(image: Image.Image) -> dict[str, float]:
+    classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+    result = classifier(image)[0]
+    return {result['label']: result['score']}
+
+# ブラーをかける関数
+def blur_image(image: Image.Image, blur_enabled: bool, nsfw_score: float) -> Image.Image:
+    if blur_enabled:
+        nsfw_result = predict_nsfw(image)
+        nsfw_score = nsfw_result.get('nsfw', 0.0)
+        if nsfw_score >= 0.75:
+            image = image.filter(ImageFilter.BLUR)
+    return image
+
 # blur_toggle_keyの初期値を設定
 blur_toggle_key = "blur_toggle_unique_key_for_sidebar"
 
@@ -21,6 +104,54 @@ def save_uploaded_photo(uploaded_photo, file_name):
 
     image.save(os.path.join(photos_folder, file_name), "PNG")
 
+# ページングと並び替えのためのデータを取得する関数
+def load_data():
+    if os.path.exists(index_file_path):
+        return pd.read_csv(index_file_path)
+    else:
+        return pd.DataFrame(columns=["File Name", "Timestamp", "Tags"])
+
+# アップロードされた写真を表示する関数
+def display_photos(photos):
+    for photo_info in photos.iterrows():
+        row = photo_info[1]
+        photo_path = os.path.join(photos_folder, row["File Name"])
+        image = Image.open(photo_path)
+
+        # タグを予測して表示
+        result_threshold, result_all, result_text = predict_tags(image, 0.7)
+
+        st.image(image, caption=row["File Name"], use_column_width=True)
+        
+        # タグを表示
+        st.write("タグ:", ", ".join(result_all.keys()))
+
+# フォルダーの中身を zip ファイルとしてダウンロード
+def download_photos_as_zip(file_paths):
+    # Zip ファイル作成
+    with BytesIO() as zip_buffer:
+        with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED) as zip_file:
+            for file_path in file_paths:
+                zip_file.write(file_path, os.path.basename(file_path))
+        
+        # ダウンロードリンク表示
+        st.markdown(
+            f"**[ダウンロード ZIPファイル](data:application/zip;base64,{base64.b64encode(zip_buffer.getvalue()).decode()})**",
+            unsafe_allow_html=True
+        )
+
+# モデルとラベルをダウンロードする関数
+def load_model():
+    path = huggingface_hub.hf_hub_download('public-data/DeepDanbooru', 'model-resnet_custom_v3.h5')
+    model = tf.keras.models.load_model(path)
+    return model
+
+def load_labels():
+    path = huggingface_hub.hf_hub_download('public-data/DeepDanbooru', 'tags.txt')
+    with open(path) as f:
+        labels = [line.strip() for line in f.readlines()]
+    return labels
+
 # Streamlit アプリケーションのメイン部分
 def main():
     st.sidebar.title("アップロードオプション")
@@ -86,6 +217,6 @@ def main():
     if st.button("写真をダウンロード (ZIP)"):
         file_paths = [os.path.join(photos_folder, file) for file in os.listdir(photos_folder)]
         download_photos_as_zip(file_paths)
-
+print("all function defined.")
 if __name__ == "__main__":
     main()