Update app.py

app.py

@@ -1,11 +1,15 @@
 import streamlit as st
 import os
 import zipfile
+import csv
+import base64  # 追加
 from PIL import Image, ImageFilter
 from gradio_client import Client
+import datetime
 
 # Initialize Gradio client
-client = Client("https://ozoneasai-falconsai-nsfw-image-detection.hf.space/--replicas/0fyhj/")
+client_nsfw = Client("https://ozoneasai-falconsai-nsfw-image-detection.hf.space/--replicas/0fyhj/")
+client_other = Client("https://ozoneasai-sanali209-imclasif-quality-v001.hf.space/--replicas/m6xqz/")
 
 # Initialize pagination
 if "page" not in st.session_state:
@@ -15,19 +19,42 @@ if "page" not in st.session_state:
 if "blur_option" not in st.session_state:
     st.session_state.blur_option = True
 
+# Initialize sort options
+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
+def get_binary_file_downloader_html(file_path, label="Download"):
+    with open(file_path, 'rb') as f:
+        data = f.read()
+    b64 = base64.b64encode(data).decode()
+    href = f'<a href="data:file/csv;base64,{b64}" download="{os.path.basename(file_path)}">{label}</a>'
+    return href
+
 st.title("ファイルアップロードと共有")
 
 uploaded_files = st.file_uploader("ファイルをアップロードしてください", type=["jpg", "png"], accept_multiple_files=True)
 
 # List to store Gradio predictions
-gradio_predictions = []
+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_prediction(image_path):
-    result = client.predict(image_path, api_name="/predict")
-    return result
+def get_gradio_nsfw_prediction(image_path):
+    result = client_nsfw.predict(image_path, api_name="/predict")
+    return result["output"]  # 追加
+
+def get_gradio_other_prediction(image_path):
+    result = client_other.predict(image_path, api_name="/predict")
+    return result["output"]  # 追加
 
 # Function to save uploaded files and get Gradio predictions
 def save_uploaded_files(uploaded_files):
@@ -39,9 +66,23 @@ def save_uploaded_files(uploaded_files):
         with open(file_path, "wb") as f:
             f.write(uploaded_file.getbuffer())
         
-        # Get Gradio prediction
-        gradio_prediction = get_gradio_prediction(file_path)
-        gradio_predictions.append((file_path, gradio_prediction))
+        # Get Gradio predictions
+        gradio_nsfw_prediction = get_gradio_nsfw_prediction(file_path)
+        gradio_other_prediction = get_gradio_other_prediction(file_path)
+        
+        gradio_nsfw_predictions.append((file_path, gradio_nsfw_prediction))
+        gradio_other_predictions.append((file_path, gradio_other_prediction))
+        
+        # Get good value
+        good_value = None
+        if "confidences" in gradio_other_prediction:
+            good_value = gradio_other_prediction["confidences"][0]["confidence"]
+        good_values.append(good_value)
+        
+        # Get timestamp
+        timestamp = datetime.datetime.now()
+        timestamps.append(timestamp)
+        
         file_paths.append(file_path)
 
 # Function to apply blur to an image
@@ -58,14 +99,21 @@ def paginate_files(files, page, files_per_page):
 
 # Display images with pagination
 def display_images(images, blur_option):
-    for file_path in images:
+    for i, file_path in enumerate(images):
         file_ext = os.path.splitext(file_path)[1].lower()
 
-        gradio_prediction = get_gradio_prediction(file_path)
+        gradio_nsfw_prediction = get_gradio_nsfw_prediction(file_path)
+        gradio_other_prediction = get_gradio_other_prediction(file_path)
+        
+        good_value = good_values[i]
+        timestamp = timestamps[i]
 
-        st.write(f"**Prediction for {os.path.basename(file_path)}:** {gradio_prediction}")
+        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}")
 
-        if blur_option and gradio_prediction.lower() == "nsfw":
+        if blur_option and gradio_nsfw_prediction.lower() == "nsfw":
             # Apply blur if NSFW prediction and blur option is enabled
             blurred_img = apply_blur(file_path)
             st.image(blurred_img, caption=os.path.basename(file_path), use_column_width=True)
@@ -106,15 +154,45 @@ 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)
 
 # Blurチェックボックス
-st.checkbox("NSFW画像にBlurをかける", key="blur_option")
+st.session_state.blur_option = st.checkbox("NSFW画像にBlurをかける", key="blur_option", value=st.session_state.blur_option)
+
+# 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_predictions:
-        f.write(f"{prediction[0]}: {prediction[1]}\n")
+    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("ファイルを一括ダウンロード"):
@@ -124,6 +202,6 @@ if st.button("ファイルを一括ダウンロード"):
             zipf.write(file_path, os.path.basename(file_path))
 
     # Zipファイルをリンクとして提供
-    st.markdown(get_binary_file_downloader_html(zip_filename), unsafe_allow_html=True)
+    st.markdown(get_binary_file_downloader_html(zip_filename, label="Zipファイルをダウンロード"), unsafe_allow_html=True)
 
 st.success("ファイルがアップロードされました。このページのURLを他のクライアントと共有してください。")