From 22ce9c3f5854af97eb2e456ab0e2bbb3e718ddd3 Mon Sep 17 00:00:00 2001 From: "K1llM@n" Date: Sat, 3 Jun 2023 00:01:44 +0300 Subject: [PATCH] Preview window --- roop/core.py | 70 +++++++++++++++++++++++++++++++++++++++++++----- roop/swapper.py | 16 ++++++----- roop/ui.py | 71 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 143 insertions(+), 14 deletions(-) create mode 100644 roop/ui.py diff --git a/roop/core.py b/roop/core.py index b9aee40..744d765 100755 --- a/roop/core.py +++ b/roop/core.py @@ -21,9 +21,10 @@ import threading from PIL import Image, ImageTk import roop.globals -from roop.swapper import process_video, process_img +from roop.swapper import process_video, process_img, process_faces from roop.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames, rreplace from roop.analyser import get_face_single +import roop.ui as ui if 'ROCMExecutionProvider' in roop.globals.providers: del torch @@ -100,13 +101,13 @@ def start_processing(): n = len(frame_paths) // (args['cores_count']) # single thread if args['gpu'] or n < 2: - process_video(args['source_img'], args["frame_paths"]) + process_video(args['source_img'], args["frame_paths"], preview.update) return # multithread if total frames to cpu cores ratio is greater than 2 if n > 2: processes = [] for i in range(0, len(frame_paths), n): - p = pool.apply_async(process_video, args=(args['source_img'], frame_paths[i:i+n],)) + p = pool.apply_async(process_video, args=(args['source_img'], frame_paths[i:i+n], preview.update,)) processes.append(p) for p in processes: p.get() @@ -125,6 +126,20 @@ def preview_image(image_path): img_label.pack() +def get_video_frame(video_path, frame_number = 1): + cap = cv2.VideoCapture(video_path) + amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) + cap.set(cv2.CAP_PROP_POS_FRAMES, min(amount_of_frames, frame_number-1)) + if not cap.isOpened(): + print("Error opening video file") + return + ret, frame = cap.read() + if ret: + return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + cap.release() + + def preview_video(video_path): cap = cv2.VideoCapture(video_path) if not cap.isOpened(): @@ -132,7 +147,7 @@ def preview_video(video_path): return ret, frame = cap.read() if ret: - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame = get_video_frame(video_path) img = Image.fromarray(frame) img = img.resize((180, 180), Image.ANTIALIAS) photo_img = ImageTk.PhotoImage(img) @@ -142,6 +157,26 @@ def preview_video(video_path): img_label.image = photo_img img_label.pack() + # 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)) + cap.release() @@ -237,8 +272,22 @@ def start(): status("swap successful!") +def preview_thread(thread_function): + threading.Thread(target=thread_function).start() + + +def open_preview(): + if (preview.visible): + preview.hide() + else: + preview.show() + if args['target_path']: + frame = get_video_frame(args['target_path']) + preview.update(frame) + + def run(): - global all_faces, keep_frames, limit_fps, status_label, window + global all_faces, keep_frames, limit_fps, status_label, window, preview pre_check() limit_resources() @@ -253,6 +302,9 @@ def run(): window.configure(bg="#2d3436") window.resizable(width=False, height=False) + # Preview window + preview = ui.PreviewWindow(window) + # Contact information support_link = tk.Label(window, text="Donate to project <3", fg="#fd79a8", bg="#2d3436", cursor="hand2", font=("Arial", 8)) support_link.place(x=180,y=20,width=250,height=30) @@ -282,8 +334,12 @@ def run(): 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()]) - start_button.place(x=240,y=560,width=120,height=49) + start_button = tk.Button(window, text="Start", bg="#f1c40f", relief="flat", borderwidth=0, highlightthickness=0, command=lambda: [save_file(), preview_thread(start)]) + start_button.place(x=170,y=560,width=120,height=49) + + # Preview button + preview_button = tk.Button(window, text="Preview", bg="#f1c40f", relief="flat", borderwidth=0, highlightthickness=0, command=lambda: [open_preview()]) + preview_button.place(x=310,y=560,width=120,height=49) # Status label status_label = tk.Label(window, width=580, justify="center", text="Status: waiting for input...", fg="#2ecc71", bg="#2d3436") diff --git a/roop/swapper.py b/roop/swapper.py index bfc4d63..ba19c23 100644 --- a/roop/swapper.py +++ b/roop/swapper.py @@ -23,25 +23,25 @@ def swap_face_in_frame(source_face, target_face, frame): def process_faces(source_face, frame, progress, all_faces=False): + progress_status = 'S' 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) + progress_status='.' 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) + progress_status='.' + + if progress: + progress.set_postfix(status=progress_status, refresh=True) return frame -def process_video(source_img, frame_paths): +def process_video(source_img, frame_paths, preview_callback): source_face = get_face_single(cv2.imread(source_img)) progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' @@ -51,6 +51,8 @@ def process_video(source_img, frame_paths): try: result = process_faces(source_face, frame, progress, roop.globals.all_faces) cv2.imwrite(frame_path, result) + if preview_callback: + preview_callback(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) except Exception: progress.set_postfix(status='E', refresh=True) pass diff --git a/roop/ui.py b/roop/ui.py new file mode 100644 index 0000000..e2c3c5b --- /dev/null +++ b/roop/ui.py @@ -0,0 +1,71 @@ +import tkinter as tk +from PIL import Image, ImageTk + + +class PreviewWindow: + def __init__(self, master): + self.master = master + self.window = tk.Toplevel(self.master) + # Override close button + self.window.protocol("WM_DELETE_WINDOW", self.hide) + self.window.withdraw() + self.window.geometry("600x700") + self.window.title("Preview") + self.window.configure(bg="red") + self.window.resizable(width=False, height=False) + + self.visible = False + self.frame = tk.Frame(self.window, background="#2d3436") + self.frame.pack_propagate(0) + self.frame.pack(fill='both', side='left', expand='True') + + # Bottom frame + buttons_frame = tk.Frame(self.frame, background="#2d3436") + buttons_frame.pack(fill='both', side='bottom') + + self.current_frame = tk.IntVar() + self.frame_slider = tk.Scale( + buttons_frame, + from_=0, + to=0, + orient='horizontal', + variable=self.current_frame, + command=self.slider_changed + ) + self.frame_slider.pack(fill='both', side='left', expand='True') + + self.test_button = tk.Button(buttons_frame, text="Test", bg="#f1c40f", relief="flat", width=15, borderwidth=0, highlightthickness=0) + self.test_button.pack( side='right', fill='y') + + def init_slider(self, frames_count, change_handler): + self.frame_change = change_handler + self.frame_slider.configure(to=frames_count) + + def slider_changed(self, event): + self.frame_change(self.frame_slider.get()) + + def set_test_handler(self, test_handler): + self.test_button.config(command = test_handler) + + # Show the window + def show(self): + self.visible = True + self.window.deiconify() + + # Hide the window + def hide(self): + self.visible = False + self.window.withdraw() + + def update(self, frame): + if not self.visible: + return + + img = Image.fromarray(frame) + img = img.resize((600, 650), Image.ANTIALIAS) + photo_img = ImageTk.PhotoImage(img) + img_frame = tk.Frame(self.frame) + img_frame.place(x=0, y=0) + img_label = tk.Label(img_frame, image=photo_img) + img_label.image = photo_img + img_label.pack(side='top') \ No newline at end of file