diff --git a/deepface/models/face_detection/YuNet.py b/deepface/models/face_detection/YuNet.py index 9075927..e3d310f 100644 --- a/deepface/models/face_detection/YuNet.py +++ b/deepface/models/face_detection/YuNet.py @@ -1,6 +1,6 @@ # built-in dependencies import os -from typing import Any, List +from typing import Any, List, Union # 3rd party dependencies import cv2 @@ -57,10 +57,27 @@ class YuNetClient(Detector): ) from err return face_detector - def detect_faces(self, img: np.ndarray) -> List[FacialAreaRegion]: + def detect_faces(self, img: Union[np.ndarray, List[np.ndarray]]) -> Union[List[FacialAreaRegion], List[List[FacialAreaRegion]]]: """ Detect and align face with yunet + Args: + img (Union[np.ndarray, List[np.ndarray]]): pre-loaded image as numpy array or a list of those + + Returns: + results (Union[List[FacialAreaRegion], List[List[FacialAreaRegion]]]): A list or a list of lists of FacialAreaRegion objects + """ + if isinstance(img, np.ndarray): + return self._process_single_image(img) + elif isinstance(img, list): + return [self._process_single_image(single_img) for single_img in img] + else: + raise ValueError("Input must be a numpy array or a list of numpy arrays.") + + def _process_single_image(self, img: np.ndarray) -> List[FacialAreaRegion]: + """ + Helper function to detect faces in a single image. + Args: img (np.ndarray): pre-loaded image as numpy array