Add support for face extraction

model/fecnet.py

@@ -5,6 +5,7 @@ from importlib.util import module_from_spec, spec_from_file_location
 from tempfile import TemporaryDirectory
 
 import cv2
+import mediapipe as mp
 import numpy as np
 import plotly.express as px
 import requests
@@ -40,6 +41,10 @@ class FECNetModel:
             self.__download_weights(repo_dir.name), fecnet_module.FECNet
         )
 
+        self.face_detector = mp.solutions.face_detection.FaceDetection(
+            min_detection_confidence=0.5
+        )
+
     def __download_weights(self, model_dir: str) -> str:
         model_path = hf_hub_download(
             "natexcvi/pretrained-fecnet",
@@ -62,12 +67,48 @@ class FECNetModel:
     def distance(a, b):
         return np.linalg.norm(a - b)
 
-    def embed_image(self, image) -> np.ndarray:
+    def embed_image(self, image, crop_face: bool = False) -> np.ndarray:
         image = cv2.imdecode(image, cv2.IMREAD_COLOR)
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        if crop_face:
+            image = self.extract_face(image)
+        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         image = cv2.resize(image, (224, 224))
         image = np.transpose(image, (2, 0, 1))
         image = np.expand_dims(image, axis=0)
         image = torch.from_numpy(image.astype(np.float32))
         pred = self.predict(image)
         return pred.detach().numpy()
+
+    def extract_face(self, image):
+        mp_face_detection = mp.solutions.face_detection
+
+        # Convert the image to RGB
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        # Initialize the face detection model
+
+        # Run the face detection model on the image
+        results = self.face_detection.process(image)
+        # If a face is detected, crop the image to the face box
+        if results.detections:
+            for detection in results.detections:
+                x, y, w, h = (
+                    int(
+                        detection.location_data.relative_bounding_box.xmin
+                        * image.shape[1]
+                    ),
+                    int(
+                        detection.location_data.relative_bounding_box.ymin
+                        * image.shape[0]
+                    ),
+                    int(
+                        detection.location_data.relative_bounding_box.width
+                        * image.shape[1]
+                    ),
+                    int(
+                        detection.location_data.relative_bounding_box.height
+                        * image.shape[0]
+                    ),
+                )
+                cropped_image = image[y : y + h, x : x + w]
+                return cv2.cvtColor(cropped_image, cv2.COLOR_RGB2BGR)

routers/fecnet_router.py

@@ -23,9 +23,10 @@ model = FECNetModel(os.getenv("HF_TOKEN", ""))
 )
 async def calculate_embedding(
     image: UploadFile = File(...),
+    should_extract_face: bool = False,
 ):
     image_arr = np.asarray(bytearray(await image.read()), dtype=np.uint8) # type: ignore
-    rep = model.embed_image(image_arr)
+    rep = model.embed_image(image_arr, should_extract_face)
     return EmbeddingResponse(embedding=rep.tolist())
 
 
@@ -37,9 +38,10 @@ async def calculate_embedding(
 async def calculate_similarity_score(
     image1: UploadFile = File(...),
     image2: UploadFile = File(...),
+    should_extract_face: bool = False,
 ):
     image1_arr = np.asarray(bytearray(await image1.read()), dtype=np.uint8) # type: ignore
     image2_arr = np.asarray(bytearray(await image2.read()), dtype=np.uint8) # type: ignore
-    rep1 = model.embed_image(image1_arr)
-    rep2 = model.embed_image(image2_arr)
+    rep1 = model.embed_image(image1_arr, should_extract_face)
+    rep2 = model.embed_image(image2_arr, should_extract_face)
     return SimilarityResponse(score=np.linalg.norm(rep1, rep2))