Preview window

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 = tk.Button(window, text="Start", bg="#f1c40f", relief="flat", borderwidth=0, highlightthickness=0, command=lambda: [save_file(), preview_thread(start)])
-    start_button.place(x=240,y=560,width=120,height=49)
+    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")

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)
+            progress_status='.'
-        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)
+            progress_status='.'
-        else:
+
-            progress.set_postfix(status='S', refresh=True)
+    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

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')