Lock face by reference (#679)

* Lock face by position of a face reference

* Prevent exception for get_many_faces

* Finalize face reference implementation

* Fix potential exception

README.md

@@ -40,6 +40,7 @@ options:
   --keep-frames                                            keep temporary frames
   --skip-audio                                             skip target audio
   --many-faces                                             process every face
+  --face-position FACE_POSITION                            position of the target face
   --video-encoder {libx264,libx265,libvpx-vp9}             adjust output video encoder
   --video-quality [0-51]                                   adjust output video quality
   --max-memory MAX_MEMORY                                  maximum amount of RAM in GB

roop/core.py

@@ -42,6 +42,7 @@ def parse_args() -> None:
     program.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true')
     program.add_argument('--skip-audio', help='skip target audio', dest='skip_audio', action='store_true')
     program.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true')
+    program.add_argument('--face-position', help='position of the target face', dest='face_position', type=int, default=0)
     program.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9'])
     program.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18, choices=range(52), metavar='[0-51]')
     program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory())
@@ -60,6 +61,7 @@ def parse_args() -> None:
     roop.globals.keep_frames = args.keep_frames
     roop.globals.skip_audio = args.skip_audio
     roop.globals.many_faces = args.many_faces
+    roop.globals.face_position = args.face_position
     roop.globals.video_encoder = args.video_encoder
     roop.globals.video_quality = args.video_quality
     roop.globals.max_memory = args.max_memory

roop/face_analyser.py

@@ -1,6 +1,7 @@
 import threading
 from typing import Any, Optional, List
 import insightface
+import numpy
 
 import roop.globals
 from roop.typing import Frame, Face
@@ -8,6 +9,8 @@ from roop.typing import Frame, Face
 FACE_ANALYSER = None
 THREAD_LOCK = threading.Lock()
 
+MAX_DISTANCE = 0.85
+
 
 def get_face_analyser() -> Any:
     global FACE_ANALYSER
@@ -19,16 +22,33 @@ def get_face_analyser() -> Any:
     return FACE_ANALYSER
 
 
-def get_one_face(frame: Frame) -> Optional[Face]:
+def clear_face_analyser() -> Any:
+    global FACE_ANALYSER
+
+    FACE_ANALYSER = None
+
+
+def get_one_face(frame: Frame, position: int = 0) -> Optional[Face]:
     faces = get_many_faces(frame)
-    try:
-        return faces[0]
-    except ValueError:
-        return None
+    if faces:
+        try:
+            return faces[position]
+        except IndexError:
+            return faces[-1]
+    return None
 
 
 def get_many_faces(frame: Frame) -> Optional[List[Face]]:
     try:
         return get_face_analyser().get(frame)
-    except IndexError:
+    except ValueError:
         return None
+
+
+def find_similar_face(frame: Frame, reference_face: Face) -> Optional[Face]:
+    faces = get_many_faces(frame)
+    for face in faces:
+        distance = numpy.sum(numpy.square(face.normed_embedding - reference_face.normed_embedding))
+        if distance < MAX_DISTANCE:
+            return face
+    return None

roop/face_reference.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+from roop.typing import Face
+
+FACE_REFERENCE = None
+
+
+def get_face_reference() -> Optional[Face]:
+    return FACE_REFERENCE
+
+
+def set_face_reference(face: Face) -> None:
+    global FACE_REFERENCE
+
+    FACE_REFERENCE = face
+
+
+def clear_face_reference() -> None:
+    global FACE_REFERENCE
+
+    FACE_REFERENCE = None

roop/globals.py

@@ -8,6 +8,7 @@ keep_fps = None
 keep_frames = None
 skip_audio = None
 many_faces = None
+face_position = None
 video_encoder = None
 video_quality = None
 max_memory = None

roop/processors/frame/face_enhancer.py

@@ -35,6 +35,12 @@ def get_device() -> str:
     return 'cpu'
 
 
+def clear_face_enhancer() -> None:
+    global FACE_ENHANCER
+
+    FACE_ENHANCER = None
+
+
 def pre_check() -> bool:
     download_directory_path = resolve_relative_path('../models')
     conditional_download(download_directory_path, ['https://huggingface.co/henryruhs/roop/resolve/main/GFPGANv1.4.pth'])
@@ -49,9 +55,7 @@ def pre_start() -> bool:
 
 
 def post_process() -> None:
-    global FACE_ENHANCER
-
-    FACE_ENHANCER = None
+    clear_face_enhancer()
 
 
 def enhance_face(temp_frame: Frame) -> Frame:

roop/processors/frame/face_swapper.py

@@ -6,7 +6,8 @@ import threading
 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.face_analyser import get_one_face, get_many_faces, find_similar_face
+from roop.face_reference import get_face_reference, set_face_reference, clear_face_reference
 from roop.typing import Face, Frame
 from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video
 
@@ -25,6 +26,12 @@ def get_face_swapper() -> Any:
     return FACE_SWAPPER
 
 
+def clear_face_swapper() -> None:
+    global FACE_SWAPPER
+
+    FACE_SWAPPER = None
+
+
 def pre_check() -> bool:
     download_directory_path = resolve_relative_path('../models')
     conditional_download(download_directory_path, ['https://huggingface.co/henryruhs/roop/resolve/main/inswapper_128.onnx'])
@@ -45,23 +52,22 @@ def pre_start() -> bool:
 
 
 def post_process() -> None:
-    global FACE_SWAPPER
-
-    FACE_SWAPPER = None
+    clear_face_swapper()
+    clear_face_reference()
 
 
 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: Face, temp_frame: Frame) -> Frame:
+def process_frame(source_face: Face, reference_face: Face, temp_frame: Frame) -> Frame:
     if roop.globals.many_faces:
         many_faces = get_many_faces(temp_frame)
         if many_faces:
             for target_face in many_faces:
                 temp_frame = swap_face(source_face, target_face, temp_frame)
     else:
-        target_face = get_one_face(temp_frame)
+        target_face = find_similar_face(temp_frame, reference_face)
         if target_face:
             temp_frame = swap_face(source_face, target_face, temp_frame)
     return temp_frame
@@ -69,9 +75,14 @@ def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
 
 def process_frames(source_path: str, temp_frame_paths: List[str], update: Callable[[], None]) -> None:
     source_face = get_one_face(cv2.imread(source_path))
+    if not get_face_reference():
+        reference_face = get_one_face(cv2.imread(temp_frame_paths[0]), roop.globals.face_position)
+        set_face_reference(reference_face)
+    else:
+        reference_face = get_face_reference()
     for temp_frame_path in temp_frame_paths:
         temp_frame = cv2.imread(temp_frame_path)
-        result = process_frame(source_face, temp_frame)
+        result = process_frame(source_face, reference_face, temp_frame)
         cv2.imwrite(temp_frame_path, result)
         if update:
             update()
@@ -80,7 +91,8 @@ def process_frames(source_path: str, temp_frame_paths: List[str], update: Callab
 def process_image(source_path: str, target_path: str, output_path: str) -> None:
     source_face = get_one_face(cv2.imread(source_path))
     target_frame = cv2.imread(target_path)
-    result = process_frame(source_face, target_frame)
+    reference_face = get_one_face(target_frame, roop.globals.face_position)
+    result = process_frame(source_face, reference_face, target_frame)
     cv2.imwrite(output_path, result)
 
 

roop/ui.py

@@ -9,6 +9,7 @@ import roop.globals
 import roop.metadata
 from roop.face_analyser import get_one_face
 from roop.capturer import get_video_frame, get_video_frame_total
+from roop.face_reference import get_face_reference, set_face_reference, clear_face_reference
 from roop.predicter import predict_frame
 from roop.processors.frame.core import get_frame_processors_modules
 from roop.utilities import is_image, is_video, resolve_relative_path
@@ -144,6 +145,7 @@ def select_target_path() -> None:
     global RECENT_DIRECTORY_TARGET
 
     PREVIEW.withdraw()
+    clear_face_reference()
     target_path = ctk.filedialog.askopenfilename(title='select an target image or video', initialdir=RECENT_DIRECTORY_TARGET)
     if is_image(target_path):
         roop.globals.target_path = target_path
@@ -199,6 +201,7 @@ def render_video_preview(video_path: str, size: Tuple[int, int], frame_number: i
 def toggle_preview() -> None:
     if PREVIEW.state() == 'normal':
         PREVIEW.withdraw()
+        clear_face_reference()
     elif roop.globals.source_path and roop.globals.target_path:
         init_preview()
         update_preview()
@@ -220,9 +223,16 @@ def update_preview(frame_number: int = 0) -> None:
         temp_frame = get_video_frame(roop.globals.target_path, frame_number)
         if predict_frame(temp_frame):
             quit()
+        source_face = get_one_face(cv2.imread(roop.globals.source_path))
+        if not get_face_reference():
+            reference_face = get_one_face(temp_frame, roop.globals.face_position)
+            set_face_reference(reference_face)
+        else:
+            reference_face = get_face_reference()
         for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
             temp_frame = frame_processor.process_frame(
-                get_one_face(cv2.imread(roop.globals.source_path)),
+                source_face,
+                reference_face,
                 temp_frame
             )
         image = Image.fromarray(cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB))