Update app.py

app.py

@@ -4,10 +4,31 @@ import cv2
 from ultralytics import YOLO
 import insightface
 
-# Load models
-yolo = YOLO("yolov8n.pt")
+# ----------------------------
+# Load Models (CPU mode)
+# ----------------------------
+
+yolo = YOLO("yolov8n.pt")  # lightweight model
+
 face_model = insightface.app.FaceAnalysis(name="buffalo_l")
-face_model.prepare(ctx_id=0)
+face_model.prepare(ctx_id=-1)  # -1 forces CPU (important for HF free tier)
+
+
+# ----------------------------
+# Utility: Normalize embedding
+# ----------------------------
+
+def normalize(vec):
+    vec = np.array(vec, dtype=np.float32)
+    norm = np.linalg.norm(vec)
+    if norm == 0:
+        return vec.tolist()
+    return (vec / norm).tolist()
+
+
+# ----------------------------
+# Main Processing Function
+# ----------------------------
 
 def process_image(image):
     image_np = np.array(image)
@@ -17,41 +38,61 @@ def process_image(image):
     faces_output = []
 
     for r in results:
-        for box, cls, conf in zip(r.boxes.xyxy, r.boxes.cls, r.boxes.conf):
+        boxes = r.boxes
 
-            # Only detect persons
+        for box, cls, conf in zip(boxes.xyxy, boxes.cls, boxes.conf):
+
+            # YOLO class 0 = person
             if int(cls) != 0:
                 continue
-            if conf < 0.4:
+
+            if float(conf) < 0.4:
                 continue
 
             xmin, ymin, xmax, ymax = box.cpu().numpy()
             xmin, ymin, xmax, ymax = map(int, [xmin, ymin, xmax, ymax])
 
+            # Safety check for valid crop
+            h, w, _ = image_np.shape
+            xmin = max(0, xmin)
+            ymin = max(0, ymin)
+            xmax = min(w, xmax)
+            ymax = min(h, ymax)
+
             person_crop = image_np[ymin:ymax, xmin:xmax]
 
+            if person_crop.size == 0:
+                continue
+
+            # Detect face inside person crop
             detected_faces = face_model.get(person_crop)
 
             for face in detected_faces:
-                embedding = face.embedding.tolist()
+                embedding = normalize(face.embedding)
 
-                cx = (xmin + xmax) / 2
-                cy = (ymin + ymax) / 2
+                # Adjust face bbox to original image coordinates
+                fxmin, fymin, fxmax, fymax = face.bbox.astype(int)
 
                 faces_output.append({
-                    "cx": float(cx),
-                    "cy": float(cy),
+                    "cx": float((fxmin + fxmax) / 2 + xmin),
+                    "cy": float((fymin + fymax) / 2 + ymin),
+                    "confidence": float(conf),
                     "box": {
-                        "xmin": xmin,
-                        "ymin": ymin,
-                        "xmax": xmax,
-                        "ymax": ymax
+                        "xmin": int(fxmin + xmin),
+                        "ymin": int(fymin + ymin),
+                        "xmax": int(fxmax + xmin),
+                        "ymax": int(fymax + ymin)
                     },
                     "embedding": embedding
                 })
 
     return faces_output
 
+
+# ----------------------------
+# Gradio Interface
+# ----------------------------
+
 iface = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),