Add preview back part1

roop/ui.py

@@ -1,7 +1,7 @@
 import os
 import tkinter as tk
 from tkinter import filedialog
-from typing import Callable, Any
+from typing import Callable, Any, Tuple
 
 import cv2
 from PIL import Image, ImageTk, ImageOps
@@ -14,60 +14,89 @@ TERTIARY_COLOR = '#f1c40f'
 ACCENT_COLOR = '#2ecc71'
 WINDOW_HEIGHT = 700
 WINDOW_WIDTH = 600
-MAX_PREVIEW_SIZE = 800
+PREVIEW_HEIGHT = 700
+PREVIEW_WIDTH = 1200
 
 
-def init(start: Callable, destroy: Callable):
+def init(start: Callable, destroy: Callable) -> tk.Tk:
-    global WINDOW, source_label, target_label, status_label
+    global ROOT, PREVIEW
 
-    WINDOW = tk.Tk()
+    ROOT = create_root(start, destroy)
-    WINDOW.minsize(WINDOW_WIDTH, WINDOW_HEIGHT)
+    PREVIEW = create_preview(ROOT)
-    WINDOW.title('roop')
-    WINDOW.configure(bg=PRIMARY_COLOR)
-    WINDOW.option_add('*Font', ('Arial', 11))
 
-    source_label = tk.Label(bg=PRIMARY_COLOR)
+    return ROOT
+
+
+def create_root(start: Callable, destroy: Callable) -> tk.Tk:
+    global source_label, target_label, status_label
+
+    root = tk.Tk()
+    root.minsize(WINDOW_WIDTH, WINDOW_HEIGHT)
+    root.title('roop')
+    root.configure(bg=PRIMARY_COLOR)
+    root.option_add('*Font', ('Arial', 11))
+
+    source_label = tk.Label(root, bg=PRIMARY_COLOR)
     source_label.place(relx=0.1, rely=0.1, relwidth=0.3, relheight=0.25)
 
-    target_label = tk.Label(bg=PRIMARY_COLOR)
+    target_label = tk.Label(root, bg=PRIMARY_COLOR)
     target_label.place(relx=0.6, rely=0.1, relwidth=0.3, relheight=0.25)
 
-    source_button = create_primary_button(WINDOW, 'Select a face', lambda: select_source_path())
+    source_button = create_primary_button(root, 'Select a face', lambda: select_source_path())
     source_button.place(relx=0.1, rely=0.4, relwidth=0.3, relheight=0.1)
 
-    target_button = create_primary_button(WINDOW, 'Select a target', lambda: select_target_path())
+    target_button = create_primary_button(root, 'Select a target', lambda: select_target_path())
     target_button.place(relx=0.6, rely=0.4, relwidth=0.3, relheight=0.1)
 
     keep_fps_value = tk.BooleanVar(value=roop.globals.keep_fps)
-    keep_fps_checkbox = create_checkbox(WINDOW, 'Limit to 30 fps', keep_fps_value, lambda: setattr(roop.globals, 'keep_fps', not roop.globals.keep_fps))
+    keep_fps_checkbox = create_checkbox(root, 'Limit to 30 fps', keep_fps_value, lambda: setattr(roop.globals, 'keep_fps', not roop.globals.keep_fps))
     keep_fps_checkbox.place(relx=0.1, rely=0.6)
 
     keep_frames_value = tk.BooleanVar(value=roop.globals.keep_frames)
-    keep_frames_checkbox = create_checkbox(WINDOW, 'Keep frames dir', keep_frames_value, lambda: setattr(roop.globals, 'keep_frames', keep_frames_value.get()))
+    keep_frames_checkbox = create_checkbox(root, 'Keep frames dir', keep_frames_value, lambda: setattr(roop.globals, 'keep_frames', keep_frames_value.get()))
     keep_frames_checkbox.place(relx=0.1, rely=0.65)
 
     keep_audio_value = tk.BooleanVar(value=roop.globals.keep_audio)
-    keep_audio_checkbox = create_checkbox(WINDOW, 'Keep original audio', keep_frames_value, lambda: setattr(roop.globals, 'keep_audio', keep_audio_value.get()))
+    keep_audio_checkbox = create_checkbox(root, 'Keep original audio', keep_frames_value, lambda: setattr(roop.globals, 'keep_audio', keep_audio_value.get()))
     keep_audio_checkbox.place(relx=0.6, rely=0.6)
 
     many_faces_value = tk.BooleanVar(value=roop.globals.many_faces)
-    many_faces_checkbox = create_checkbox(WINDOW, 'Replace all faces', many_faces_value, lambda: setattr(roop.globals, 'many_faces', keep_audio_value.get()))
+    many_faces_checkbox = create_checkbox(root, 'Replace all faces', many_faces_value, lambda: setattr(roop.globals, 'many_faces', keep_audio_value.get()))
     many_faces_checkbox.place(relx=0.6, rely=0.65)
 
-    start_button = create_secondary_button(WINDOW, 'Start', lambda: select_output_path(start))
+    start_button = create_secondary_button(root, 'Start', lambda: select_output_path(start))
     start_button.place(relx=0.15, rely=0.75, relwidth=0.2, relheight=0.05)
 
-    stop_button = create_secondary_button(WINDOW, 'Destroy', lambda: destroy())
+    stop_button = create_secondary_button(root, 'Destroy', lambda: destroy())
     stop_button.place(relx=0.4, rely=0.75, relwidth=0.2, relheight=0.05)
 
-    preview_button = create_secondary_button(WINDOW, 'Preview', lambda: None)
+    preview_button = create_secondary_button(root, 'Preview', lambda: toggle_preview())
     preview_button.place(relx=0.65, rely=0.75, relwidth=0.2, relheight=0.05)
-    preview_button.config(state='disabled')
 
-    status_label = tk.Label(WINDOW, justify='center', text='Status: UI under heavy development, more features will soon be (re)added', fg=ACCENT_COLOR, bg=PRIMARY_COLOR)
+    status_label = tk.Label(root, justify='center', text='Status: UI under heavy development, more features will soon be (re)added', fg=ACCENT_COLOR, bg=PRIMARY_COLOR)
     status_label.place(relx=0.1, rely=0.9)
 
-    return WINDOW
+    return root
+
+
+def create_preview(parent) -> tk.Toplevel:
+    global preview_label
+
+    preview = tk.Toplevel(parent)
+    preview.withdraw()
+    preview.title('Preview')
+    preview.configure(bg=PRIMARY_COLOR)
+    preview.option_add('*Font', ('Arial', 11))
+    preview.minsize(PREVIEW_WIDTH, PREVIEW_HEIGHT)
+
+    preview_label = tk.Label(preview, bg=PRIMARY_COLOR)
+    preview_label.pack(fill='both', expand=True)
+
+    frame_value = tk.IntVar()
+    frame_slider = tk.Scale(preview, orient='horizontal', variable=frame_value)
+    frame_slider.pack(fill='x')
+
+    return preview
 
 
 def create_primary_button(parent: Any, text: str, command: Callable) -> tk.Button:
@@ -116,14 +145,14 @@ def create_checkbox(parent: Any, text: str, variable: tk.BooleanVar, command: Ca
 
 def update_status(text: str) -> None:
     status_label['text'] = text
-    WINDOW.update()
+    ROOT.update()
 
 
 def select_source_path():
     source_path = filedialog.askopenfilename(title='Select an face image')
     if is_image(source_path):
         roop.globals.source_path = source_path
-        image = render_image_preview(roop.globals.source_path)
+        image = render_image_preview(roop.globals.source_path, (200, 200))
         source_label.configure(image=image)
         source_label.image = image
     else:
@@ -141,7 +170,7 @@ def select_target_path():
         target_label.image = image
     elif is_video(target_path):
         roop.globals.target_path = target_path
-        video_frame = render_video_preview(target_path)
+        video_frame = render_video_preview(target_path, (200, 200))
         target_label.configure(image=video_frame)
         target_label.image = video_frame
     else:
@@ -157,21 +186,29 @@ def select_output_path(start):
         start()
 
 
-def render_image_preview(image_path: str) -> ImageTk.PhotoImage:
+def render_image_preview(image_path: str, dimensions: Tuple[int, int] = None) -> ImageTk.PhotoImage:
     image = Image.open(image_path)
-    image = ImageOps.fit(image, (200, 200), Image.LANCZOS)
+    if dimensions:
+        image = ImageOps.fit(image, dimensions, Image.LANCZOS)
     return ImageTk.PhotoImage(image)
 
 
-def render_video_preview(target_path: str) -> ImageTk.PhotoImage:
+def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, frame: int = 1) -> ImageTk.PhotoImage:
-    capture = cv2.VideoCapture(target_path)
+    capture = cv2.VideoCapture(video_path)
-    total_frames = capture.get(cv2.CAP_PROP_FRAME_COUNT)
+    if frame:
-    capture.set(cv2.CAP_PROP_POS_FRAMES, total_frames / 2)
+        capture.set(cv2.CAP_PROP_POS_FRAMES, frame)
     has_frame, frame = capture.read()
     if has_frame:
-        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        if dimensions:
-        image = ImageOps.fit(image, (200, 200), Image.LANCZOS)
+            image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        image = ImageOps.fit(image, dimensions, Image.LANCZOS)
         return ImageTk.PhotoImage(image)
     capture.release()
     cv2.destroyAllWindows()
 
+
+def toggle_preview():
+    if PREVIEW.state() == 'normal':
+        PREVIEW.withdraw()
+    else:
+        PREVIEW.deiconify()