diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..09916c4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +.idea +__pycache__ +*.onnx \ No newline at end of file diff --git a/core/analyser.py b/core/analyser.py index 8a3d908..85a0b7c 100644 --- a/core/analyser.py +++ b/core/analyser.py @@ -12,9 +12,16 @@ def get_face_analyser(): return FACE_ANALYSER -def get_face(img_data): +def get_face_single(img_data): face = get_face_analyser().get(img_data) try: return sorted(face, key=lambda x: x.bbox[0])[0] except IndexError: return None + + +def get_face_many(img_data): + try: + return get_face_analyser().get(img_data) + except IndexError: + return None diff --git a/core/globals.py b/core/globals.py index 3525c9d..b237e8a 100644 --- a/core/globals.py +++ b/core/globals.py @@ -1,7 +1,8 @@ -import onnxruntime - -use_gpu = False -providers = onnxruntime.get_available_providers() - -if 'TensorrtExecutionProvider' in providers: - providers.remove('TensorrtExecutionProvider') +import onnxruntime + +use_gpu = False +all_faces = False +providers = onnxruntime.get_available_providers() + +if 'TensorrtExecutionProvider' in providers: + providers.remove('TensorrtExecutionProvider') diff --git a/core/swapper.py b/core/swapper.py index 188b71f..fdfd646 100644 --- a/core/swapper.py +++ b/core/swapper.py @@ -3,7 +3,7 @@ from tqdm import tqdm import cv2 import insightface import core.globals -from core.analyser import get_face +from core.analyser import get_face_single, get_face_many FACE_SWAPPER = None @@ -16,19 +16,41 @@ def get_face_swapper(): return FACE_SWAPPER +def swap_face_in_frame(source_face, target_face, frame): + if target_face: + return get_face_swapper().get(frame, target_face, source_face, paste_back=True) + return frame + + +def process_faces(source_face, frame, progress, all_faces=False): + if all_faces: + many_faces = get_face_many(frame) + if many_faces: + for face in many_faces: + frame = swap_face_in_frame(source_face, face, frame) + progress.set_postfix(status='.', refresh=True) + else: + progress.set_postfix(status='S', refresh=True) + else: + face = get_face_single(frame) + if face: + frame = swap_face_in_frame(source_face, face, frame) + progress.set_postfix(status='.', refresh=True) + else: + progress.set_postfix(status='S', refresh=True) + return frame + + def process_video(source_img, frame_paths): - source_face = get_face(cv2.imread(source_img)) - with tqdm(total=len(frame_paths), desc="Processing", unit="frame", dynamic_ncols=True, bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]') as progress: + source_face = get_face_single(cv2.imread(source_img)) + progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' + + with tqdm(total=len(frame_paths), desc="Processing", unit="frame", dynamic_ncols=True, bar_format=progress_bar_format) as progress: for frame_path in frame_paths: frame = cv2.imread(frame_path) try: - face = get_face(frame) - if face: - result = get_face_swapper().get(frame, face, source_face, paste_back=True) - cv2.imwrite(frame_path, result) - progress.set_postfix(status='.', refresh=True) - else: - progress.set_postfix(status='S', refresh=True) + result = process_faces(source_face, frame, progress, core.globals.all_faces) + cv2.imwrite(frame_path, result) except Exception: progress.set_postfix(status='E', refresh=True) pass @@ -37,8 +59,8 @@ def process_video(source_img, frame_paths): def process_img(source_img, target_path, output_file): frame = cv2.imread(target_path) - face = get_face(frame) - source_face = get_face(cv2.imread(source_img)) + face = get_face_single(frame) + source_face = get_face_single(cv2.imread(source_img)) result = get_face_swapper().get(frame, face, source_face, paste_back=True) cv2.imwrite(output_file, result) print("\n\nImage saved as:", output_file, "\n\n") diff --git a/run.py b/run.py index abca34b..2099fdc 100755 --- a/run.py +++ b/run.py @@ -22,7 +22,7 @@ from PIL import Image, ImageTk import core.globals from core.swapper import process_video, process_img from core.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames, rreplace -from core.analyser import get_face +from core.analyser import get_face_single if 'ROCMExecutionProvider' in core.globals.providers: del torch @@ -40,6 +40,7 @@ parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', type=int) parser.add_argument('--max-cores', help='number of cores to be use for CPU mode', dest='cores_count', type=int, default=max(psutil.cpu_count() - 2, 2)) +parser.add_argument('--all-faces', help='swap all faces in frame', dest='all_faces', action='store_true', default=False) for name, value in vars(parser.parse_args()).items(): args[name] = value @@ -149,6 +150,10 @@ def toggle_fps_limit(): args['keep_fps'] = int(limit_fps.get() != True) +def toggle_all_faces(): + core.globals.all_faces = True if all_faces.get() == 1 else False + + def toggle_keep_frames(): args['keep_frames'] = int(keep_frames.get()) @@ -181,12 +186,12 @@ def start(): global pool pool = mp.Pool(args['cores_count']) target_path = args['target_path'] - test_face = get_face(cv2.imread(args['source_img'])) + test_face = get_face_single(cv2.imread(args['source_img'])) if not test_face: print("\n[WARNING] No face detected in source image. Please try with another one.\n") return if is_img(target_path): - if predict_image(target_path) > 0.7: + if predict_image(target_path) > 0.85: quit() process_img(args['source_img'], target_path, args['output_file']) status("swap successful!") @@ -196,7 +201,7 @@ def start(): quit() video_name_full = target_path.split("/")[-1] video_name = os.path.splitext(video_name_full)[0] - output_dir = os.path.dirname(target_path) + "/" + video_name + output_dir = os.path.dirname(target_path) + "/" + video_name if os.path.dirname(target_path) else video_name Path(output_dir).mkdir(exist_ok=True) status("detecting video's FPS...") fps, exact_fps = detect_fps(target_path) @@ -252,15 +257,20 @@ if __name__ == "__main__": target_button = tk.Button(window, text="Select a target", command=select_target, bg="#2d3436", fg="#74b9ff", highlightthickness=4, relief="flat", highlightbackground="#74b9ff", activebackground="#74b9ff", borderwidth=4) target_button.place(x=360,y=320,width=180,height=80) + # All faces checkbox + all_faces = tk.IntVar() + all_faces_checkbox = tk.Checkbutton(window, anchor="w", relief="groove", activebackground="#2d3436", activeforeground="#74b9ff", selectcolor="black", text="Process all faces in frame", fg="#dfe6e9", borderwidth=0, highlightthickness=0, bg="#2d3436", variable=all_faces, command=toggle_all_faces) + all_faces_checkbox.place(x=60,y=500,width=240,height=31) + # FPS limit checkbox limit_fps = tk.IntVar(None, not args['keep_fps']) - fps_checkbox = tk.Checkbutton(window, relief="groove", activebackground="#2d3436", activeforeground="#74b9ff", selectcolor="black", text="Limit FPS to 30", fg="#dfe6e9", borderwidth=0, highlightthickness=0, bg="#2d3436", variable=limit_fps, command=toggle_fps_limit) - fps_checkbox.place(x=30,y=500,width=240,height=31) + fps_checkbox = tk.Checkbutton(window, anchor="w", relief="groove", activebackground="#2d3436", activeforeground="#74b9ff", selectcolor="black", text="Limit FPS to 30", fg="#dfe6e9", borderwidth=0, highlightthickness=0, bg="#2d3436", variable=limit_fps, command=toggle_fps_limit) + fps_checkbox.place(x=60,y=475,width=240,height=31) # Keep frames checkbox keep_frames = tk.IntVar(None, args['keep_frames']) - frames_checkbox = tk.Checkbutton(window, relief="groove", activebackground="#2d3436", activeforeground="#74b9ff", selectcolor="black", text="Keep frames dir", fg="#dfe6e9", borderwidth=0, highlightthickness=0, bg="#2d3436", variable=keep_frames, command=toggle_keep_frames) - frames_checkbox.place(x=37,y=450,width=240,height=31) + frames_checkbox = tk.Checkbutton(window, anchor="w", relief="groove", activebackground="#2d3436", activeforeground="#74b9ff", selectcolor="black", text="Keep frames dir", fg="#dfe6e9", borderwidth=0, highlightthickness=0, bg="#2d3436", variable=keep_frames, command=toggle_keep_frames) + frames_checkbox.place(x=60,y=450,width=240,height=31) # Start button start_button = tk.Button(window, text="Start", bg="#f1c40f", relief="flat", borderwidth=0, highlightthickness=0, command=lambda: [save_file(), start()]) @@ -269,5 +279,5 @@ if __name__ == "__main__": # Status label status_label = tk.Label(window, width=580, justify="center", text="Status: waiting for input...", fg="#2ecc71", bg="#2d3436") status_label.place(x=10,y=640,width=580,height=30) - + window.mainloop()