Update app.py

app.py

@@ -3,10 +3,12 @@ import os
 import zipfile
 import csv
 from PIL import Image, ImageFilter
-from gradio_client import Client
 import base64
 import datetime
 
+if 'blur_option' not in st.session_state:
+    st.session_state.blur_option = True  
+
 # Initialize Gradio client
 client_nsfw = Client("https://ozoneasai-falconsai-nsfw-image-detection.hf.space/--replicas/0fyhj/")
 
@@ -16,7 +18,11 @@ if 'page' not in st.session_state:
 
 # Initialize blur option as a toggle
 if 'blur_option' not in st.session_state:
-    st.session_state.blur_option = st.checkbox("NSFW画像にBlurをかける", value=True, key="blur_toggle")
+    st.session_state.blur_option = st.checkbox("NSFW画像にBlurをかける", value=st.session_state.blur_option, key="blur_toggle")
+
+# Initialize sort options
+if 'sort_option' not in st.session_state:
+    st.session_state.sort_option = {"column": "timestamp", "ascending": True}
 
 # Lists to store Gradio predictions, good values, timestamps, and file paths
 gradio_nsfw_predictions = []
@@ -24,12 +30,7 @@ good_values = []
 timestamps = []
 file_paths = []
 
-# Function to get Gradio NSFW prediction
-def get_gradio_nsfw_prediction(image_path):
-    result = client_nsfw.predict(image_path, api_name="/predict")
-    return result  # Return the entire dictionary
-
-# Define the function to generate a download link for a file
+# Function to generate a download link for a file
 def get_binary_file_downloader_html(file_path, label="Download"):
     with open(file_path, 'rb') as f:
         data = f.read()
@@ -37,6 +38,11 @@ def get_binary_file_downloader_html(file_path, label="Download"):
     href = f'<a href="data:file/csv;base64,{b64}" download="{os.path.basename(file_path)}">{label}</a>'
     return href
 
+# Function to get Gradio NSFW prediction
+def get_gradio_nsfw_prediction(image_path):
+    result = client_nsfw.predict(image_path, api_name="/predict")
+    return result["label"]  # Assuming the label is the prediction result
+
 # Function to save uploaded files and get Gradio predictions
 def save_uploaded_files(uploaded_files):
     if not os.path.exists("temp"):
@@ -47,20 +53,21 @@ def save_uploaded_files(uploaded_files):
         with open(file_path, "wb") as f:
             f.write(uploaded_file.getbuffer())
         
-        # Get Gradio NSFW prediction
+        # Get Gradio predictions
         gradio_nsfw_prediction = get_gradio_nsfw_prediction(file_path)
         
-        # Extract good value confidence
-        good_value_confidence = None
+        gradio_nsfw_predictions.append((file_path, gradio_nsfw_prediction))
+        
+        # Get good value
+        good_value = None
         if "good" in gradio_nsfw_prediction and "confidences" in gradio_nsfw_prediction["good"]:
             confidences = gradio_nsfw_prediction["good"]["confidences"]
             for label_confidence in confidences:
                 if label_confidence["label"] == "good":
-                    good_value_confidence = label_confidence["confidence"]
+                    good_value = label_confidence["confidence"]
                     break
         
-        gradio_nsfw_predictions.append((file_path, gradio_nsfw_prediction))
-        good_values.append(good_value_confidence)
+        good_values.append(good_value)
         
         # Get timestamp
         timestamp = datetime.datetime.now()
@@ -74,7 +81,7 @@ def apply_blur(image_path):
     img = img.filter(ImageFilter.GaussianBlur(radius=5))
     return img
 
-# Function to paginate files for display
+# Function to paginate files
 def paginate_files(files, page, files_per_page):
     start_index = (page - 1) * files_per_page
     end_index = start_index + files_per_page
@@ -94,7 +101,7 @@ def display_images(images):
             st.write(f"**Good value:** {good_values[i]}")
             st.write(f"**Timestamp:** {timestamp}")
 
-            if st.session_state.blur_option and gradio_nsfw_prediction.get("label", "").lower() == "nsfw":
+            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:
@@ -109,13 +116,9 @@ def display_images(images):
                         os.remove(file_path)
                         st.experimental_rerun()
 
-# Main part of the Streamlit app
-st.title("ファイルアップロードと共有")
-
-uploaded_files = st.file_uploader("ファイルをアップロードしてください", type=["jpg", "png"], accept_multiple_files=True)
-
-if uploaded_files:
-    save_uploaded_files(uploaded_files)
+# Main part of the app
+if 'uploaded_files' in st.session_state:
+    save_uploaded_files(st.session_state.uploaded_files)
 
 files_per_page = 20
 num_pages = (len(file_paths) - 1) // files_per_page + 1
@@ -142,6 +145,9 @@ st.write(f"現在のページ: {st.session_state.page}")
 
 current_page_files = paginate_files(file_paths, st.session_state.page, files_per_page)
 
+st.write("ファイル一覧:")
+st.write(current_page_files)
+
 # Display images with predictions and options
 display_images(current_page_files)
 
@@ -155,7 +161,10 @@ if st.button("CSVに保存"):
         csv_writer = csv.writer(csv_file)
         csv_writer.writerow(['ファイル名', 'NSFWの予測', 'Goodの値', '作成時間'])
         for i, file_path in enumerate(file_paths):
-            csv_writer.writerow([os.path.basename(file_path), gradio_nsfw_predictions[i][1], good_values[i], timestamps[i]])
+            gradio_nsfw_prediction = gradio_nsfw_predictions[i]
+            good_value = good_values[i]
+            timestamp = timestamps[i]
+            csv_writer.writerow([os.path.basename(file_path), gradio_nsfw_prediction, good_value, timestamp])
 
     st.markdown(get_binary_file_downloader_html(csv_filename, label="CSVをダウンロード"), unsafe_allow_html=True)