Update app.py

app.py

@@ -11,8 +11,8 @@ import gradio as gr
 # ------------------------------
 # MODEL REPOSITORIES
 # ------------------------------
-YOLO_REPO = "arnabdhar/YOLOv8-Face-Detection"   # Face detection
-ARCFACE_ONNX_REPO = "garavv/arcface-onnx"       # Face embeddings
+YOLO_REPO = "arnabdhar/YOLOv8-Face-Detection"   # Face detection model
+ARCFACE_ONNX_REPO = "garavv/arcface-onnx"       # ArcFace ONNX model
 
 # ------------------------------
 # DOWNLOAD MODELS FROM HUGGING FACE
@@ -40,21 +40,23 @@ arcface_input = arcface_sess.get_inputs()[0].name
 # ------------------------------
 # HELPER FUNCTIONS
 # ------------------------------
-def preprocess_face(img):
-    img = cv2.resize(img, (112, 112))
+
+def get_embedding(face):
+    """
+    Preprocess the face and compute embedding using ArcFace ONNX model.
+    Fixed to match NHWC format (1, 112, 112, 3).
+    """
+    img = cv2.resize(face, (112, 112))
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     img = (img - 127.5) / 128.0
-    img = np.transpose(img, (2, 0, 1))[None, ...].astype(np.float32)
-    return img
+    img = np.expand_dims(img.astype(np.float32), axis=0)  # (1, 112, 112, 3)
 
-def get_embedding(face):
-    x = preprocess_face(face)
-    emb = arcface_sess.run(None, {arcface_input: x})[0][0]
+    emb = arcface_sess.run(None, {arcface_input: img})[0][0]
     emb = emb / np.linalg.norm(emb)
     return emb
 
 # ------------------------------
-# KNOWN FACES DATABASE
+# KNOWN FACES DATABASE (in memory)
 # ------------------------------
 known_faces = {}
 
@@ -78,7 +80,7 @@ def register_face(name, image):
         else:
             img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
 
-        # Try computing embedding
+        # Compute embedding
         emb = get_embedding(img)
         if emb is None or np.isnan(emb).any():
             return f"❌ Failed to compute embedding for {name}. Try another image."