Improve face reference by introducing more CLI args

README.md

@@ -40,7 +40,9 @@ 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
+  --reference-face-position REFERENCE_FACE_POSITION        position of the reference face
+  --reference-frame-number REFERENCE_FRAME_NUMBER          number of the reference frame
+  --similar-face-distance SIMILAR_FACE_DISTANCE            face distance used for recognition
   --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,7 +42,9 @@ 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('--reference-face-position', help='position of the reference face', dest='reference_face_position', type=int, default=0)
+    program.add_argument('--reference-frame-number', help='number of the reference frame', dest='reference_frame_number', type=int, default=0)
+    program.add_argument('--similar-face-distance', help='face distance used for recognition', dest='similar_face_distance', type=float, default=0.85)
     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())
@@ -61,7 +63,9 @@ 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.reference_face_position = args.reference_face_position
+    roop.globals.reference_frame_number = args.reference_frame_number
+    roop.globals.similar_face_distance = args.similar_face_distance
     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

@@ -9,8 +9,6 @@ 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
@@ -49,6 +47,6 @@ 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:
+        if distance < roop.globals.similar_face_distance:
             return face
     return None

roop/globals.py

@@ -8,7 +8,9 @@ keep_fps = None
 keep_frames = None
 skip_audio = None
 many_faces = None
-face_position = None
+reference_face_position = None
+reference_frame_number = None
+similar_face_distance = None
 video_encoder = None
 video_quality = None
 max_memory = None

roop/processors/frame/face_swapper.py

@@ -76,7 +76,8 @@ def process_frame(source_face: Face, reference_face: Face, temp_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)
+        reference_frame = cv2.imread(temp_frame_paths[roop.globals.reference_frame_number])
+        reference_face = get_one_face(reference_frame, roop.globals.reference_face_position)
         set_face_reference(reference_face)
     else:
         reference_face = get_face_reference()
@@ -91,7 +92,7 @@ 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)
-    reference_face = get_one_face(target_frame, roop.globals.face_position)
+    reference_face = get_one_face(target_frame, roop.globals.reference_face_position)
     result = process_frame(source_face, reference_face, target_frame)
     cv2.imwrite(output_path, result)
 

roop/ui.py

@@ -205,7 +205,7 @@ def toggle_preview() -> None:
         clear_face_reference()
     elif roop.globals.source_path and roop.globals.target_path:
         init_preview()
-        update_preview()
+        update_preview(roop.globals.reference_frame_number)
         PREVIEW.deiconify()
 
 
@@ -216,7 +216,7 @@ def init_preview() -> None:
         video_frame_total = get_video_frame_total(roop.globals.target_path)
         preview_slider.configure(to=video_frame_total)
         preview_slider.pack(fill='x')
-        preview_slider.set(0)
+        preview_slider.set(roop.globals.reference_frame_number)
 
 
 def update_preview(frame_number: int = 0) -> None:
@@ -226,7 +226,7 @@ def update_preview(frame_number: int = 0) -> None:
             sys.exit()
         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)
+            reference_face = get_one_face(temp_frame, roop.globals.reference_face_position)
             set_face_reference(reference_face)
         else:
             reference_face = get_face_reference()