Update app.py

app.py

@@ -1,4 +1,7 @@
-import os, zipfile, tempfile, io
+import os
+import zipfile
+import tempfile
+import requests
 import numpy as np
 import pandas as pd
 from PIL import Image
@@ -8,13 +11,21 @@ from torchvision import transforms
 from torchvision.models import resnet50, ResNet50_Weights
 from sklearn.cluster import MiniBatchKMeans
 import matplotlib.pyplot as plt
+import io
+from datetime import datetime
+
 import gradio as gr
+
+# Face analysis
+from deepface import DeepFace
 import cv2
-import requests
 
 # ---------------------------
-# Device
+# 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")
 
 # ---------------------------
@@ -23,35 +34,53 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 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, closest_color = float("inf"), 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_))]
@@ -60,151 +89,152 @@ def get_dominant_color(image,num_colors=5):
     return dominant_color, hex_color
 
 # ---------------------------
-# OpenCV DNN Face + Gender
-# ---------------------------
-os.makedirs("models", exist_ok=True)
-
-# Face detection model
-FACE_PROTO = "models/deploy.prototxt"
-FACE_MODEL = "models/res10_300x300_ssd_iter_140000_fp16.caffemodel"
-if not os.path.exists(FACE_PROTO):
-    r = requests.get("https://raw.githubusercontent.com/opencv/opencv/master/samples/dnn/face_detector/deploy.prototxt"); open(FACE_PROTO,"wb").write(r.content)
-if not os.path.exists(FACE_MODEL):
-    r = requests.get("https://raw.githubusercontent.com/opencv/opencv_3rdparty/master/res10_300x300_ssd_iter_140000_fp16.caffemodel"); open(FACE_MODEL,"wb").write(r.content)
-
-# Gender model
-GENDER_PROTO = "models/deploy_gender.prototxt"
-GENDER_MODEL = "models/gender_net.caffemodel"
-if not os.path.exists(GENDER_PROTO):
-    r = requests.get("https://raw.githubusercontent.com/spmallick/learnopencv/master/AgeGender/deploy_gender.prototxt"); open(GENDER_PROTO,"wb").write(r.content)
-if not os.path.exists(GENDER_MODEL):
-    r = requests.get("https://raw.githubusercontent.com/spmallick/learnopencv/master/AgeGender/gender_net.caffemodel"); open(GENDER_MODEL,"wb").write(r.content)
-
-face_net = cv2.dnn.readNet(FACE_MODEL, FACE_PROTO)
-gender_net = cv2.dnn.readNet(GENDER_MODEL, GENDER_PROTO)
-GENDER_LIST = ["Homme","Femme"]
-
-def detect_faces_and_gender(image):
-    img = np.array(image)[:, :, ::-1]  # PIL RGB -> BGR
-    h, w = img.shape[:2]
-    blob = cv2.dnn.blobFromImage(img, 1.0, (300,300), [104,117,123], swapRB=False)
-    face_net.setInput(blob)
-    detections = face_net.forward()
-    faces_data = []
-
-    for i in range(detections.shape[2]):
-        confidence = detections[0,0,i,2]
-        if confidence > 0.5:
-            box = detections[0,0,i,3:7] * np.array([w,h,w,h])
-            x1,y1,x2,y2 = box.astype(int)
-            x1,y1,x2,y2 = max(0,x1), max(0,y1), min(w,x2), min(h,y2)
-            face_img = img[y1:y2, x1:x2]
-            if face_img.size == 0:
-                continue
-            face_blob = cv2.dnn.blobFromImage(face_img, 1.0, (227,227),
-                                             [78.4263377603, 87.7689143744, 114.895847746], swapRB=False)
-            gender_net.setInput(face_blob)
-            gender_preds = gender_net.forward()
-            gender = GENDER_LIST[gender_preds[0].argmax()]
-            faces_data.append({"bbox":(x1,y1,x2,y2),"gender":gender})
-    return faces_data
-
-# ---------------------------
-# Core analysis
+# Core function
 # ---------------------------
 def classify_zip_and_analyze_color(zip_file):
     results = []
-    images_list = []
+
     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 not fname.lower().endswith(('.png','.jpg','.jpeg')):
-                continue
-            img_path = os.path.join(tmpdir,fname)
-            try:
-                image = Image.open(img_path).convert("RGB")
-                images_list.append((image.copy(), fname))
-            except:
-                continue
-
-            # Image classification
-            input_tensor = transform(image).unsqueeze(0).to(device)
-            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)]
-
-            # Dominant color
-            rgb, hex_color = get_dominant_color(image)
-            basic_color = closest_basic_color(rgb)
-
-            # Face + gender detection
-            faces_data = detect_faces_and_gender(image)
-            faces_str = "; ".join([f"Gender: {f['gender']}" for f in faces_data])
-
-            results.append((
-                fname,
-                ", ".join([p[0] for p in preds]),
-                ", ".join([p[1] for p in preds]),
-                hex_color,
-                basic_color,
-                faces_str
-            ))
-
-    df = pd.DataFrame(results, columns=["Filename","Top 3 Predictions","Confidence","Dominant Color","Basic Color","Face Info"])
-    out_xlsx = os.path.join(tempfile.gettempdir(), f"{zip_name}_results.xlsx")
-    df.to_excel(out_xlsx,index=False)
+            if fname.lower().endswith(('.png', '.jpg', '.jpeg')):
+                img_path = os.path.join(tmpdir, fname)
+                try:
+                    image = Image.open(img_path).convert("RGB")
+                except Exception:
+                    continue
+
+                # Classification
+                input_tensor = transform(image).unsqueeze(0).to(device)
+                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)]
+
+                # Dominant color
+                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)
+                    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": "Homme" if f["gender"]=="Man" else "Femme",
+                                "emotion": f["dominant_emotion"]
+                            })
+                    else:
+                        faces_data.append({
+                            "age": detected_faces["age"],
+                            "gender": "Homme" if detected_faces["gender"]=="Man" else "Femme",
+                            "emotion": detected_faces["dominant_emotion"]
+                        })
+                except Exception:
+                    faces_data = []
+
+                results.append((
+                    fname,
+                    ", ".join([p[0] for p in preds]),
+                    ", ".join([p[1] for p in preds]),
+                    hex_color,
+                    basic_color,
+                    faces_data
+                ))
+
+    # Build dataframe
+    df = pd.DataFrame(results, columns=[
+        "Filename", "Top 3 Predictions", "Confidence",
+        "Dominant Color", "Basic Color", "Face Info"
+    ])
+
+    # Save XLSX with zip name + date
+    out_xlsx = os.path.join(tempfile.gettempdir(), f"{zip_name}_{date_str}_results.xlsx")
+    df.to_excel(out_xlsx, index=False)
 
     # ---------------------------
-    # Plots
+    # 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)
+    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(", "))
+    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)
+    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)
+
+    # ---------------------------
+    # Plot 3: Gender distribution
+    # ---------------------------
+    gender_counts = {"Homme":0, "Femme":0}
+    for face_list in df["Face Info"]:
+        for face in face_list:
+            gender_counts[face["gender"]] += 1
 
-    # Gender distribution
-    gender_counts = [df["Face Info"].str.count("Homme").sum(), df["Face Info"].str.count("Femme").sum()]
     fig3, ax3 = plt.subplots()
-    ax3.bar(["Homme","Femme"], gender_counts, color=["lightblue","pink"])
-    ax3.set_title("Gender Distribution"); ax3.set_ylabel("Count")
-    buf3 = io.BytesIO(); plt.savefig(buf3, format="png"); plt.close(fig3); buf3.seek(0); plot3_img = Image.open(buf3)
+    ax3.bar(gender_counts.keys(), gender_counts.values(), color=["lightblue","pink"])
+    ax3.set_title("Gender Distribution")
+    ax3.set_ylabel("Count")
+    buf3 = io.BytesIO()
+    plt.savefig(buf3, format="png")
+    plt.close(fig3)
+    buf3.seek(0)
+    plot3_img = Image.open(buf3)
 
-    return df, images_list, out_xlsx, plot1_img, plot2_img, plot3_img
+    return df, out_xlsx, plot1_img, plot2_img, plot3_img
 
 # ---------------------------
-# Gradio interface
+# 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")
-
-    output_df = gr.Dataframe(headers=["Filename","Top 3 Predictions","Confidence","Dominant Color","Basic Color","Face Info"])
-    image_gallery = gr.Gallery(label="Preview Images", columns=4, show_label=True)
-    download_file = gr.File(label="Download XLSX")
-
-    plot1 = gr.Image(label="Basic Color Frequency")
-    plot2 = gr.Image(label="Top Prediction Distribution")
-    plot3 = gr.Image(label="Gender Distribution")
-
-    analyze_btn.click(
-        classify_zip_and_analyze_color,
-        inputs=uploaded_zip,
-        outputs=[output_df, image_gallery, download_file, plot1, plot2, plot3]
-    )
-
-demo.launch(server_name="0.0.0.0", server_port=7860)
+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"],
+            datatype=["str","str","str","str","str","str"]
+        ),
+        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"),
+    ],
+    title="Image Classifier with Color & Face Analysis",
+    description="Upload a ZIP of images. Classifies images, analyzes dominant color, detects/characterizes faces (age, gender, emotion).",
+)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)