Move ui from core
This commit is contained in:
K1llM@n
parent
b8a158d3aa
commit
6288cdce65
193
roop/core.py
193
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)
|
||||
|
||||
# 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"))
|
||||
|
||||
# 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 = 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
|
||||
)
|
||||
|
||||
window.mainloop()
|
||||
226
roop/ui.py
226
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:
|
||||
|
||||
@ -84,3 +90,223 @@ class PreviewWindow:
|
||||
photo_img = ImageTk.PhotoImage(img)
|
||||
self.img_label.configure(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
|
||||
Loading…
Reference in New Issue
Block a user