Update app.py

app.py

@@ -1,97 +1,12 @@
-import os
-import zipfile
-import tempfile
-import requests
-import numpy as np
-import pandas as pd
-from PIL import Image
-import torch
-import torch.nn.functional as F
-from torchvision import transforms
-from torchvision.models import resnet50, ResNet50_Weights
-from sklearn.cluster import MiniBatchKMeans
-import matplotlib.pyplot as plt
-import io
-
-import gradio as gr
-
-# Face analysis
-from deepface import DeepFace
-import cv2
-
-# ---------------------------
-# Force CPU if no CUDA
-# ---------------------------
-if not torch.cuda.is_available():
-    os.environ["CUDA_VISIBLE_DEVICES"] = ""
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# ---------------------------
-# Load ResNet50
-# ---------------------------
-weights = ResNet50_Weights.DEFAULT
-model = resnet50(weights=weights).to(device)
-model.eval()
-
-# ---------------------------
-# Transformations
-# ---------------------------
-transform = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std=[0.229, 0.224, 0.225]),
-])
-
-# ---------------------------
-# ImageNet labels
-# ---------------------------
-LABELS_URL = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
-imagenet_classes = [line.strip() for line in requests.get(LABELS_URL).text.splitlines()]
-
 # ---------------------------
-# Color utilities
-# ---------------------------
-BASIC_COLORS = {
-    "Red": (255, 0, 0),
-    "Green": (0, 255, 0),
-    "Blue": (0, 0, 255),
-    "Yellow": (255, 255, 0),
-    "Cyan": (0, 255, 255),
-    "Magenta": (255, 0, 255),
-    "Black": (0, 0, 0),
-    "White": (255, 255, 255),
-    "Gray": (128, 128, 128),
-}
-
-def closest_basic_color(rgb):
-    r, g, b = rgb
-    min_dist = float("inf")
-    closest_color = None
-    for name, (cr, cg, cb) in BASIC_COLORS.items():
-        dist = (r - cr) ** 2 + (g - cg) ** 2 + (b - cb) ** 2
-        if dist < min_dist:
-            min_dist = dist
-            closest_color = name
-    return closest_color
-
-def get_dominant_color(image, num_colors=5):
-    image = image.resize((100, 100))
-    pixels = np.array(image).reshape(-1, 3)
-    kmeans = MiniBatchKMeans(n_clusters=num_colors, random_state=0, n_init=5)
-    kmeans.fit(pixels)
-    dominant_color = kmeans.cluster_centers_[np.argmax(np.bincount(kmeans.labels_))]
-    dominant_color = tuple(dominant_color.astype(int))
-    hex_color = f"#{dominant_color[0]:02x}{dominant_color[1]:02x}{dominant_color[2]:02x}"
-    return dominant_color, hex_color
-
-# ---------------------------
-# Core function
+# Core function with gallery and renamed XLSX
 # ---------------------------
 def classify_zip_and_analyze_color(zip_file):
     results = []
+    thumbnails = []
+
+    # Get base name of zip to rename XLSX
+    zip_basename = os.path.splitext(os.path.basename(zip_file.name))[0]
 
     with tempfile.TemporaryDirectory() as tmpdir:
         with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
@@ -105,6 +20,11 @@ def classify_zip_and_analyze_color(zip_file):
                 except Exception:
                     continue
 
+                # Create small thumbnail for gallery
+                thumb = image.copy()
+                thumb.thumbnail((100, 100))
+                thumbnails.append(thumb)
+
                 # Classification
                 input_tensor = transform(image).unsqueeze(0).to(device)
                 with torch.no_grad():
@@ -118,9 +38,7 @@ def classify_zip_and_analyze_color(zip_file):
                 rgb, hex_color = get_dominant_color(image)
                 basic_color = closest_basic_color(rgb)
 
-                # ---------------------------
                 # Face detection & characterization
-                # ---------------------------
                 faces_data = []
                 try:
                     img_cv2 = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
@@ -131,13 +49,13 @@ def classify_zip_and_analyze_color(zip_file):
                         for f in detected_faces:
                             faces_data.append({
                                 "age": f["age"],
-                                "gender": f["gender"],  # dict of probabilities
+                                "gender": f["gender"],
                                 "emotion": f["dominant_emotion"]
                             })
                     else:
                         faces_data.append({
                             "age": detected_faces["age"],
-                            "gender": detected_faces["gender"],  # dict
+                            "gender": detected_faces["gender"],
                             "emotion": detected_faces["dominant_emotion"]
                         })
                 except Exception:
@@ -155,13 +73,14 @@ def classify_zip_and_analyze_color(zip_file):
     # Build dataframe
     df = pd.DataFrame(results, columns=["Filename", "Top 3 Predictions", "Confidence", "Dominant Color", "Basic Color", "Face Info"])
 
-    # Save XLSX
-    out_xlsx = os.path.join(tempfile.gettempdir(), "results.xlsx")
+    # Save XLSX with zip name
+    out_xlsx = os.path.join(tempfile.gettempdir(), f"{zip_basename}_results.xlsx")
     df.to_excel(out_xlsx, index=False)
 
     # ---------------------------
-    # Plot 1: Basic color frequency
+    # Plotting code (same as before)
     # ---------------------------
+    # Basic color frequency
     fig1, ax1 = plt.subplots()
     color_counts = df["Basic Color"].value_counts()
     ax1.bar(color_counts.index, color_counts.values, color="skyblue")
@@ -173,9 +92,7 @@ def classify_zip_and_analyze_color(zip_file):
     buf1.seek(0)
     plot1_img = Image.open(buf1)
 
-    # ---------------------------
-    # Plot 2: Top prediction distribution
-    # ---------------------------
+    # Top prediction distribution
     fig2, ax2 = plt.subplots()
     preds_flat = []
     for p in df["Top 3 Predictions"]:
@@ -190,30 +107,24 @@ def classify_zip_and_analyze_color(zip_file):
     buf2.seek(0)
     plot2_img = Image.open(buf2)
 
-    # ---------------------------
-    # Extract age and weighted gender (confidence ≤ 0.9)
-    # ---------------------------
+    # Gender distribution
     ages = []
     gender_confidence = {"Man": 0, "Woman": 0}
-
     for face_list in df["Face Info"]:
-        for face in face_list:  # each face is a dict
+        for face in face_list:
             ages.append(face["age"])
-            gender_dict = face["gender"]  # dict of probabilities
+            gender_dict = face["gender"]
             gender = max(gender_dict, key=gender_dict.get)
-            conf = float(gender_dict[gender]) / 100  # convert % to 0-1
+            conf = float(gender_dict[gender]) / 100
             weight = min(conf, 0.9)
             if gender in gender_confidence:
                 gender_confidence[gender] += weight
             else:
                 gender_confidence[gender] = weight
 
-    # ---------------------------
-    # Plot 3: Gender distribution (weighted ≤ 0.9)
-    # ---------------------------
     fig3, ax3 = plt.subplots()
     ax3.bar(gender_confidence.keys(), gender_confidence.values(), color=["lightblue", "pink"])
-    ax3.set_title("Gender Distribution (Weighted ≤90% Confidence)")
+    ax3.set_title("Gender Distribution (Weighted ≤90%)")
     ax3.set_ylabel("Sum of Confidence")
     buf3 = io.BytesIO()
     plt.savefig(buf3, format="png")
@@ -221,9 +132,7 @@ def classify_zip_and_analyze_color(zip_file):
     buf3.seek(0)
     plot3_img = Image.open(buf3)
 
-    # ---------------------------
-    # Plot 4: Age distribution
-    # ---------------------------
+    # Age distribution
     fig4, ax4 = plt.subplots()
     ax4.hist(ages, bins=range(0, 101, 5), color="lightgreen", edgecolor="black")
     ax4.set_title("Age Distribution")
@@ -235,10 +144,10 @@ def classify_zip_and_analyze_color(zip_file):
     buf4.seek(0)
     plot4_img = Image.open(buf4)
 
-    return df, out_xlsx, plot1_img, plot2_img, plot3_img, plot4_img
+    return df, out_xlsx, thumbnails, plot1_img, plot2_img, plot3_img, plot4_img
 
 # ---------------------------
-# Gradio Interface
+# Gradio Interface with gallery
 # ---------------------------
 demo = gr.Interface(
     fn=classify_zip_and_analyze_color,
@@ -246,14 +155,15 @@ demo = gr.Interface(
     outputs=[
         gr.Dataframe(headers=["Filename", "Top 3 Predictions", "Confidence", "Dominant Color", "Basic Color", "Face Info"]),
         gr.File(label="Download XLSX"),
+        gr.Gallery(label="Thumbnails", elem_id="thumbnail-gallery").style(grid=[5], height="auto"),
         gr.Image(type="pil", label="Basic Color Frequency"),
         gr.Image(type="pil", label="Top Prediction Distribution"),
         gr.Image(type="pil", label="Gender Distribution (Weighted ≤90%)"),
         gr.Image(type="pil", label="Age Distribution"),
     ],
     title="Image Classifier with Color & Face Analysis",
-    description="Upload a ZIP of images. Classifies images, analyzes dominant color, and detects/characterizes faces (age, gender, emotion).",
+    description="Upload a ZIP of images. Classifies images, analyzes dominant color, detects faces, and displays thumbnails.",
 )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)