From 6288cdce65e9da922128dca3d2067e33ebff52c4 Mon Sep 17 00:00:00 2001 From: "K1llM@n" Date: Sat, 3 Jun 2023 16:02:51 +0300 Subject: [PATCH] Move ui from core --- roop/core.py | 195 ++++++++++++------------------------------- roop/ui.py | 228 ++++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 280 insertions(+), 143 deletions(-) diff --git a/roop/core.py b/roop/core.py index e3ef2d8..adbc4c9 100755 --- a/roop/core.py +++ b/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() - - -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")]) - + return (amount_of_frames, frame) def status(string): + value = "Status: " + string if 'cli_mode' in args: - print("Status: " + string) + print(value) else: - status_label["text"] = "Status: " + string - window.update() + ui.update_status_label(value) 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): - threading.Thread(target=thread_function).start() +def select_face_handler(path: str): + args['source_img'] = path -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 select_target_handler(path: str): + args['target_path'] = path + return preview_video(args['target_path']) +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 - preview = ui.PreviewWindow(window) + window = ui.init( + { + '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 - 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("", 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() + window.mainloop() \ No newline at end of file diff --git a/roop/ui.py b/roop/ui.py index 399878f..8b2bcac 100644 --- a/roop/ui.py +++ b/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 \ No newline at end of file + 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("", 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 \ No newline at end of file