diff --git a/roop/core.py b/roop/core.py index e685a60..79350a6 100755 --- a/roop/core.py +++ b/roop/core.py @@ -126,6 +126,23 @@ def get_video_frame(video_path, frame_number = 1): cap.release() +def update_slider(video_path): + return lambda frame_number: preview.update(get_video_frame(video_path, frame_number)) + + +def process_test_preview(video_path): + test_frame = process_faces( + get_face_single(cv2.imread(args['source_img'])), + get_video_frame(video_path, preview.current_frame.get()), + None + ) + preview.update(test_frame) + + +def preview_handler(video_path): + return lambda: preview_thread(process_test_preview(video_path)) + + def preview_video(video_path): cap = cv2.VideoCapture(video_path) if not cap.isOpened(): @@ -146,22 +163,8 @@ def preview_video(video_path): # Preview preview.update(frame) amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) - - def update_slider(frame_number): - preview.update(get_video_frame(video_path, frame_number)) - - preview.init_slider(amount_of_frames, update_slider) - - def test_handler(): - test_frame = process_faces( - get_face_single(cv2.imread(args['source_img'])), - get_video_frame(video_path, preview.current_frame.get()), - None, - roop.globals.all_faces - ) - preview.update(test_frame) - - preview.set_test_handler(lambda: preview_thread(test_handler)) + preview.init_slider(amount_of_frames, update_slider(video_path)) + preview.set_preview_handler(preview_handler(video_path)) cap.release() diff --git a/roop/swapper.py b/roop/swapper.py index 0ca8544..20294cf 100644 --- a/roop/swapper.py +++ b/roop/swapper.py @@ -29,9 +29,9 @@ def swap_face_in_frame(source_face, target_face, frame): return frame -def process_faces(source_face, frame, progress, all_faces=False): +def process_faces(source_face, frame, progress): progress_status = 'S' - if all_faces: + if roop.globals.all_faces: many_faces = get_face_many(frame) if many_faces: for face in many_faces: @@ -56,7 +56,7 @@ def process_video(source_img, frame_paths, preview_callback): for frame_path in frame_paths: frame = cv2.imread(frame_path) try: - result = process_faces(source_face, frame, progress, roop.globals.all_faces) + result = process_faces(source_face, frame, progress) cv2.imwrite(frame_path, result) if preview_callback: preview_callback(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) diff --git a/roop/ui.py b/roop/ui.py index c257c64..399878f 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -50,7 +50,7 @@ class PreviewWindow: def slider_changed(self, event): self.frame_change(self.frame_slider.get()) - def set_test_handler(self, test_handler): + def set_preview_handler(self, test_handler): self.test_button.config(command = test_handler) # Show the window