From 01d7fa2272a10a9dabf1eb5cdb67be8c67ef6517 Mon Sep 17 00:00:00 2001 From: Daniel Dunderfelt Date: Mon, 24 Jun 2024 04:20:37 +0300 Subject: [PATCH] Fix face detect when no faces are detected. Install only onnxruntime-gpu, otherwise CUDA doesn't work. (#85) --- hallo/datasets/image_processor.py | 16 ++++++++++++---- requirements.txt | 2 +- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/hallo/datasets/image_processor.py b/hallo/datasets/image_processor.py index 1a824220..2c093907 100644 --- a/hallo/datasets/image_processor.py +++ b/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. diff --git a/requirements.txt b/requirements.txt index 5b06d280..5c4a9a58 100644 --- a/requirements.txt +++ b/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