Update app.py

app.py

@@ -16,8 +16,6 @@ import cv2
 # ---------------------------
 # Device
 # ---------------------------
-if not torch.cuda.is_available():
-    os.environ["CUDA_VISIBLE_DEVICES"] = ""
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # ---------------------------
@@ -76,50 +74,65 @@ def classify_zip_and_analyze_color(zip_file):
             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)
-                try:
-                    image = Image.open(img_path).convert("RGB")
-                    images_list.append((image.copy(), fname))
-                except:
-                    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)]
-
-                rgb, hex_color = get_dominant_color(image)
-                basic_color = closest_basic_color(rgb)
-
-                # Face analysis
-                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": f["gender"], "emotion": f["dominant_emotion"]})
+            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 analysis
+            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 not isinstance(detected_faces, list):
+                    detected_faces = [detected_faces]
+                for f in detected_faces:
+                    gender = f["gender"].lower()
+                    if gender in ["man", "male"]:
+                        gender_fr = "Homme"
+                    elif gender in ["woman", "female"]:
+                        gender_fr = "Femme"
                     else:
-                        faces_data.append({"age": detected_faces["age"], "gender": detected_faces["gender"], "emotion": detected_faces["dominant_emotion"]})
-                except:
-                    faces_data=[]
-
-                # Convert faces data to readable string
-                faces_str = "; ".join([f"Age: {face['age']}, Gender: {'Homme' if face['gender']=='Man' else 'Femme'}, Emotion: {face['emotion']}" for face 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
-                ))
-
+                        gender_fr = "Inconnu"
+                    faces_data.append({
+                        "age": f["age"],
+                        "gender": gender_fr,
+                        "emotion": f["dominant_emotion"]
+                    })
+            except:
+                faces_data=[]
+
+            faces_str = "; ".join([
+                f"Age: {face['age']}, Gender: {face['gender']}, Emotion: {face['emotion']}"
+                for face 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
+            ))
+
+    # Create DataFrame
     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}_{date_str}_results.xlsx")
     df.to_excel(out_xlsx,index=False)
@@ -127,12 +140,14 @@ def classify_zip_and_analyze_color(zip_file):
     # ---------------------------
     # Plots
     # ---------------------------
+    # 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)
 
+    # Top prediction distribution
     fig2, ax2 = plt.subplots()
     preds_flat = []
     for p in df["Top 3 Predictions"]: preds_flat.extend(p.split(", "))
@@ -142,8 +157,8 @@ def classify_zip_and_analyze_color(zip_file):
     buf2 = io.BytesIO(); plt.savefig(buf2, format="png", bbox_inches="tight"); plt.close(fig2); buf2.seek(0); plot2_img = Image.open(buf2)
 
     # Gender and age
-    ages_male, ages_female = [], []
-    gender_confidence = {"Homme":0, "Femme":0}
+    ages_male = []
+    ages_female = []
     for face_list in df["Face Info"]:
         if face_list.strip()=="":
             continue
@@ -151,16 +166,18 @@ def classify_zip_and_analyze_color(zip_file):
             parts = face_str.split(", ")
             age = int(parts[0].split(": ")[1])
             gender = parts[1].split(": ")[1]
-            conf = 0.9  # approximation for histogram
-            gender_confidence[gender] += conf
             if gender=="Homme": ages_male.append(age)
-            else: ages_female.append(age)
+            elif gender=="Femme": ages_female.append(age)
+
+    gender_counts = {"Homme": len(ages_male), "Femme": len(ages_female)}
 
+    # Gender distribution
     fig3, ax3 = plt.subplots()
-    ax3.bar(gender_confidence.keys(), gender_confidence.values(), color=["lightblue","pink"])
-    ax3.set_title("Gender Distribution"); ax3.set_ylabel("Sum of Confidence")
+    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)
 
+    # 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")
@@ -185,12 +202,8 @@ with gr.Blocks() as demo:
     plot3 = gr.Image(label="Gender Distribution")
     plot4 = gr.Image(label="Age Distribution by Gender")
 
-    def run_analysis(zip_file):
-        df, images_list, out_xlsx, p1, p2, p3, p4 = classify_zip_and_analyze_color(zip_file)
-        return df, images_list, out_xlsx, p1, p2, p3, p4
-
     analyze_btn.click(
-        run_analysis,
+        classify_zip_and_analyze_color,
         inputs=uploaded_zip,
         outputs=[output_df, image_gallery, download_file, plot1, plot2, plot3, plot4]
     )