Update roop/face_analyser.py

roop/face_analyser.py

@@ -1,5 +1,5 @@
 import threading
-from typing import Any
+from typing import Any, List, Optional
 import insightface
 
 import roop.globals
@@ -10,25 +10,56 @@ THREAD_LOCK = threading.Lock()
 
 
 def get_face_analyser() -> Any:
+    """
+    Initialize and return the face analyser.
+
+    This function ensures that the face analyser is initialized only once,
+    using a thread-safe mechanism.
+
+    Returns:
+        Any: An instance of the FaceAnalysis class from insightface.
+    """
     global FACE_ANALYSER
 
     with THREAD_LOCK:
         if FACE_ANALYSER is None:
-            FACE_ANALYSER = insightface.app.FaceAnalysis(name='buffalo_l', providers=roop.globals.execution_providers)
+            FACE_ANALYSER = insightface.app.FaceAnalysis(
+                name='buffalo_l',
+                providers=roop.globals.execution_providers
+            )
             FACE_ANALYSER.prepare(ctx_id=0, det_size=(640, 640))
     return FACE_ANALYSER
 
 
-def get_one_face(frame: Frame) -> Any:
-    face = get_face_analyser().get(frame)
-    try:
-        return min(face, key=lambda x: x.bbox[0])
-    except ValueError:
+def get_one_face(frame: Frame) -> Optional[Any]:
+    """
+    Extract the face with the smallest bounding box from the given frame.
+
+    Args:
+        frame (Frame): The image or video frame to analyze.
+
+    Returns:
+        Optional[Any]: The face with the smallest bounding box or None if no face is found.
+    """
+    faces = get_face_analyser().get(frame)
+    if not faces:
         return None
 
+    return min(faces, key=lambda x: x.bbox[0])
 
-def get_many_faces(frame: Frame) -> Any:
+
+def get_many_faces(frame: Frame) -> List[Any]:
+    """
+    Extract all detected faces from the given frame.
+
+    Args:
+        frame (Frame): The image or video frame to analyze.
+
+    Returns:
+        List[Any]: A list of detected faces or an empty list if no faces are found.
+    """
     try:
         return get_face_analyser().get(frame)
-    except IndexError:
-        return None
+    except Exception as e:
+        print(f"Error getting faces: {e}")
+        return []