Update app.py

app.py

@@ -1,45 +1,59 @@
 import gradio as gr
 import numpy as np
+import cv2
 from ultralytics import YOLO
+import insightface
 
-model = YOLO("yolov8n.pt")
+# Load models
+yolo = YOLO("yolov8n.pt")
+face_model = insightface.app.FaceAnalysis(name="buffalo_l")
+face_model.prepare(ctx_id=0)
 
-def detect_persons(image):
+def process_image(image):
     image_np = np.array(image)
-    results = model(image_np)
 
-    persons = []
+    results = yolo(image_np)
+
+    faces_output = []
 
     for r in results:
         for box, cls, conf in zip(r.boxes.xyxy, r.boxes.cls, r.boxes.conf):
 
-            # class 0 in COCO = person
+            # Only detect persons
             if int(cls) != 0:
                 continue
-
             if conf < 0.4:
                 continue
 
             xmin, ymin, xmax, ymax = box.cpu().numpy()
-            cx = (xmin + xmax) / 2
-            cy = (ymin + ymax) / 2
-
-            persons.append({
-                "cx": float(cx),
-                "cy": float(cy),
-                "confidence": float(conf),
-                "box": {
-                    "xmin": float(xmin),
-                    "ymin": float(ymin),
-                    "xmax": float(xmax),
-                    "ymax": float(ymax)
-                }
-            })
-
-    return persons
+            xmin, ymin, xmax, ymax = map(int, [xmin, ymin, xmax, ymax])
+
+            person_crop = image_np[ymin:ymax, xmin:xmax]
+
+            detected_faces = face_model.get(person_crop)
+
+            for face in detected_faces:
+                embedding = face.embedding.tolist()
+
+                cx = (xmin + xmax) / 2
+                cy = (ymin + ymax) / 2
+
+                faces_output.append({
+                    "cx": float(cx),
+                    "cy": float(cy),
+                    "box": {
+                        "xmin": xmin,
+                        "ymin": ymin,
+                        "xmax": xmax,
+                        "ymax": ymax
+                    },
+                    "embedding": embedding
+                })
+
+    return faces_output
 
 iface = gr.Interface(
-    fn=detect_persons,
+    fn=process_image,
     inputs=gr.Image(type="pil"),
     outputs="json"
 )