diff --git a/deepface/commons/weight_utils.py b/deepface/commons/weight_utils.py index d320702..dfac1fa 100644 --- a/deepface/commons/weight_utils.py +++ b/deepface/commons/weight_utils.py @@ -127,7 +127,7 @@ def download_all_models_in_one_shot() -> None: MODEL_URL as SSD_MODEL, WEIGHTS_URL as SSD_WEIGHTS, ) - from deepface.models.YoloModel import ( + from deepface.models.face_detection.Yolo import ( WEIGHT_URLS as YOLO_WEIGHTS, WEIGHT_NAMES as YOLO_WEIGHT_NAMES, YoloModel diff --git a/deepface/models/YoloClientBase.py b/deepface/models/YoloClientBase.py deleted file mode 100644 index 83ad324..0000000 --- a/deepface/models/YoloClientBase.py +++ /dev/null @@ -1,37 +0,0 @@ -# built-in dependencies -from typing import Any - -# project dependencies -from deepface.models.YoloModel import YoloModel, WEIGHT_URLS, WEIGHT_NAMES -from deepface.commons import weight_utils -from deepface.commons.logger import Logger - -logger = Logger() - - -class YoloClientBase: - def __init__(self, model: YoloModel): - self.model = self.build_model(model) - - def build_model(self, model: YoloModel) -> Any: - """ - Build a yolo detector model - Returns: - model (Any) - """ - - # Import the optional Ultralytics YOLO model - try: - from ultralytics import YOLO - except ModuleNotFoundError as e: - raise ImportError( - "Yolo is an optional detector, ensure the library is installed. " - "Please install using 'pip install ultralytics'" - ) from e - - weight_file = weight_utils.download_weights_if_necessary( - file_name=WEIGHT_NAMES[model.value], source_url=WEIGHT_URLS[model.value] - ) - - # Return face_detector - return YOLO(weight_file) diff --git a/deepface/models/YoloModel.py b/deepface/models/YoloModel.py deleted file mode 100644 index 93f2a74..0000000 --- a/deepface/models/YoloModel.py +++ /dev/null @@ -1,21 +0,0 @@ -from enum import Enum - - -class YoloModel(Enum): - V8N = 0 - V11N = 1 - V11S = 2 - V11M = 3 - - -# Model's weights paths -WEIGHT_NAMES = ["yolov8n-face.pt", - "yolov11n-face.pt", - "yolov11s-face.pt", - "yolov11m-face.pt"] - -# Google Drive URL from repo (https://github.com/derronqi/yolov8-face) ~6MB -WEIGHT_URLS = ["https://drive.google.com/uc?id=1qcr9DbgsX3ryrz2uU8w4Xm3cOrRywXqb", - "https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11n-face.pt", - "https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11s-face.pt", - "https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11m-face.pt"] diff --git a/deepface/models/face_detection/Yolo.py b/deepface/models/face_detection/Yolo.py index 1f29781..548e4f7 100644 --- a/deepface/models/face_detection/Yolo.py +++ b/deepface/models/face_detection/Yolo.py @@ -1,22 +1,66 @@ # built-in dependencies import os -from typing import List +from typing import List, Any +from enum import Enum # 3rd party dependencies import numpy as np # project dependencies -from deepface.models.YoloClientBase import YoloClientBase -from deepface.models.YoloModel import YoloModel from deepface.models.Detector import Detector, FacialAreaRegion from deepface.commons.logger import Logger +from deepface.commons import weight_utils logger = Logger() -class YoloDetectorClient(YoloClientBase, Detector): +class YoloModel(Enum): + V8N = 0 + V11N = 1 + V11S = 2 + V11M = 3 + + +# Model's weights paths +WEIGHT_NAMES = ["yolov8n-face.pt", + "yolov11n-face.pt", + "yolov11s-face.pt", + "yolov11m-face.pt"] + +# Google Drive URL from repo (https://github.com/derronqi/yolov8-face) ~6MB +WEIGHT_URLS = ["https://drive.google.com/uc?id=1qcr9DbgsX3ryrz2uU8w4Xm3cOrRywXqb", + "https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11n-face.pt", + "https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11s-face.pt", + "https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11m-face.pt"] + + +class YoloDetectorClient(Detector): def __init__(self, model: YoloModel): - super().__init__(model) + super().__init__() + self.model = self.build_model(model) + + def build_model(self, model: YoloModel) -> Any: + """ + Build a yolo detector model + Returns: + model (Any) + """ + + # Import the optional Ultralytics YOLO model + try: + from ultralytics import YOLO + except ModuleNotFoundError as e: + raise ImportError( + "Yolo is an optional detector, ensure the library is installed. " + "Please install using 'pip install ultralytics'" + ) from e + + weight_file = weight_utils.download_weights_if_necessary( + file_name=WEIGHT_NAMES[model.value], source_url=WEIGHT_URLS[model.value] + ) + + # Return face_detector + return YOLO(weight_file) def detect_faces(self, img: np.ndarray) -> List[FacialAreaRegion]: """ diff --git a/deepface/models/facial_recognition/Yolo.py b/deepface/models/facial_recognition/Yolo.py index e8f6d59..fb8054d 100644 --- a/deepface/models/facial_recognition/Yolo.py +++ b/deepface/models/facial_recognition/Yolo.py @@ -1,25 +1,68 @@ # built-in dependencies -from typing import List +from typing import List, Any +from enum import Enum # 3rd party dependencies import numpy as np # project dependencies -from deepface.models.YoloClientBase import YoloClientBase -from deepface.models.YoloModel import YoloModel from deepface.models.FacialRecognition import FacialRecognition from deepface.commons.logger import Logger +from deepface.commons import weight_utils logger = Logger() -class YoloFacialRecognitionClient(YoloClientBase, FacialRecognition): +class YoloModel(Enum): + V8N = 0 + V11N = 1 + V11S = 2 + V11M = 3 + + +# Model's weights paths +WEIGHT_NAMES = ["yolov8n-face.pt", + "yolov11n-face.pt", + "yolov11s-face.pt", + "yolov11m-face.pt"] + +# Google Drive URL from repo (https://github.com/derronqi/yolov8-face) ~6MB +WEIGHT_URLS = ["https://drive.google.com/uc?id=1qcr9DbgsX3ryrz2uU8w4Xm3cOrRywXqb", + "https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11n-face.pt", + "https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11s-face.pt", + "https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11m-face.pt"] + + +class YoloFacialRecognitionClient(FacialRecognition): def __init__(self, model: YoloModel): super().__init__(model) self.model_name = "Yolo" self.input_shape = None self.output_shape = 512 + def build_model(self, model: YoloModel) -> Any: + """ + Build a yolo detector model + Returns: + model (Any) + """ + + # Import the optional Ultralytics YOLO model + try: + from ultralytics import YOLO + except ModuleNotFoundError as e: + raise ImportError( + "Yolo is an optional detector, ensure the library is installed. " + "Please install using 'pip install ultralytics'" + ) from e + + weight_file = weight_utils.download_weights_if_necessary( + file_name=WEIGHT_NAMES[model.value], source_url=WEIGHT_URLS[model.value] + ) + + # Return face_detector + return YOLO(weight_file) + def forward(self, img: np.ndarray) -> List[float]: return self.model.embed(img)[0].tolist()