Update app.py

app.py

@@ -2,11 +2,40 @@ import streamlit as st
 import os
 import zipfile
 import base64
+from PIL import Image
+from transformers import pipeline
+
+# Function to save NSFW data to a CSV file
+def save_nsfw_data_to_csv(data):
+    import csv
+
+    with open("nsfw_data.csv", mode="w", newline="") as file:
+        writer = csv.writer(file)
+        writer.writerow(["Image_Path", "NSFW_Score"])
+        for row in data:
+            writer.writerow(row)
+
+# Function to apply blur to an image
+def apply_blur(image_path):
+    img = Image.open(image_path)
+    img = img.filter(ImageFilter.GaussianBlur(radius=5))
+    return img
+
+# Function to get NSFW score using the provided pipeline
+def get_nsfw_score(image_path):
+    img = Image.open(image_path)
+    classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+    result = classifier(img)
+    nsfw_score = result[0]["score"]
+    return nsfw_score
 
 st.title("ファイルアップロードと共有")
 
 uploaded_files = st.file_uploader("ファイルをアップロードしてください（複数選択可）", type=["jpg", "png", "pdf"], accept_multiple_files=True)
 
+# List to store NSFW data for CSV
+nsfw_data = []
+
 def save_uploaded_files(uploaded_files):
     if not os.path.exists("temp"):
         os.makedirs("temp")
@@ -16,14 +45,12 @@ def save_uploaded_files(uploaded_files):
         with open(file_path, "wb") as f:
             f.write(uploaded_file.getbuffer())
         file_paths.append(file_path)
-    return file_paths
 
-# ダウンロード用のHTMLを生成
-def get_binary_file_downloader_html(file_path):
-    with open(file_path, 'rb') as f:
-        data = f.read()
-    b64 = base64.b64encode(data).decode()
-    return f'<a href="data:application/zip;base64,{b64}" download="{file_path}">ダウンロード</a>'
+        # Get NSFW score
+        nsfw_score = get_nsfw_score(file_path)
+        nsfw_data.append((file_path, nsfw_score))
+
+    return file_paths
 
 if uploaded_files:
     file_paths = save_uploaded_files(uploaded_files)
@@ -50,21 +77,29 @@ if file_paths:
 
     for file_path in file_paths[start_index:end_index]:
         file_ext = os.path.splitext(file_path)[1].lower()
-        if file_ext in [".jpg", ".png"]:
-            st.image(file_path, caption=os.path.basename(file_path), use_column_width=True)
-        elif file_ext == ".pdf":
-            st.write(os.path.basename(file_path))
-        col1, col2 = st.columns([4, 1])
-        with col1:
+
+        nsfw_score = get_nsfw_score(file_path)
+
+        if nsfw_score >= 0.8 and not st.session_state.display_options:
+            # Apply blur if NSFW score is high and display options are hidden
+            blurred_img = apply_blur(file_path)
+            st.image(blurred_img, caption=os.path.basename(file_path), use_column_width=True)
+        else:
             if file_ext in [".jpg", ".png"]:
                 st.image(file_path, caption=os.path.basename(file_path), use_column_width=True)
-            else:
+            elif file_ext == ".pdf":
                 st.write(os.path.basename(file_path))
-                st.write(f"URL: {file_path}")
-        with col2:
-            if col2.button("削除", key=file_path):
-                os.remove(file_path)
-                st.experimental_rerun()
+            col1, col2 = st.columns([4, 1])
+            with col1:
+                if file_ext in [".jpg", ".png"]:
+                    st.image(file_path, caption=os.path.basename(file_path), use_column_width=True)
+                else:
+                    st.write(os.path.basename(file_path))
+                    st.write(f"URL: {file_path}")
+            with col2:
+                if col2.button("削除", key=file_path):
+                    os.remove(file_path)
+                    st.experimental_rerun()
 
     # 一括ダウンロードボタンとtempフォルダーをZipファイルに作成
     if st.button("フォルダーを一括ダウンロード"):
@@ -79,4 +114,12 @@ if file_paths:
         # Zipファイルをリンクとして提供
         st.markdown(get_binary_file_downloader_html(zip_filename), unsafe_allow_html=True)
 
+    # トグルボタンを表示
+    st.checkbox("表示オプション", key="display_options")
+
+    if st.session_state.display_options:
+        # NSFWデータをCSVに保存
+        save_nsfw_data_to_csv(nsfw_data)
+        st.success("NSFWデータがCSVファイルに保存されました。")
+
     st.success("ファイルがアップロードされました。このページのURLを他のクライアントと共有してください。")