Update app.py

app.py

@@ -1,11 +1,100 @@
+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()]
+
+# ---------------------------
+# Basic 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 with gallery and renamed XLSX
+# Core function
 # ---------------------------
 def classify_zip_and_analyze_color(zip_file):
     results = []
     thumbnails = []
 
-    # Get base name of zip to rename XLSX
+    # Name XLSX after zip
     zip_basename = os.path.splitext(os.path.basename(zip_file.name))[0]
 
     with tempfile.TemporaryDirectory() as tmpdir:
@@ -20,7 +109,7 @@ def classify_zip_and_analyze_color(zip_file):
                 except Exception:
                     continue
 
-                # Create small thumbnail for gallery
+                # Thumbnail for gallery
                 thumb = image.copy()
                 thumb.thumbnail((100, 100))
                 thumbnails.append(thumb)
@@ -73,14 +162,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 with zip name
+    # Save XLSX
     out_xlsx = os.path.join(tempfile.gettempdir(), f"{zip_basename}_results.xlsx")
     df.to_excel(out_xlsx, index=False)
 
     # ---------------------------
-    # Plotting code (same as before)
+    # Plots
     # ---------------------------
-    # Basic color frequency
+    # 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")
@@ -92,7 +181,7 @@ def classify_zip_and_analyze_color(zip_file):
     buf1.seek(0)
     plot1_img = Image.open(buf1)
 
-    # Top prediction distribution
+    # 2. Top prediction distribution
     fig2, ax2 = plt.subplots()
     preds_flat = []
     for p in df["Top 3 Predictions"]:
@@ -107,7 +196,7 @@ def classify_zip_and_analyze_color(zip_file):
     buf2.seek(0)
     plot2_img = Image.open(buf2)
 
-    # Gender distribution
+    # 3. Gender distribution
     ages = []
     gender_confidence = {"Man": 0, "Woman": 0}
     for face_list in df["Face Info"]:
@@ -132,7 +221,7 @@ def classify_zip_and_analyze_color(zip_file):
     buf3.seek(0)
     plot3_img = Image.open(buf3)
 
-    # Age distribution
+    # 4. Age distribution
     fig4, ax4 = plt.subplots()
     ax4.hist(ages, bins=range(0, 101, 5), color="lightgreen", edgecolor="black")
     ax4.set_title("Age Distribution")
@@ -147,7 +236,7 @@ def classify_zip_and_analyze_color(zip_file):
     return df, out_xlsx, thumbnails, plot1_img, plot2_img, plot3_img, plot4_img
 
 # ---------------------------
-# Gradio Interface with gallery
+# Gradio Interface
 # ---------------------------
 demo = gr.Interface(
     fn=classify_zip_and_analyze_color,