Update handler.py

handler.py

@@ -154,15 +154,26 @@ class EndpointHandler:
         image = image[np.newaxis, :, :, :].astype(np.float32) / 127.5 - 1.0  # Normalize to [-1, 1]
         return image
 
-    def get_face_embedding(self, image):
+    def get_face_info(self, image):
         # Preprocess the image
         image = self.preprocess(image)
 
-        # Run the ONNX model to get the face embedding
+        # Run the ONNX model to get the face detection results
         input_name = self.ort_session.get_inputs()[0].name
-        embedding = self.ort_session.run(None, {input_name: image})[0]
-
-        return embedding
+        outputs = self.ort_session.run(None, {input_name: image})
+
+        # Process the output to extract face information
+        bboxes = outputs[0][0]  # Adjust based on model output structure
+        face_info_list = []
+        for bbox in bboxes:
+            score = bbox[2]
+            if score > 0.5:  # Confidence threshold
+                x1, y1, x2, y2 = bbox[3:7] * [320, 240, 320, 240]  # Scale coordinates
+                face_info_list.append({
+                    "bbox": [x1, y1, x2, y2],
+                    "embedding": self.get_face_embedding(image[:, :, int(y1):int(y2), int(x1):int(x2)])
+                })
+        return face_info_list
 
     def __call__(self, data):
 
@@ -251,17 +262,30 @@ class EndpointHandler:
         height, width, _ = face_image_cv2.shape
 
         # Extract face features using the ONNX model
-        face_emb = self.get_face_embedding(face_image_cv2)
+        face_info_list = self.get_face_info(face_image_cv2)
+
+        if len(face_info_list) == 0:
+            return {"error": "No faces detected."}
 
+        # Use the largest face detected
+        face_info = max(face_info_list, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))
+        face_emb = face_info["embedding"]
         face_kps = draw_kps(convert_from_cv2_to_image(face_image_cv2), face_info["kps"])
         img_controlnet = face_image
+
         if pose_image:
             pose_image = resize_img(pose_image, max_side=1024)
             img_controlnet = pose_image
             pose_image_cv2 = convert_from_image_to_cv2(pose_image)
 
             # Extract face features from pose image using the ONNX model
-            face_emb = self.get_face_embedding(pose_image_cv2)
+            face_info_list = self.get_face_info(pose_image_cv2)
+
+            if len(face_info_list) == 0:
+                return {"error": "No faces detected in pose image."}
+
+            face_info = max(face_info_list, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))
+            face_emb = face_info["embedding"]
             face_kps = draw_kps(pose_image, face_info["kps"])
 
             width, height = face_kps.size