diff --git a/core/config.py b/core/config.py index 151ad33..38075b1 100644 --- a/core/config.py +++ b/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 diff --git a/core/globals.py b/core/globals.py index cbd26c2..8a1c473 100644 --- a/core/globals.py +++ b/core/globals.py @@ -2,3 +2,4 @@ import onnxruntime use_gpu = False providers = onnxruntime.get_available_providers() +all_faces = False diff --git a/core/processor.py b/core/processor.py index d8ba29f..9fb62de 100644 --- a/core/processor.py +++ b/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 diff --git a/run.py b/run.py index f14e318..255432f 100644 --- a/run.py +++ b/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