fix:YoloModel removal

This commit is contained in:
roberto-corno-nttdata 2024-12-09 11:51:56 +01:00
parent aa57d879b6
commit 01bf48dff8
5 changed files with 97 additions and 68 deletions

View File

@ -127,7 +127,7 @@ def download_all_models_in_one_shot() -> None:
MODEL_URL as SSD_MODEL, MODEL_URL as SSD_MODEL,
WEIGHTS_URL as SSD_WEIGHTS, WEIGHTS_URL as SSD_WEIGHTS,
) )
from deepface.models.YoloModel import ( from deepface.models.face_detection.Yolo import (
WEIGHT_URLS as YOLO_WEIGHTS, WEIGHT_URLS as YOLO_WEIGHTS,
WEIGHT_NAMES as YOLO_WEIGHT_NAMES, WEIGHT_NAMES as YOLO_WEIGHT_NAMES,
YoloModel YoloModel

View File

@ -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)

View File

@ -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"]

View File

@ -1,22 +1,66 @@
# built-in dependencies # built-in dependencies
import os import os
from typing import List from typing import List, Any
from enum import Enum
# 3rd party dependencies # 3rd party dependencies
import numpy as np import numpy as np
# project dependencies # project dependencies
from deepface.models.YoloClientBase import YoloClientBase
from deepface.models.YoloModel import YoloModel
from deepface.models.Detector import Detector, FacialAreaRegion from deepface.models.Detector import Detector, FacialAreaRegion
from deepface.commons.logger import Logger from deepface.commons.logger import Logger
from deepface.commons import weight_utils
logger = Logger() 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): 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]: def detect_faces(self, img: np.ndarray) -> List[FacialAreaRegion]:
""" """

View File

@ -1,25 +1,68 @@
# built-in dependencies # built-in dependencies
from typing import List from typing import List, Any
from enum import Enum
# 3rd party dependencies # 3rd party dependencies
import numpy as np import numpy as np
# project dependencies # project dependencies
from deepface.models.YoloClientBase import YoloClientBase
from deepface.models.YoloModel import YoloModel
from deepface.models.FacialRecognition import FacialRecognition from deepface.models.FacialRecognition import FacialRecognition
from deepface.commons.logger import Logger from deepface.commons.logger import Logger
from deepface.commons import weight_utils
logger = Logger() 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): def __init__(self, model: YoloModel):
super().__init__(model) super().__init__(model)
self.model_name = "Yolo" self.model_name = "Yolo"
self.input_shape = None self.input_shape = None
self.output_shape = 512 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]: def forward(self, img: np.ndarray) -> List[float]:
return self.model.embed(img)[0].tolist() return self.model.embed(img)[0].tolist()