Update app.py

app.py

@@ -15,72 +15,48 @@ import io
 from datetime import datetime
 
 import gradio as gr
-
-# Face analysis
 from deepface import DeepFace
 import cv2
 
-# ---------------------------
-# Force CPU if no CUDA
-# ---------------------------
+# CUDA setup
 if not torch.cuda.is_available():
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ---------------------------
-# Load ResNet50
-# ---------------------------
+# Load model
 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]),
+    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),
+    "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
+    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)
+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_))]
@@ -88,25 +64,24 @@ 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
 
-# ---------------------------
-# Core function
-# ---------------------------
+# Main function
 def classify_zip_and_analyze_color(zip_file):
     results = []
-
+    images_dict = {}  # store images by filename for preview
     zip_name = os.path.splitext(os.path.basename(zip_file.name))[0]
     date_str = datetime.now().strftime("%Y%m%d")
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
+        with zipfile.ZipFile(zip_file.name,'r') as zip_ref:
             zip_ref.extractall(tmpdir)
 
         for fname in sorted(os.listdir(tmpdir)):
-            if fname.lower().endswith(('.png', '.jpg', '.jpeg')):
-                img_path = os.path.join(tmpdir, fname)
+            if fname.lower().endswith(('.png','.jpg','.jpeg')):
+                img_path = os.path.join(tmpdir,fname)
                 try:
                     image = Image.open(img_path).convert("RGB")
-                except Exception:
+                    images_dict[fname] = image.copy()  # save for preview
+                except:
                     continue
 
                 # Classification
@@ -114,40 +89,25 @@ def classify_zip_and_analyze_color(zip_file):
                 with torch.no_grad():
                     output = model(input_tensor)
                     probs = F.softmax(output, dim=1)[0]
-
-                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)]
+                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 & characterization
+                # Face detection
                 faces_data = []
                 try:
                     img_cv2 = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-                    detected_faces = DeepFace.analyze(
-                        img_cv2, actions=["age", "gender", "emotion"], enforce_detection=False
-                    )
+                    detected_faces = DeepFace.analyze(img_cv2, actions=["age","gender","emotion"], enforce_detection=False)
                     if isinstance(detected_faces, list):
                         for f in detected_faces:
-                            faces_data.append({
-                                "age": f["age"],
-                                "gender": f["gender"],
-                                "emotion": f["dominant_emotion"]
-                            })
+                            faces_data.append({"age": f["age"], "gender": f["gender"], "emotion": f["dominant_emotion"]})
                     else:
-                        faces_data.append({
-                            "age": detected_faces["age"],
-                            "gender": detected_faces["gender"],
-                            "emotion": detected_faces["dominant_emotion"]
-                        })
-                except Exception:
-                    faces_data = []
-
-                # Thumbnail preview
-                thumbnail = image.copy()
-                thumbnail.thumbnail((64, 64))
+                        faces_data.append({"age": detected_faces["age"], "gender": detected_faces["gender"], "emotion": detected_faces["dominant_emotion"]})
+                except:
+                    faces_data=[]
 
                 results.append((
                     fname,
@@ -155,123 +115,41 @@ def classify_zip_and_analyze_color(zip_file):
                     ", ".join([p[1] for p in preds]),
                     hex_color,
                     basic_color,
-                    faces_data,
-                    thumbnail
+                    faces_data
                 ))
 
-    # Build dataframe
-    df = pd.DataFrame(results, columns=[
-        "Filename", "Top 3 Predictions", "Confidence",
-        "Dominant Color", "Basic Color", "Face Info", "Thumbnail"
-    ])
+    # DataFrame
+    df = pd.DataFrame(results, columns=["Filename","Top 3 Predictions","Confidence","Dominant Color","Basic Color","Face Info"])
 
-    # Save XLSX with zip name + date
+    # XLSX output
     out_xlsx = os.path.join(tempfile.gettempdir(), f"{zip_name}_{date_str}_results.xlsx")
-    df.to_excel(out_xlsx, index=False)
-
-    # ---------------------------
-    # Plot 1: Basic color frequency
-    # ---------------------------
-    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)
-
-    # ---------------------------
-    # Plot 2: Top prediction distribution
-    # ---------------------------
-    fig2, ax2 = plt.subplots()
-    preds_flat = []
-    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)
+    df.to_excel(out_xlsx,index=False)
 
-    # ---------------------------
-    # Extract ages and genders
-    # ---------------------------
-    ages_male, ages_female = [], []
-    gender_confidence = {"Homme": 0, "Femme": 0}
+    return df, images_dict, out_xlsx
 
-    for face_list in df["Face Info"]:
-        for face in face_list:
-            age = face["age"]
-            gender_dict = face["gender"]
-            gender = max(gender_dict, key=gender_dict.get)
-            conf = float(gender_dict[gender]) / 100
-            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)
+# Callback to show image preview when clicking filename
+def show_preview(filename, images_dict):
+    if filename in images_dict:
+        return images_dict[filename]
+    else:
+        return None
 
-    # ---------------------------
-    # Plot 3: 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)
+# 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")
 
-    # ---------------------------
-    # Plot 4: Age distribution by gender
-    # ---------------------------
-    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)
+    # 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
 
-    return df, out_xlsx, plot1_img, plot2_img, plot3_img, plot4_img
+    analyze_btn = gr.Button("Run Analysis")
+    analyze_btn.click(run_analysis, inputs=uploaded_zip, outputs=[output_df, "state", download_file])
 
-# ---------------------------
-# Gradio Interface
-# ---------------------------
-demo = gr.Interface(
-    fn=classify_zip_and_analyze_color,
-    inputs=gr.File(file_types=[".zip"], label="Upload ZIP of images"),
-    outputs=[
-        gr.Dataframe(
-            headers=["Filename", "Top 3 Predictions", "Confidence",
-                     "Dominant Color", "Basic Color", "Face Info", "Thumbnail"],
-            datatype=["str","str","str","str","str","str","pil"]
-        ),
-        gr.File(label="Download XLSX"),
-        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 by Gender"),
-    ],
-    title="Image Classifier with Color & Face Analysis",
-    description="Upload a ZIP of images. Classifies images, analyzes dominant color, detects/characterizes faces (age, gender, emotion), and shows thumbnails.",
-)
+    # Update preview when clicking filename
+    output_df.select(show_preview, inputs=[output_df, "state"], outputs=image_preview)
 
-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)