- allow processing all faces within a frame

core/config.py

@@ -11,3 +11,11 @@ def get_face(img_data):
         return sorted(analysed, key=lambda x: x.bbox[0])[0]
     except IndexError:
         return None
+
+
+def get_all_faces(img_data):
+    analysed = face_analyser.get(img_data)
+    try:
+        return analysed
+    except IndexError:
+        return None

core/globals.py

@@ -2,3 +2,4 @@ import onnxruntime
 
 use_gpu = False
 providers = onnxruntime.get_available_providers()
+all_faces = False

core/processor.py

@@ -16,13 +16,26 @@ def process_video(source_img, frame_paths):
     for frame_path in frame_paths:
         frame = cv2.imread(frame_path)
         try:
-            face = get_face(frame)
-            if face:
-                result = face_swapper.get(frame, face, source_face, paste_back=True)
-                cv2.imwrite(frame_path, result)
-                print('.', end='', flush=True)
+            if core.globals.all_faces:
+                all_faces = get_all_faces(frame)
+                result = frame
+                for singleFace in all_faces:
+                    if singleFace:
+                        result = face_swapper.get(result, singleFace, source_face, paste_back=True)
+                        print('.', end='', flush=True)
+                    else:
+                        print('S', end='', flush=True)
+                if result is not None:
+                    cv2.imwrite(frame_path, result)
             else:
-                print('S', end='', flush=True)
+                face = get_face(frame)
+                if face:
+                    result = face_swapper.get(frame, face, source_face, paste_back=True)
+                    cv2.imwrite(frame_path, result)
+                    print('.', end='', flush=True)
+                else:
+                    print('S', end='', flush=True)
+
         except Exception as e:
             print('E', end='', flush=True)
             pass

run.py

@@ -13,6 +13,9 @@ elif 'ROCMExecutionProvider' not in core.globals.providers:
     import torch
     if not torch.cuda.is_available():
         quit("You are using --gpu flag but CUDA isn't available or properly installed on your system.")
+if '--all-faces' in sys.argv or '-a' in sys.argv:
+    core.globals.all_faces = True
+
 
 import glob
 import argparse
@@ -41,6 +44,7 @@ parser.add_argument('-o', '--output', help='save output to this file', dest='out
 parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False)
 parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False)
 parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False)
+parser.add_argument('-a', '--all-faces', help='swap all faces in frame', dest='all_faces', default=False)
 
 for name, value in vars(parser.parse_args()).items():
     args[name] = value