Update app.py

app.py

@@ -73,7 +73,7 @@ def get_dominant_color(image,num_colors=5):
 # ---------------------------
 def classify_zip_and_analyze_color(zip_file):
     results = []
-    images_dict = {}  # store images for preview
+    images_dict = {}
     zip_name = os.path.splitext(os.path.basename(zip_file.name))[0]
     date_str = datetime.now().strftime("%Y%m%d")
 
@@ -90,7 +90,6 @@ def classify_zip_and_analyze_color(zip_file):
                 except:
                     continue
 
-                # Classification
                 input_tensor = transform(image).unsqueeze(0).to(device)
                 with torch.no_grad():
                     output = model(input_tensor)
@@ -98,11 +97,9 @@ def classify_zip_and_analyze_color(zip_file):
                 top3_prob, top3_idx = torch.topk(probs,3)
                 preds = [(imagenet_classes[idx], f"{prob.item()*100:.2f}%") for idx,prob in zip(top3_idx, top3_prob)]
 
-                # Dominant color
                 rgb, hex_color = get_dominant_color(image)
                 basic_color = closest_basic_color(rgb)
 
-                # Face detection
                 faces_data = []
                 try:
                     img_cv2 = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
@@ -129,39 +126,27 @@ def classify_zip_and_analyze_color(zip_file):
     df.to_excel(out_xlsx,index=False)
 
     # ---------------------------
-    # Plot 1: Basic color frequency
+    # Plots
     # ---------------------------
+    # Basic color
     fig1, ax1 = plt.subplots()
     color_counts = df["Basic Color"].value_counts()
     ax1.bar(color_counts.index, color_counts.values, color="skyblue")
     ax1.set_title("Basic Color Frequency")
     ax1.set_ylabel("Count")
-    buf1 = io.BytesIO()
-    plt.savefig(buf1, format="png")
-    plt.close(fig1)
-    buf1.seek(0)
-    plot1_img = Image.open(buf1)
+    buf1 = io.BytesIO(); plt.savefig(buf1, format="png"); plt.close(fig1); buf1.seek(0); plot1_img = Image.open(buf1)
 
-    # ---------------------------
-    # Plot 2: Top prediction distribution
-    # ---------------------------
+    # Top predictions
     fig2, ax2 = plt.subplots()
     preds_flat = []
-    for p in df["Top 3 Predictions"]:
-        preds_flat.extend(p.split(", "))
+    for p in df["Top 3 Predictions"]: preds_flat.extend(p.split(", "))
     pred_counts = pd.Series(preds_flat).value_counts().head(20)
     ax2.barh(pred_counts.index[::-1], pred_counts.values[::-1], color="salmon")
     ax2.set_title("Top Prediction Distribution")
     ax2.set_xlabel("Count")
-    buf2 = io.BytesIO()
-    plt.savefig(buf2, format="png", bbox_inches="tight")
-    plt.close(fig2)
-    buf2.seek(0)
-    plot2_img = Image.open(buf2)
+    buf2 = io.BytesIO(); plt.savefig(buf2, format="png", bbox_inches="tight"); plt.close(fig2); buf2.seek(0); plot2_img = Image.open(buf2)
 
-    # ---------------------------
-    # Gender & Age
-    # ---------------------------
+    # Gender and age
     ages_male, ages_female = [], []
     gender_confidence = {"Homme":0, "Femme":0}
     for face_list in df["Face Info"]:
@@ -173,55 +158,41 @@ def classify_zip_and_analyze_color(zip_file):
             weight = min(conf,0.9)
             gender_trans = "Homme" if gender=="Man" else "Femme"
             gender_confidence[gender_trans] += weight
-            if gender_trans=="Homme":
-                ages_male.append(age)
-            else:
-                ages_female.append(age)
+            if gender_trans=="Homme": ages_male.append(age)
+            else: ages_female.append(age)
 
-    # Gender distribution
     fig3, ax3 = plt.subplots()
     ax3.bar(gender_confidence.keys(), gender_confidence.values(), color=["lightblue","pink"])
     ax3.set_title("Gender Distribution (Weighted ≤90%)")
     ax3.set_ylabel("Sum of Confidence")
-    buf3 = io.BytesIO()
-    plt.savefig(buf3, format="png")
-    plt.close(fig3)
-    buf3.seek(0)
-    plot3_img = Image.open(buf3)
+    buf3 = io.BytesIO(); plt.savefig(buf3, format="png"); plt.close(fig3); buf3.seek(0); plot3_img = Image.open(buf3)
 
-    # Age distribution
     fig4, ax4 = plt.subplots()
     bins = range(0,101,5)
     ax4.hist([ages_male, ages_female], bins=bins, color=["lightblue","pink"], label=["Homme","Femme"], edgecolor="black")
     ax4.set_title("Age Distribution by Gender")
-    ax4.set_xlabel("Age")
-    ax4.set_ylabel("Count")
-    ax4.legend()
-    buf4 = io.BytesIO()
-    plt.savefig(buf4, format="png")
-    plt.close(fig4)
-    buf4.seek(0)
-    plot4_img = Image.open(buf4)
+    ax4.set_xlabel("Age"); ax4.set_ylabel("Count"); ax4.legend()
+    buf4 = io.BytesIO(); plt.savefig(buf4, format="png"); plt.close(fig4); buf4.seek(0); plot4_img = Image.open(buf4)
 
-    return df, list(images_dict.keys()), images_dict, out_xlsx, plot1_img, plot2_img, plot3_img, plot4_img
+    return df, images_dict, out_xlsx, plot1_img, plot2_img, plot3_img, plot4_img
 
 # ---------------------------
 # Preview callback
 # ---------------------------
-def show_preview(selected_file, images_state):
-    if images_state is None or selected_file is None:
+def show_preview(selected_row, images_state):
+    if images_state is None or selected_row is None:
         return None
-    return images_state.get(selected_file, None)
+    filename = selected_row[0]  # first column is filename
+    return images_state.get(filename, None)
 
 # ---------------------------
 # Gradio interface
 # ---------------------------
 with gr.Blocks() as demo:
     uploaded_zip = gr.File(label="Upload ZIP of images", file_types=[".zip"])
-    analyze_btn = gr.Button("Run Analysis")  # Run button just after upload
+    analyze_btn = gr.Button("Run Analysis")
 
-    output_df = gr.Dataframe(headers=["Filename","Top 3 Predictions","Confidence","Dominant Color","Basic Color","Face Info"])
-    image_dropdown = gr.Dropdown(label="Select image to preview")
+    output_df = gr.Dataframe(headers=["Filename","Top 3 Predictions","Confidence","Dominant Color","Basic Color","Face Info"], interactive=True)
     image_preview = gr.Image(label="Image Preview")
     download_file = gr.File(label="Download XLSX")
     images_state = gr.State()
@@ -232,15 +203,15 @@ with gr.Blocks() as demo:
     plot4 = gr.Image(label="Age Distribution by Gender")
 
     def run_analysis(zip_file):
-        df, filenames, images_dict, out_xlsx, p1, p2, p3, p4 = classify_zip_and_analyze_color(zip_file)
-        return df, filenames, images_dict, out_xlsx, p1, p2, p3, p4
+        df, images_dict, out_xlsx, p1, p2, p3, p4 = classify_zip_and_analyze_color(zip_file)
+        return df, images_dict, out_xlsx, p1, p2, p3, p4
 
     analyze_btn.click(
         run_analysis,
         inputs=uploaded_zip,
-        outputs=[output_df, image_dropdown, images_state, download_file, plot1, plot2, plot3, plot4]
+        outputs=[output_df, images_state, download_file, plot1, plot2, plot3, plot4]
     )
 
-    image_dropdown.change(show_preview, inputs=[image_dropdown, images_state], outputs=image_preview)
+    output_df.select(show_preview, inputs=[output_df, images_state], outputs=image_preview)
 
 demo.launch(server_name="0.0.0.0", server_port=7860)