From 80f3870228dca1947f9aff0e02798713e6ca26c0 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 7 Jun 2023 09:03:46 +0200 Subject: [PATCH] Add preview back part1 --- roop/ui.py | 105 ++++++++++++++++++++++++++++++++++++----------------- 1 file changed, 71 insertions(+), 34 deletions(-) diff --git a/roop/ui.py b/roop/ui.py index b1a9a43..ca9e78a 100644 --- a/roop/ui.py +++ b/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): - global WINDOW, source_label, target_label, status_label +def init(start: Callable, destroy: Callable) -> tk.Tk: + global ROOT, PREVIEW - WINDOW = tk.Tk() - WINDOW.minsize(WINDOW_WIDTH, WINDOW_HEIGHT) - WINDOW.title('roop') - WINDOW.configure(bg=PRIMARY_COLOR) - WINDOW.option_add('*Font', ('Arial', 11)) + ROOT = create_root(start, destroy) + PREVIEW = create_preview(ROOT) - 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: - capture = cv2.VideoCapture(target_path) - total_frames = capture.get(cv2.CAP_PROP_FRAME_COUNT) - capture.set(cv2.CAP_PROP_POS_FRAMES, total_frames / 2) +def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, frame: int = 1) -> ImageTk.PhotoImage: + capture = cv2.VideoCapture(video_path) + if frame: + 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)) - image = ImageOps.fit(image, (200, 200), Image.LANCZOS) + if dimensions: + 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()