Update roop/face_analyser.py

roop/face_analyser.py

@@ -1,34 +1,55 @@
-import threading
-from typing import Any
-import insightface
-
-import roop.globals
-from roop.typing import Frame
-
-FACE_ANALYSER = None
-THREAD_LOCK = threading.Lock()
-
-
-def get_face_analyser() -> Any:
-    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.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:
-        return None
-
-
-def get_many_faces(frame: Frame) -> Any:
-    try:
-        return get_face_analyser().get(frame)
-    except IndexError:
-        return None
+import threading
+from typing import Any, Optional, List
+import insightface
+import numpy
+import spaces
+
+import roop.globals
+from roop.typing import Frame, Face
+
+FACE_ANALYSER = None
+THREAD_LOCK = threading.Lock()
+
+@spaces.GPU()
+def get_face_analyser() -> Any:
+    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.prepare(ctx_id=0)
+    return FACE_ANALYSER
+
+
+def clear_face_analyser() -> Any:
+    global FACE_ANALYSER
+
+    FACE_ANALYSER = None
+
+
+def get_one_face(frame: Frame, position: int = 0) -> Optional[Face]:
+    many_faces = get_many_faces(frame)
+    if many_faces:
+        try:
+            return many_faces[position]
+        except IndexError:
+            return many_faces[-1]
+    return None
+
+
+def get_many_faces(frame: Frame) -> Optional[List[Face]]:
+    try:
+        return get_face_analyser().get(frame)
+    except ValueError:
+        return None
+
+
+def find_similar_face(frame: Frame, reference_face: Face) -> Optional[Face]:
+    many_faces = get_many_faces(frame)
+    if many_faces:
+        for face in many_faces:
+            if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'):
+                distance = numpy.sum(numpy.square(face.normed_embedding - reference_face.normed_embedding))
+                if distance < roop.globals.similar_face_distance:
+                    return face
+    return None