yunet pseudobatching

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