Update app.py

app.py

@@ -2,7 +2,6 @@ import streamlit as st
 import os
 import zipfile
 import csv
-import base64  # 追加
 from PIL import Image, ImageFilter
 from gradio_client import Client
 import datetime
@@ -12,20 +11,15 @@ client_nsfw = Client("https://ozoneasai-falconsai-nsfw-image-detection.hf.space/
 client_other = Client("https://ozoneasai-sanali209-imclasif-quality-v001.hf.space/--replicas/m6xqz/")
 
 # Initialize pagination
-if "page" not in st.session_state:
+if 'page' not in st.session_state:
     st.session_state.page = 1
 
-# Initialize blur option
-if "blur_option" not in st.session_state:
-    st.session_state.blur_option = True
-
-# Create a new checkbox
-blur_option = st.checkbox("NSFW画像にBlurをかける", value=st.session_state.blur_option)
-
-# ... (rest of the code remains unchanged)
+# Initialize blur option as a toggle
+if 'blur_option' not in st.session_state:
+    st.session_state.blur_option = st.checkbox("NSFW画像にBlurをかける")
 
 # Initialize sort options
-if "sort_option" not in st.session_state:
+if 'sort_option' not in st.session_state:
     st.session_state.sort_option = {"column": "timestamp", "ascending": True}
 
 # Define the function to generate a download link for a file
@@ -40,28 +34,21 @@ st.title("ファイルアップロードと共有")
 
 uploaded_files = st.file_uploader("ファイルをアップロードしてください", type=["jpg", "png"], accept_multiple_files=True)
 
-# List to store Gradio predictions
+# Lists to store Gradio predictions, good values, timestamps, and file paths
 gradio_nsfw_predictions = []
 gradio_other_predictions = []
-
-# List to store good values
 good_values = []
-
-# List to store timestamps
 timestamps = []
-
-# List to store file paths
 file_paths = []
 
 def get_gradio_nsfw_prediction(image_path):
     result = client_nsfw.predict(image_path, api_name="/predict")
-    return result["output"]  # 追加
+    return result
 
 def get_gradio_other_prediction(image_path):
     result = client_other.predict(image_path, api_name="/predict")
-    return result["output"]  # 追加
+    return result
 
-# Function to save uploaded files and get Gradio predictions
 def save_uploaded_files(uploaded_files):
     if not os.path.exists("temp"):
         os.makedirs("temp")
@@ -80,8 +67,8 @@ def save_uploaded_files(uploaded_files):
         
         # Get good value
         good_value = None
-        if "confidences" in gradio_other_prediction:
-            good_value = gradio_other_prediction["confidences"][0]["confidence"]
+        if "good" in gradio_other_prediction and "confidences" in gradio_other_prediction["good"]:
+            good_value = gradio_other_prediction["good"]["confidences"][0]["confidence"]
         good_values.append(good_value)
         
         # Get timestamp
@@ -90,20 +77,17 @@ def save_uploaded_files(uploaded_files):
         
         file_paths.append(file_path)
 
-# 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
 
-# Paginate files
 def paginate_files(files, page, files_per_page):
     start_index = (page - 1) * files_per_page
     end_index = start_index + files_per_page
     return files[start_index:end_index]
 
-# Display images with pagination
-def display_images(images, blur_option):
+def display_images(images):
     for i, file_path in enumerate(images):
         file_ext = os.path.splitext(file_path)[1].lower()
 
@@ -113,91 +97,89 @@ def display_images(images, blur_option):
         good_value = good_values[i]
         timestamp = timestamps[i]
 
-        st.write(f"**ファイル名:** {os.path.basename(file_path)}")
-        st.write(f"**NSFWの予測:** {gradio_nsfw_prediction}")
-        st.write(f"**Otherの予測:** {gradio_other_prediction}")
-        st.write(f"**Goodの値:** {good_value}")
-        st.write(f"**作成時間:** {timestamp}")
+        st.write(f"**Prediction for {os.path.basename(file_path)} (NSFW):** {gradio_nsfw_prediction}")
+        st.write(f"**Prediction for {os.path.basename(file_path)} (Other):** {gradio_other_prediction}")
+        st.write(f"**Good value:** {good_value}")
+        st.write(f"**Timestamp:** {timestamp}")
 
-        # Apply blur if NSFW prediction is above a threshold and blur option is enabled
-        if blur_option and gradio_nsfw_prediction == "nsfw" and good_value is not None and good_value >= 0.8:
-            st.image(apply_blur(file_path), caption="Blur applied", use_column_width=True)
+        if st.session_state.blur_option and gradio_nsfw_prediction.lower() == "nsfw":
+            blurred_img = apply_blur(file_path)
+            st.image(blurred_img, caption=os.path.basename(file_path), use_column_width=True)
         else:
-            st.image(file_path, caption="Original", use_column_width=True)
+            if file_ext in [".jpg", ".png"]:
+                st.image(file_path, caption=os.path.basename(file_path), use_column_width=True)
+            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)
+            with col2:
+                if col2.button("削除", key=file_path):
+                    os.remove(file_path)
+                    st.experimental_rerun()
 
-# Save uploaded files and get Gradio predictions
 if uploaded_files:
     save_uploaded_files(uploaded_files)
 
-# Display pagination controls
-if file_paths:
-    file_paths.sort(key=lambda x: os.path.getmtime(x), reverse=True)
-    files_per_page = 10
-    num_pages = (len(file_paths) - 1) // files_per_page + 1
-    page = st.session_state.get("page", 1)
+files_per_page = 20
+num_pages = (len(file_paths) - 1) // files_per_page + 1
 
-    # Next page button
+col1, col2 = st.columns(2)
+with col1:
+    if st.button("前へ", key="prev_page"):
+        st.session_state.page = max(1, st.session_state.page - 1)
+with col2:
     if st.button("次へ", key="next_page"):
         st.session_state.page = min(num_pages, st.session_state.page + 1)
 
-    # Previous page button
-    if st.button("前へ", key="prev_page"):
-        st.session_state.page = max(1, st.session_state.page - 1)
+selected_page = st.number_input("移動するページを指定してください", min_value=1, max_value=num_pages, value=st.session_state.page, key="selected_page")
+if selected_page != st.session_state.page:
+    st.session_state.page = selected_page
+
+st.write(f"現在のページ: {st.session_state.page}")
+
+current_page_files = paginate_files(file_paths, st.session_state.page, files_per_page)
+
+sort_column = st.selectbox("並び替えオプション", ["timestamp", "good"])
+ascending = st.checkbox("昇順", value=st.session_state.sort_option["ascending"])
+st.session_state.sort_option = {"column": sort_column, "ascending": ascending}
+
+if sort_column == "timestamp":
+    current_page_files.sort(key=lambda x: timestamps[file_paths.index(x)], reverse=not ascending)
+elif sort_column == "good":
+    current_page_files.sort(key=lambda x: good_values[file_paths.index(x)], reverse=not ascending)
+
+display_images(current_page_files)
+
+# Blurのトグルを更新
+st.session_state.blur_option = st.checkbox("NSFW画像にBlurをかける", value=st.session_state.blur_option)
+
+# CSVに保存ボタン
+if st.button("CSVに保存"):
+    csv_filename = "gradio_predictions_and_good_values.csv"
+    with open(csv_filename, mode='w', newline='', encoding='utf-8') as csv_file:
+        csv_writer = csv.writer(csv_file)
+        csv_writer.writerow(['ファイル名', 'NSFWの予測', 'Otherの予測', 'Goodの値', '作成時間'])
+        for i, file_path in enumerate(file_paths):
+            gradio_nsfw_prediction = gradio_nsfw_predictions[i]
+            gradio_other_prediction = gradio_other_predictions[i]
+            good_value = good_values[i]
+            timestamp = timestamps[i]
+            csv_writer.writerow([os.path.basename(file_path), gradio_nsfw_prediction, gradio_other_prediction, good_value, timestamp])
+
+    st.markdown(get_binary_file_downloader_html(csv_filename, label="CSVをダウンロード"), unsafe_allow_html=True)
+
+with open("gradio_predictions.txt", "w") as f:
+    for prediction in gradio_nsfw_predictions:
+        f.write(f"{prediction[0]} (NSFW): {prediction[1]}\n")
+    for prediction in gradio_other_predictions:
+        f.write(f"{prediction[0]} (Other): {prediction[1]}\n")
+
+if st.button("ファイルを一括ダウンロード"):
+    zip_filename = "files.zip"
+    with zipfile.ZipFile(zip_filename, "w") as zipf:
+        for file_path in file_paths:
+            zipf.write(file_path, os.path.basename(file_path))
 
-    # Page number selection
-    selected_page = st.number_input("移動するページを指定してください", min_value=1, max_value=num_pages, value=st.session_state.page, key="selected_page")
-    if selected_page != st.session_state.page:
-        st.session_state.page = selected_page
-
-    st.write(f"現在のページ: {st.session_state.page}")
-
-    # Display images for the current page
-    current_page_files = paginate_files(file_paths, st.session_state.page, files_per_page)
-
-    # Sort files based on the selected option
-    sort_column = st.selectbox("並び替えオプション", ["timestamp", "good"])
-    ascending = st.checkbox("昇順", value=st.session_state.sort_option["ascending"])
-    st.session_state.sort_option = {"column": sort_column, "ascending": ascending}
-
-    # Sort files based on the selected option
-    if sort_column == "timestamp":
-        current_page_files.sort(key=lambda x: timestamps[file_paths.index(x)], reverse=not ascending)
-    elif sort_column == "good":
-        current_page_files.sort(key=lambda x: good_values[file_paths.index(x)], reverse=not ascending)
-
-    display_images(current_page_files, st.session_state.blur_option_checkbox)
-
-    # Save Gradio predictions, good values, and timestamps to CSV
-    if st.button("CSVに保存"):
-        csv_filename = "gradio_predictions_and_good_values.csv"
-        with open(csv_filename, mode='w', newline='', encoding='utf-8') as csv_file:
-            csv_writer = csv.writer(csv_file)
-            csv_writer.writerow(['ファイル名', 'NSFWの予測', 'Otherの予測', 'Goodの値', '作成時間'])
-            for i, file_path in enumerate(file_paths):
-                gradio_nsfw_prediction = gradio_nsfw_predictions[i]
-                gradio_other_prediction = gradio_other_predictions[i]
-                good_value = good_values[i]
-                timestamp = timestamps[i]
-                csv_writer.writerow([os.path.basename(file_path), gradio_nsfw_prediction, gradio_other_prediction, good_value, timestamp])
-
-        # CSVファイルをリンクとして提供
-        st.markdown(get_binary_file_downloader_html(csv_filename, label="CSVをダウンロード"), unsafe_allow_html=True)
-
-    # Gradio予測データを保存
-    with open("gradio_predictions.txt", "w") as f:
-        for prediction in gradio_nsfw_predictions:
-            f.write(f"{prediction[0]} (NSFW): {prediction[1]}\n")
-        for prediction in gradio_other_predictions:
-            f.write(f"{prediction[0]} (Other): {prediction[1]}\n")
-
-    # Zipファイル作成とダウンロード
-    if st.button("ファイルを一括ダウンロード"):
-        zip_filename = "files.zip"
-        with zipfile.ZipFile(zip_filename, "w") as zipf:
-            for file_path in file_paths:
-                zipf.write(file_path, os.path.basename(file_path))
-
-        # Zipファイルをリンクとして提供
-        st.markdown(get_binary_file_downloader_html(zip_filename, label="Zipファイルをダウンロード"), unsafe_allow_html=True)
+    st.markdown(get_binary_file_downloader_html(zip_filename, label="Zipファイルをダウンロード"), unsafe_allow_html=True)
 
 st.success("ファイルがアップロードされました。このページのURLを他のクライアントと共有してください。")