diff --git a/deepface/DeepFace.py b/deepface/DeepFace.py index 0f41a4e..5848d7b 100644 --- a/deepface/DeepFace.py +++ b/deepface/DeepFace.py @@ -48,13 +48,19 @@ if tf_version == 2: folder_utils.initialize_folder() -def build_model(task: str, model_name: str) -> Any: +def build_model(model_name: str, task: str = "facial_recognition") -> Any: """ - This function builds a deepface model + This function builds a pre-trained model Args: - model_name (string): face recognition, facial attribute or face detector model - VGG-Face, Facenet, OpenFace, DeepFace, DeepID for face recognition - Age, Gender, Emotion, Race for facial attributes + model_name (str): model identifier + - VGG-Face, Facenet, Facenet512, OpenFace, DeepFace, DeepID, Dlib, + ArcFace, SFace, GhostFaceNet for face recognition + - Age, Gender, Emotion, Race for facial attributes + - opencv, mtcnn, ssd, dlib, retinaface, mediapipe, yolov8, yunet, + fastmtcnn or centerface for face detectors + - Fasnet for spoofing + task (str): facial_recognition, facial_attribute, face_detector, spoofing + default is facial_recognition Returns: built_model """ diff --git a/deepface/modules/modeling.py b/deepface/modules/modeling.py index 250102b..d6cf6a9 100644 --- a/deepface/modules/modeling.py +++ b/deepface/modules/modeling.py @@ -97,6 +97,6 @@ def build_model(task: str, model_name: str) -> Any: if model: model_obj[task][model_name] = model() else: - raise ValueError(f"Invalid model_name passed - {model_name}") + raise ValueError(f"Invalid model_name passed - {task}/{model_name}") return model_obj[task][model_name] diff --git a/tests/face-recognition-how.py b/tests/face-recognition-how.py index f33e1af..cb26f2f 100644 --- a/tests/face-recognition-how.py +++ b/tests/face-recognition-how.py @@ -16,7 +16,7 @@ logger = Logger() model_name = "VGG-Face" -model: FacialRecognition = DeepFace.build_model(model_name=model_name) +model: FacialRecognition = DeepFace.build_model(task="facial_recognition", model_name=model_name) target_size = model.input_shape