Update app.py

app.py

@@ -18,12 +18,16 @@ import gradio as gr
 from deepface import DeepFace
 import cv2
 
-# CUDA setup
+# ---------------------------
+# Device setup
+# ---------------------------
 if not torch.cuda.is_available():
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load model
+# ---------------------------
+# Load ResNet50
+# ---------------------------
 weights = ResNet50_Weights.DEFAULT
 model = resnet50(weights=weights).to(device)
 model.eval()
@@ -64,10 +68,12 @@ def get_dominant_color(image,num_colors=5):
     hex_color = f"#{dominant_color[0]:02x}{dominant_color[1]:02x}{dominant_color[2]:02x}"
     return dominant_color, hex_color
 
-# Main function
+# ---------------------------
+# Core analysis
+# ---------------------------
 def classify_zip_and_analyze_color(zip_file):
     results = []
-    images_dict = {}  # store images by filename for preview
+    images_dict = {}  # store images for preview
     zip_name = os.path.splitext(os.path.basename(zip_file.name))[0]
     date_str = datetime.now().strftime("%Y%m%d")
 
@@ -80,7 +86,7 @@ def classify_zip_and_analyze_color(zip_file):
                 img_path = os.path.join(tmpdir,fname)
                 try:
                     image = Image.open(img_path).convert("RGB")
-                    images_dict[fname] = image.copy()  # save for preview
+                    images_dict[fname] = image.copy()
                 except:
                     continue
 
@@ -118,38 +124,34 @@ def classify_zip_and_analyze_color(zip_file):
                     faces_data
                 ))
 
-    # DataFrame
     df = pd.DataFrame(results, columns=["Filename","Top 3 Predictions","Confidence","Dominant Color","Basic Color","Face Info"])
 
-    # XLSX output
     out_xlsx = os.path.join(tempfile.gettempdir(), f"{zip_name}_{date_str}_results.xlsx")
     df.to_excel(out_xlsx,index=False)
 
     return df, images_dict, out_xlsx
 
-# Callback to show image preview when clicking filename
+# Callback for preview
 def show_preview(filename, images_dict):
-    if filename in images_dict:
-        return images_dict[filename]
-    else:
-        return None
+    return images_dict.get(filename, None)
 
+# ---------------------------
 # Gradio interface
+# ---------------------------
 with gr.Blocks() as demo:
     uploaded_zip = gr.File(label="Upload ZIP of images", file_types=[".zip"])
     output_df = gr.Dataframe(headers=["Filename","Top 3 Predictions","Confidence","Dominant Color","Basic Color","Face Info"])
     image_preview = gr.Image(label="Image Preview")
     download_file = gr.File(label="Download XLSX")
+    images_state = gr.State()  # store images dict
+
+    analyze_btn = gr.Button("Run Analysis")
 
-    # Run analysis
     def run_analysis(zip_file):
         df, images_dict, out_xlsx = classify_zip_and_analyze_color(zip_file)
         return df, images_dict, out_xlsx
 
-    analyze_btn = gr.Button("Run Analysis")
-    analyze_btn.click(run_analysis, inputs=uploaded_zip, outputs=[output_df, "state", download_file])
-
-    # Update preview when clicking filename
-    output_df.select(show_preview, inputs=[output_df, "state"], outputs=image_preview)
+    analyze_btn.click(run_analysis, inputs=uploaded_zip, outputs=[output_df, images_state, download_file])
+    output_df.select(show_preview, inputs=[output_df, images_state], outputs=image_preview)
 
 demo.launch(server_name="0.0.0.0", server_port=7860)