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,
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

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

View File

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