Move ui from core

roop/core.py

@@ -9,14 +9,8 @@ import argparse
 import os
 import torch
 from pathlib import Path
-import tkinter as tk
-from tkinter import filedialog
 from opennsfw2 import predict_video_frames, predict_image
-from tkinter.filedialog import asksaveasfilename
-import webbrowser
 import cv2
-import threading
-from PIL import Image, ImageTk
 
 import roop.globals
 from roop.swapper import process_video, process_img, process_faces
@@ -101,17 +95,6 @@ def pre_check():
         roop.globals.providers = ['CPUExecutionProvider']
 
 
-def preview_image(image_path):
-    img = Image.open(image_path)
-    img = img.resize((180, 180), Image.ANTIALIAS)
-    photo_img = ImageTk.PhotoImage(img)
-    left_frame = tk.Frame(window)
-    left_frame.place(x=60, y=100)
-    img_label = tk.Label(left_frame, image=photo_img)
-    img_label.image = photo_img
-    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)
@@ -126,84 +109,25 @@ 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():
         print("Error opening video file")
-        return
+        return 0
+    amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
     ret, frame = cap.read()
     if ret:
         frame = get_video_frame(video_path)
-        img = Image.fromarray(frame)
-        img = img.resize((180, 180), Image.ANTIALIAS)
-        photo_img = ImageTk.PhotoImage(img)
-        right_frame = tk.Frame(window)
-        right_frame.place(x=360, y=100)
-        img_label = tk.Label(right_frame, image=photo_img)
-        img_label.image = photo_img
-        img_label.pack()
-
-        # Preview
-        preview.update(frame)
-        amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
-        preview.init_slider(amount_of_frames, update_slider(video_path))
-        preview.set_preview_handler(preview_handler(video_path))
 
     cap.release()
-
+    return (amount_of_frames, frame)
-
-def select_face():
-    args['source_img'] = filedialog.askopenfilename(title="Select a face")
-    preview_image(args['source_img'])
-
-
-def select_target():
-    args['target_path'] = filedialog.askopenfilename(title="Select a target")
-    threading.Thread(target=preview_video, args=(args['target_path'],)).start()
-
-
-def toggle_fps_limit():
-    args['keep_fps'] = int(limit_fps.get() != True)
-
-
-def toggle_all_faces():
-    roop.globals.all_faces = True if all_faces.get() == 1 else False
-
-
-def toggle_keep_frames():
-    args['keep_frames'] = int(keep_frames.get())
-
-
-def save_file():
-    filename, ext = 'output.mp4', '.mp4'
-    if is_img(args['target_path']):
-        filename, ext = 'output.png', '.png'
-    args['output_file'] = asksaveasfilename(initialfile=filename, defaultextension=ext, filetypes=[("All Files","*.*"),("Videos","*.mp4")])
-
 
 def status(string):
+    value = "Status: " + string
     if 'cli_mode' in args:
-        print("Status: " + string)
+        print(value)
     else:
-        status_label["text"] = "Status: " + string
+        ui.update_status_label(value)
-        window.update()
 
 
 def start():
@@ -249,7 +173,7 @@ def start():
         key=lambda x: int(x.split(sep)[-1].replace(".png", ""))
     ))
     status("swapping in progress...")
-    process_video(args['source_img'], args["frame_paths"], preview.update if preview else None)
+    process_video(args['source_img'], args["frame_paths"], ui.preview.update if ui.preview else None)
     status("creating video...")
     create_video(video_name, exact_fps, output_dir)
     status("adding audio...")
@@ -259,22 +183,40 @@ def start():
     status("swap successful!")
 
 
-def preview_thread(thread_function):
+def select_face_handler(path: str):
-    threading.Thread(target=thread_function).start()
+    args['source_img'] = path
 
 
-def open_preview():
+def select_target_handler(path: str):
-    if (preview.visible):
+    args['target_path'] = path
-        preview.hide()
+    return preview_video(args['target_path'])
-    else:
-        preview.show()
-        if args['target_path']:
-            frame = get_video_frame(args['target_path'])
-            preview.update(frame)
 
 
+def toggle_all_faces_handler(value: int):
+    roop.globals.all_faces = True if value == 1 else False
+
+
+def toggle_fps_limit_handler(value: int):
+    args['keep_fps'] = int(value != 1)
+
+
+def toggle_keep_frames_handler(value: int):
+    args['keep_frames'] = value
+
+
+def save_file_handler(path: str):
+    args['output_file'] = path
+
+
+def create_test_preview(frame_number):
+    return process_faces(
+        get_face_single(cv2.imread(args['source_img'])), 
+        get_video_frame(args['target_path'], frame_number),
+        None
+    )
+
 def run():
-    global all_faces, keep_frames, limit_fps, status_label, window, preview
+    global all_faces, keep_frames, limit_fps
 
     pre_check()
     limit_resources()
@@ -282,53 +224,22 @@ def run():
         args['cli_mode'] = True
         start()
         quit()
-    window = tk.Tk()
-    window.geometry("600x700")
-    window.title("roop")
-    window.configure(bg="#2d3436")
-    window.resizable(width=False, height=False)
 
-    # Preview window
+    window = ui.init(
-    preview = ui.PreviewWindow(window)
+        {
+            'all_faces': roop.globals.all_faces,
+            'keep_fps': args['keep_fps'],
+            'keep_frames': args['keep_frames']
+        },
+        select_face_handler,
+        select_target_handler,
+        toggle_all_faces_handler,
+        toggle_fps_limit_handler,
+        toggle_keep_frames_handler,
+        save_file_handler,
+        start,
+        get_video_frame,
+        create_test_preview
+    )
 
-    # Contact information
+    window.mainloop()
-    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)
-    support_link.bind("<Button-1>", lambda e: webbrowser.open("https://github.com/sponsors/s0md3v"))
-
-    # Select a face button
-    face_button = tk.Button(window, text="Select a face", command=select_face, bg="#2d3436", fg="#74b9ff", highlightthickness=4, relief="flat", highlightbackground="#74b9ff", activebackground="#74b9ff", borderwidth=4)
-    face_button.place(x=60,y=320,width=180,height=80)
-
-    # Select a target button
-    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, 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, 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(), 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")
-    status_label.place(x=10,y=640,width=580,height=30)
-
-    window.mainloop()

roop/ui.py

@@ -1,6 +1,12 @@
 import tkinter as tk
+from typing import Any, Callable, Tuple
 from PIL import Image, ImageTk
+import webbrowser
+from tkinter import filedialog
+from tkinter.filedialog import asksaveasfilename
+import threading
 
+from roop.utils import is_img
 
 class PreviewWindow:
 
@@ -83,4 +89,224 @@ class PreviewWindow:
         )
         photo_img = ImageTk.PhotoImage(img)
         self.img_label.configure(image=photo_img)
-        self.img_label.image = photo_img
+        self.img_label.image = photo_img
+
+
+def select_face(select_face_handler: Callable[[str], None]):
+    if select_face_handler:
+        path = filedialog.askopenfilename(title="Select a face")
+        preview_face(path)
+        return select_face_handler(path)
+    return None
+
+
+def update_slider_handler(get_video_frame, video_path):
+    return lambda frame_number: preview.update(get_video_frame(video_path, frame_number))
+
+def test_preview(create_test_preview):
+    frame = create_test_preview(preview.current_frame.get())
+    preview.update(frame)
+
+def update_slider(get_video_frame, create_test_preview, video_path, frames_amount):
+    preview.init_slider(frames_amount, update_slider_handler(get_video_frame, video_path))
+    preview.set_preview_handler(lambda: preview_thread(test_preview(create_test_preview)))
+
+
+def analyze_target(select_target_handler: Callable[[str], Tuple[int, Any]], target_path: tk.StringVar, frames_amount: tk.IntVar):    
+    path = filedialog.askopenfilename(title="Select a target")
+    target_path.set(path)
+    amount, frame = select_target_handler(path)
+    frames_amount.set(amount)
+    preview_target(frame)
+    preview.update(frame)
+
+
+def select_target(select_target_handler: Callable[[str], Tuple[int, Any]], target_path: tk.StringVar, frames_amount: tk.IntVar):
+    if select_target_handler:
+        analyze_target(select_target_handler, target_path, frames_amount)
+
+
+def save_file(save_file_handler: Callable[[str], None], target_path: str):
+    filename, ext = 'output.mp4', '.mp4'
+
+    if is_img(target_path):
+        filename, ext = 'output.png', '.png'
+
+    if save_file_handler:
+        return save_file_handler(asksaveasfilename(initialfile=filename, defaultextension=ext, filetypes=[("All Files","*.*"),("Videos","*.mp4")]))
+    return None
+
+def toggle_all_faces(toggle_all_faces_handler: Callable[[int], None], variable: tk.IntVar):
+    if toggle_all_faces_handler:
+        return lambda: toggle_all_faces_handler(variable.get())
+    return None
+
+
+def toggle_fps_limit(toggle_all_faces_handler: Callable[[int], None], variable: tk.IntVar):
+    if toggle_all_faces_handler:
+        return lambda: toggle_all_faces_handler(variable.get())
+    return None
+
+
+def toggle_keep_frames(toggle_keep_frames_handler: Callable[[int], None], variable: tk.IntVar):
+    if toggle_keep_frames_handler:
+        return lambda: toggle_keep_frames_handler(variable.get())
+    return None
+
+
+def create_button(parent, text, command):
+    return tk.Button(
+        parent, 
+        text=text, 
+        command=command,
+        bg="#f1c40f", 
+        relief="flat", 
+        borderwidth=0, 
+        highlightthickness=0
+    )
+
+
+def create_background_button(parent, text, command):
+    button = create_button(parent, text, command)
+    button.configure(
+        bg="#2d3436", 
+        fg="#74b9ff", 
+        highlightthickness=4, 
+        highlightbackground="#74b9ff", 
+        activebackground="#74b9ff", 
+        borderwidth=4
+    )
+    return button
+
+
+def create_check(parent, text, variable, command):
+    return tk.Checkbutton(
+        parent, 
+        anchor="w", 
+        relief="groove", 
+        activebackground="#2d3436", 
+        activeforeground="#74b9ff", 
+        selectcolor="black", 
+        text=text, 
+        fg="#dfe6e9", 
+        borderwidth=0, 
+        highlightthickness=0, 
+        bg="#2d3436", 
+        variable=variable, 
+        command=command
+    )
+
+
+def preview_thread(thread_function):
+    threading.Thread(target=thread_function).start()
+
+
+def open_preview_window(get_video_frame, target_path):
+    if (preview.visible):
+        preview.hide()
+    else:
+        preview.show()
+        if target_path:
+            frame = get_video_frame(target_path)
+            preview.update(frame)
+
+def preview_face(path):
+    img = Image.open(path)
+    img = img.resize((180, 180), Image.ANTIALIAS)
+    photo_img = ImageTk.PhotoImage(img)
+    face_label.configure(image=photo_img)
+    face_label.image = photo_img
+
+def preview_target(frame):
+    img = Image.fromarray(frame)
+    img = img.resize((180, 180), Image.ANTIALIAS)
+    photo_img = ImageTk.PhotoImage(img)
+    target_label.configure(image=photo_img)
+    target_label.image = photo_img
+
+def update_status_label(value):
+    status_label["text"] = value
+    window.update()
+
+def init(
+    initial_values: dict,
+    select_face_handler: Callable[[str], None],
+    select_target_handler: Callable[[str], Tuple[int, Any]],
+    toggle_all_faces_handler: Callable[[int], None],
+    toggle_fps_limit_handler: Callable[[int], None],
+    toggle_keep_frames_handler: Callable[[int], None],
+    save_file_handler: Callable[[str], None],
+    start: Callable[[], None],
+    get_video_frame: Callable[[str, int], None],
+    create_test_preview: Callable[[int], Any],
+):
+    global window, preview, face_label, target_label, status_label
+
+    window = tk.Tk()
+    window.geometry("600x700")
+    window.title("roop")
+    window.configure(bg="#2d3436")
+    window.resizable(width=False, height=False)
+
+    target_path = tk.StringVar()
+    frames_amount = tk.IntVar()
+
+    # Preview window
+    preview = 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)
+    support_link.bind("<Button-1>", lambda e: webbrowser.open("https://github.com/sponsors/s0md3v"))
+
+    left_frame = tk.Frame(window)
+    left_frame.place(x=60, y=100)
+    face_label = tk.Label(left_frame)
+    face_label.pack(fill='both', side='top', expand=True)
+
+    right_frame = tk.Frame(window)
+    right_frame.place(x=360, y=100)
+    target_label = tk.Label(right_frame)
+    target_label.pack(fill='both', side='top', expand=True)
+
+    # Select a face button
+    face_button = create_background_button(window, "Select a face", lambda: [
+        select_face(select_face_handler)
+    ])
+    face_button.place(x=60,y=320,width=180,height=80)
+
+    # Select a target button
+    target_button = create_background_button(window, "Select a target", lambda: [
+        select_target(select_target_handler, target_path, frames_amount),
+        update_slider(get_video_frame, create_test_preview, target_path.get(), frames_amount.get())
+    ])
+    target_button.place(x=360,y=320,width=180,height=80)
+
+    # All faces checkbox
+    all_faces = tk.IntVar(None, initial_values['all_faces'])
+    all_faces_checkbox = create_check(window, "Process all faces in frame", all_faces, toggle_all_faces(toggle_all_faces_handler, all_faces))
+    all_faces_checkbox.place(x=60,y=500,width=240,height=31)
+
+    # FPS limit checkbox
+    limit_fps = tk.IntVar(None, not initial_values['keep_fps'])
+    fps_checkbox = create_check(window, "Limit FPS to 30", limit_fps, toggle_fps_limit(toggle_fps_limit_handler, limit_fps))
+    fps_checkbox.place(x=60,y=475,width=240,height=31)
+
+    # Keep frames checkbox
+    keep_frames = tk.IntVar(None, initial_values['keep_frames'])
+    frames_checkbox = create_check(window, "Keep frames dir", keep_frames, toggle_keep_frames(toggle_keep_frames_handler, keep_frames))
+    frames_checkbox.place(x=60,y=450,width=240,height=31)
+
+    # Start button
+    start_button = create_button(window, "Start", lambda: [save_file(save_file_handler, target_path.get()), preview_thread(start)])
+    start_button.place(x=170,y=560,width=120,height=49)
+
+    # Preview button
+    preview_button = create_button(window, "Preview", lambda: open_preview_window(get_video_frame, target_path.get()))
+    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")
+    status_label.place(x=10,y=640,width=580,height=30)
+
+    return window