Fix face detect when no faces are detected. Install only onnxruntime-…

…gpu, otherwise CUDA doesn't work. (#85)

hallo/datasets/image_processor.py

@@ -117,11 +117,19 @@ def preprocess(self, source_image_path: str, cache_dir: str, face_region_ratio:
         # 1. image augmentation
         pixel_values_ref_img = self._augmentation(ref_image_pil, self.pixel_transform)
 
-
         # 2.1 detect face
         faces = self.face_analysis.get(cv2.cvtColor(np.array(ref_image_pil.copy()), cv2.COLOR_RGB2BGR))
-        # use max size face
-        face = sorted(faces, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))[-1]
+        if not faces:
+            print("No faces detected in the image. Using the entire image as the face region.")
+            # Use the entire image as the face region
+            face = {
+                "bbox": [0, 0, ref_image_pil.width, ref_image_pil.height],
+                "embedding": np.zeros(512)
+            }
+        else:
+            # Sort faces by size and select the largest one
+            faces_sorted = sorted(faces, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]), reverse=True)
+            face = faces_sorted[0]  # Select the largest face
 
         # 2.2 face embedding
         face_emb = face["embedding"]
@@ -173,7 +181,7 @@ def preprocess(self, source_image_path: str, cache_dir: str, face_region_ratio:
     def close(self):
         """
         Closes the ImageProcessor and releases any resources held by the FaceAnalysis instance.
-        
+
         Args:
             self: The ImageProcessor instance.
 

requirements.txt

@@ -16,7 +16,7 @@ numpy==1.26.4
 omegaconf==2.3.0
 onnx2torch==1.5.14
 onnx==1.16.1
-onnxruntime==1.18.0
+onnxruntime-gpu==1.18.0
 opencv-contrib-python==4.9.0.80
 opencv-python-headless==4.9.0.80
 opencv-python==4.9.0.80