Refactoring and requirements.

deepface/detectors/FaceDetector.py

@@ -9,7 +9,7 @@ from deepface.detectors import (
     MtcnnWrapper,
     RetinaFaceWrapper,
     MediapipeWrapper,
-    Yolov8faceWrapper,
+    YoloWrapper,
 )
 
 
@@ -23,7 +23,7 @@ def build_model(detector_backend):
         "mtcnn": MtcnnWrapper.build_model,
         "retinaface": RetinaFaceWrapper.build_model,
         "mediapipe": MediapipeWrapper.build_model,
-        "yolov8n": Yolov8faceWrapper.build_model("yolov8n"),
+        "yolov8n": YoloWrapper.build_model,
     }
 
     if not "face_detector_obj" in globals():
@@ -66,7 +66,7 @@ def detect_faces(face_detector, detector_backend, img, align=True):
         "mtcnn": MtcnnWrapper.detect_face,
         "retinaface": RetinaFaceWrapper.detect_face,
         "mediapipe": MediapipeWrapper.detect_face,
-        "yolov8n": Yolov8faceWrapper.detect_face,
+        "yolov8n": YoloWrapper.detect_face,
     }
 
     detect_face_fn = backends.get(detector_backend)

deepface/detectors/YoloWrapper.py

@@ -0,0 +1,61 @@
+from deepface.detectors import FaceDetector
+
+# Model's weights paths
+PATH = "/.deepface/weights/yolov8n-face.pt"
+
+# Google Drive URL
+WEIGHT_URL = "https://drive.google.com/uc?id=1qcr9DbgsX3ryrz2uU8w4Xm3cOrRywXqb"
+
+# Confidence thresholds for landmarks detection
+# used in alignment_procedure function
+LANDMARKS_CONFIDENCE_THRESHOLD = 0.5
+
+
+def build_model():
+    """Build YOLO (yolov8n-face) model"""
+    import gdown
+    import os
+
+    # Import the Ultralytics YOLO model
+    from ultralytics import YOLO
+
+    from deepface.commons.functions import get_deepface_home
+    weight_path = f"{get_deepface_home()}{PATH}"
+
+    # Download the model's weights if they don't exist
+    if not os.path.isfile(weight_path):
+        gdown.download(WEIGHT_URL, weight_path, quiet=False)
+        print(f"Downloaded YOLO model {os.path.basename(weight_path)}")
+
+    # Return face_detector
+    return YOLO(weight_path)
+
+
+def detect_face(face_detector, img, align=False):
+    resp = []
+
+    # Detect faces
+    results = face_detector.predict(
+        img, verbose=False, show=False, conf=0.25)[0]
+
+    # For each face, extract the bounding box, the landmarks and confidence
+    for result in results:
+        # Extract the bounding box and the confidence
+        x, y, w, h = result.boxes.xywh.tolist()[0]
+        confidence = result.boxes.conf.tolist()[0]
+
+        x, y, w, h = int(x - w / 2), int(y - h / 2), int(w), int(h)
+        detected_face = img[y: y + h, x: x + w].copy()
+
+        if align:
+            # Extract landmarks
+            left_eye, right_eye, _, _, _ = result.keypoints.tolist()
+            # Check the landmarks confidence before alignment
+            if left_eye[2] > LANDMARKS_CONFIDENCE_THRESHOLD and right_eye[2] > LANDMARKS_CONFIDENCE_THRESHOLD:
+                detected_face = FaceDetector.alignment_procedure(
+                    detected_face, left_eye[:2], right_eye[:2]
+                )
+
+        resp.append((detected_face, [x, y, w, h], confidence))
+
+    return resp

deepface/detectors/Yolov8faceWrapper.py

@@ -1,71 +0,0 @@
-from deepface.detectors import FaceDetector
-
-# Models names and paths
-PATHS = {
-    "yolov8n": "/.deepface/weights/yolov8n-face.pt",
-}
-
-# Google Drive base URL
-BASE_URL = "https://drive.google.com/uc?id="
-
-# Models' Google Drive IDs
-IDS = {
-    "yolov8n": "1qcr9DbgsX3ryrz2uU8w4Xm3cOrRywXqb",
-}
-
-
-def build_model(model: str):
-    """Function factory for YOLO models"""
-    from deepface.commons.functions import get_deepface_home
-
-    # Get model's weights path and Google Drive URL
-    func_weights_path = f"{get_deepface_home()}{PATHS[model]}"
-    func_url = f"{BASE_URL}{IDS[model]}"
-
-    # Define function to build the model
-    def _build_model(weights_path: str = func_weights_path, url: str = func_url):
-        import gdown
-        import os
-
-        # Import the Ultralytics YOLO model
-        from ultralytics import YOLO
-
-        # Download the model's weights if they don't exist
-        if not os.path.isfile(weights_path):
-            gdown.download(url, weights_path, quiet=False)
-            print(f"Downloaded YOLO model {os.path.basename(PATHS[model])}")
-
-        # Return face_detector
-        return YOLO(weights_path)
-
-    return _build_model
-
-
-def detect_face(face_detector, img, align=False):
-    resp = []
-
-    # Detect faces
-    results = face_detector.predict(
-        img, verbose=False, show=False, conf=0.25)[0]
-
-    # For each face, extract the bounding box, the landmarks and confidence
-    for result in results:
-        # Extract the bounding box and the confidence
-        x, y, w, h = result.boxes.xywh.tolist()[0]
-        confidence = result.boxes.conf.tolist()[0]
-
-        x, y, w, h = int(x - w / 2), int(y - h / 2), int(w), int(h)
-        detected_face = img[y: y + h, x: x + w].copy()
-
-        if align:
-            # Extract landmarks
-            left_eye, right_eye, _, _, _ = result.keypoints.tolist()
-            # Check the landmarks confidence before alignment
-            if left_eye[2] > 0.5 and right_eye[2] > 0.5:
-                detected_face = FaceDetector.alignment_procedure(
-                    detected_face, left_eye[:2], right_eye[:2]
-                )
-
-        resp.append((detected_face, [x, y, w, h], confidence))
-
-    return resp

requirements_additional.txt

@@ -1,3 +1,4 @@
 opencv-contrib-python>=4.3.0.36
 mediapipe>=0.8.7.3
 dlib>=19.20.0
+ultralytics @ git+https://github.com/derronqi/yolov8-face.git@b623989575bdb78601b5ca717851e3d63ca9e01c