Introduce Face and Frame typing

roop/face_analyser.py

@@ -1,6 +1,8 @@
 from typing import Any
 import insightface
+
 import roop.globals
+from roop.typing import Frame
 
 FACE_ANALYSER = None
 
@@ -14,7 +16,7 @@ def get_face_analyser() -> Any:
     return FACE_ANALYSER
 
 
-def get_one_face(frame: Any) -> Any:
+def get_one_face(frame: Frame) -> Any:
     face = get_face_analyser().get(frame)
     try:
         return min(face, key=lambda x: x.bbox[0])
@@ -22,7 +24,7 @@ def get_one_face(frame: Any) -> Any:
         return None
 
 
-def get_many_faces(frame: Any) -> Any:
+def get_many_faces(frame: Frame) -> Any:
     try:
         return get_face_analyser().get(frame)
     except IndexError:

roop/predicter.py

@@ -2,10 +2,12 @@ import numpy
 import opennsfw2
 from PIL import Image
 
+from roop.typing import Frame
+
 MAX_PROBABILITY = 0.85
 
 
-def predict_frame(target_frame: Image) -> bool:
+def predict_frame(target_frame: Frame) -> bool:
     image = Image.fromarray(target_frame)
     image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO)
     model = opennsfw2.make_open_nsfw_model()

roop/processors/frame/face_enhancer.py

@@ -7,6 +7,7 @@ import roop.globals
 import roop.processors.frame.core
 from roop.core import update_status
 from roop.face_analyser import get_one_face
+from roop.typing import Frame, Face
 from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video
 
 FACE_ENHANCER = None
@@ -39,7 +40,7 @@ def get_face_enhancer() -> Any:
     return FACE_ENHANCER
 
 
-def enhance_face(temp_frame: Any) -> Any:
+def enhance_face(temp_frame: Frame) -> Frame:
     with THREAD_SEMAPHORE:
         _, _, temp_frame = get_face_enhancer().enhance(
             temp_frame,
@@ -48,7 +49,7 @@ def enhance_face(temp_frame: Any) -> Any:
     return temp_frame
 
 
-def process_frame(source_face: Any, temp_frame: Any) -> Any:
+def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
     target_face = get_one_face(temp_frame)
     if target_face:
         temp_frame = enhance_face(temp_frame)

roop/processors/frame/face_swapper.py

@@ -7,6 +7,7 @@ import roop.globals
 import roop.processors.frame.core
 from roop.core import update_status
 from roop.face_analyser import get_one_face, get_many_faces
+from roop.typing import Face, Frame
 from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video
 
 FACE_SWAPPER = None
@@ -43,11 +44,11 @@ def get_face_swapper() -> Any:
     return FACE_SWAPPER
 
 
-def swap_face(source_face: Any, target_face: Any, temp_frame: Any) -> Any:
+def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
     return get_face_swapper().get(temp_frame, target_face, source_face, paste_back=True)
 
 
-def process_frame(source_face: Any, temp_frame: Any) -> Any:
+def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
     if roop.globals.many_faces:
         many_faces = get_many_faces(temp_frame)
         if many_faces:

roop/typing.py

@@ -0,0 +1,7 @@
+from typing import Any
+
+from insightface.app.common import Face
+import numpy
+
+Face = Face
+Frame = numpy.ndarray[Any, Any]