Update app.py

app.py

@@ -8,13 +8,18 @@ import mediapipe as mp
 from collections import defaultdict
 
 # -------------------- CONFIG --------------------
-skin_model_path = "pores.pt"  # replace with your model in HF Space
-conf_threshold = 0.05
+model_paths = {
+    "pores": "pores.pt",
+    "newpig": "newpig.pt",
+    "wrinkle": "wrinkle.pt"
+}
+default_conf_threshold = 0.05
 imgsz = 1024
 
 # -------------------- INIT --------------------
 app = Flask(__name__)
-skin_model = YOLO(skin_model_path)
+models = {name: YOLO(path) for name, path in model_paths.items()}
+
 mp_face_mesh = mp.solutions.face_mesh
 face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True)
 
@@ -55,35 +60,42 @@ def analyze():
             hull = cv2.convexHull(points)
             face_area = cv2.contourArea(hull)
             face_polygon = Polygon(hull.reshape(-1, 2))
-            break  # only first face
+            break  # only process the first face
     else:
         return jsonify({"success": False, "error": "No face detected"}), 400
 
-    # -------------------- SKIN DETECTION --------------------
-    results = skin_model(img, conf=conf_threshold, imgsz=imgsz)
-    boxes_xy = results[0].boxes.xyxy.cpu().numpy()
-    boxes_conf = results[0].boxes.conf.cpu().numpy()
-    boxes_cls = results[0].boxes.cls.cpu().numpy().astype(int)
-
-    class_polygons = defaultdict(list)
-
-    for i, (xy, cls_id) in enumerate(zip(boxes_xy, boxes_cls)):
-        x1, y1, x2, y2 = xy.astype(int)
-        det_poly = shapely_box(x1, y1, x2, y2)
-        if face_polygon.intersects(det_poly):
-            intersection = det_poly.intersection(face_polygon)
-            if intersection.area > 0:
-                class_polygons[cls_id].append(intersection)
-
-    # -------------------- CALCULATE PERCENTAGES --------------------
-    skin_percentages = {}
-    for cls_id, polys in class_polygons.items():
-        union_poly = unary_union(polys)
-        pixels = union_poly.area
-        percentage = (pixels / face_area) * 100 if face_area > 0 else 0.0
-        skin_percentages[skin_model.names.get(cls_id, str(cls_id))] = round(percentage, 2)
-
-    return jsonify({"success": True, "percentages": skin_percentages})
+    # -------------------- RUN ALL MODELS --------------------
+    all_percentages = {}
+
+    for model_name, model in models.items():
+        # Set confidence threshold per model
+        conf = 0.03 if model_name == "newpig" else default_conf_threshold
+
+        results = model(img, conf=conf, imgsz=imgsz)
+        boxes_xy = results[0].boxes.xyxy.cpu().numpy()
+        boxes_cls = results[0].boxes.cls.cpu().numpy().astype(int)
+
+        class_polygons = defaultdict(list)
+
+        for i, cls_id in enumerate(boxes_cls):
+            x1, y1, x2, y2 = boxes_xy[i].astype(int)
+            det_poly = shapely_box(x1, y1, x2, y2)
+            if face_polygon.intersects(det_poly):
+                intersection = det_poly.intersection(face_polygon)
+                if intersection.area > 0:
+                    class_polygons[cls_id].append(intersection)
+
+        # Calculate percentages
+        skin_percentages = {}
+        for cls_id, polys in class_polygons.items():
+            union_poly = unary_union(polys)
+            pixels = union_poly.area
+            percentage = (pixels / face_area) * 100 if face_area > 0 else 0.0
+            skin_percentages[model.names.get(cls_id, str(cls_id))] = round(percentage, 2)
+
+        all_percentages[model_name] = skin_percentages
+
+    return jsonify({"success": True, "percentages": all_percentages})
 
 # -------------------- RUN --------------------
 if __name__ == "__main__":