.gitignore

@@ -1,5 +1,6 @@
 data
 .idea
+.yoloface
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

face_detect.py

@@ -1,5 +1,8 @@
 import cv2
 import numpy as np
+from yoloface import face_analysis
+
+face_detector = face_analysis()
 
 
 def equalize_image(data: np.ndarray):
@@ -8,6 +11,13 @@ def equalize_image(data: np.ndarray):
   return cv2.cvtColor(data_hsv, cv2.COLOR_HSV2RGB)
 
 
+def find_face_bbox_yolo(data: np.ndarray):
+  _, box, conf = face_detector.face_detection(frame_arr=data, frame_status=True, model='full')
+  if len(box) < 1:
+    return None, None
+  return box, conf
+
+
 def find_face_bbox(data: np.ndarray):
   classifier_files = [
     'haarcascades/haarcascade_frontalface_default.xml',
@@ -19,16 +29,26 @@ def find_face_bbox(data: np.ndarray):
   ]
   data_equalized = equalize_image(data)
   data_gray = cv2.cvtColor(data_equalized, cv2.COLOR_RGB2GRAY)
-  face_coords = None
+  face_coords, conf = find_face_bbox_yolo(cv2.cvtColor(data_equalized, cv2.COLOR_RGB2BGR))
+  if face_coords is not None:
+    return face_coords[0]
+
   for classifier in classifier_files:
     face_cascade = cv2.CascadeClassifier(classifier)
     face_coords = face_cascade.detectMultiScale(data_gray, 1.1, 3)
     if face_coords is not None:
       break
-  return max(face_coords, key=len)
+  return max(face_coords, key=lambda v: v[2]*v[3])
 
 
 def crop_face(data: np.ndarray, bounding_box) -> np.ndarray:
   x, y, w, h = bounding_box
+  # Extending the boxes
+  factor = 0.4
+  x, y = round(x - factor * w), round(y - factor * h)
+  w, h = round(w + factor * w * 2), round(h + factor * h * 2)
+  y = max(y, 0)
+  x = max(x, 0)
+
   face = data[y:y + h, x:x + w]
   return face

requirements.txt

@@ -8,4 +8,6 @@ lxml==4.9.2
 opencv-python==4.7.0.68
 torch==1.13.1
 matplotlib==3.6.3
-scikit-image==0.19.3
+scikit-image==0.19.3
+yoloface==0.0.4
+ipython==8.9.0