upagge/roop
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Preview window

Browse Source
This commit is contained in:
K1llM@n 2023-06-03 00:01:44 +03:00
parent 103cb34902
commit 22ce9c3f58
3 changed files with 143 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
roop/core.py
View File

@ -21,9 +21,10 @@ import threading
from PIL import Image, ImageTk
import roop.globals
from roop.swapper import process_video, process_img
from roop.swapper import process_video, process_img, process_faces
from roop.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames, rreplace
from roop.analyser import get_face_single
import roop.ui as ui
if 'ROCMExecutionProvider' in roop.globals.providers:
del torch
@ -100,13 +101,13 @@ def start_processing():
n = len(frame_paths) // (args['cores_count'])
# single thread
if args['gpu'] or n < 2:
process_video(args['source_img'], args["frame_paths"])
process_video(args['source_img'], args["frame_paths"], preview.update)
return
# multithread if total frames to cpu cores ratio is greater than 2
if n > 2:
processes = []
for i in range(0, len(frame_paths), n):
p = pool.apply_async(process_video, args=(args['source_img'], frame_paths[i:i+n],))
p = pool.apply_async(process_video, args=(args['source_img'], frame_paths[i:i+n], preview.update,))
processes.append(p)
for p in processes:
p.get()
@ -125,6 +126,20 @@ def preview_image(image_path):
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)
cap.set(cv2.CAP_PROP_POS_FRAMES, min(amount_of_frames, frame_number-1))
if not cap.isOpened():
print("Error opening video file")
return
ret, frame = cap.read()
if ret:
return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
cap.release()
def preview_video(video_path):
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
@ -132,7 +147,7 @@ def preview_video(video_path):
return
ret, frame = cap.read()
if ret:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame = get_video_frame(video_path)
img = Image.fromarray(frame)
img = img.resize((180, 180), Image.ANTIALIAS)
photo_img = ImageTk.PhotoImage(img)
@ -142,6 +157,26 @@ def preview_video(video_path):
img_label.image = photo_img
img_label.pack()
# Preview
preview.update(frame)
amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
def update_slider(frame_number):
preview.update(get_video_frame(video_path, frame_number))
preview.init_slider(amount_of_frames, update_slider)
def test_handler():
test_frame = process_faces(
get_face_single(cv2.imread(args['source_img'])),
get_video_frame(video_path, preview.current_frame.get()),
None,
roop.globals.all_faces
)
preview.update(test_frame)
preview.set_test_handler(lambda: preview_thread(test_handler))
cap.release()
@ -237,8 +272,22 @@ def start():
status("swap successful!")
def preview_thread(thread_function):
threading.Thread(target=thread_function).start()
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 run():
global all_faces, keep_frames, limit_fps, status_label, window
global all_faces, keep_frames, limit_fps, status_label, window, preview
pre_check()
limit_resources()
@ -253,6 +302,9 @@ def run():
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)
@ -282,8 +334,12 @@ def run():
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(), start()])
start_button.place(x=240,y=560,width=120,height=49)
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")

16
roop/swapper.py
View File

@ -23,25 +23,25 @@ def swap_face_in_frame(source_face, target_face, frame):
def process_faces(source_face, frame, progress, all_faces=False):
progress_status = 'S'
if all_faces:
many_faces = get_face_many(frame)
if many_faces:
for face in many_faces:
frame = swap_face_in_frame(source_face, face, frame)
progress.set_postfix(status='.', refresh=True)
else:
progress.set_postfix(status='S', refresh=True)
progress_status='.'
else:
face = get_face_single(frame)
if face:
frame = swap_face_in_frame(source_face, face, frame)
progress.set_postfix(status='.', refresh=True)
else:
progress.set_postfix(status='S', refresh=True)
progress_status='.'
if progress:
progress.set_postfix(status=progress_status, refresh=True)
return frame
def process_video(source_img, frame_paths):
def process_video(source_img, frame_paths, preview_callback):
source_face = get_face_single(cv2.imread(source_img))
progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
@ -51,6 +51,8 @@ def process_video(source_img, frame_paths):
try:
result = process_faces(source_face, frame, progress, roop.globals.all_faces)
cv2.imwrite(frame_path, result)
if preview_callback:
preview_callback(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
except Exception:
progress.set_postfix(status='E', refresh=True)
pass

71
roop/ui.py Normal file
View File

@ -0,0 +1,71 @@
import tkinter as tk
from PIL import Image, ImageTk
class PreviewWindow:
def __init__(self, master):
self.master = master
self.window = tk.Toplevel(self.master)
# Override close button
self.window.protocol("WM_DELETE_WINDOW", self.hide)
self.window.withdraw()
self.window.geometry("600x700")
self.window.title("Preview")
self.window.configure(bg="red")
self.window.resizable(width=False, height=False)
self.visible = False
self.frame = tk.Frame(self.window, background="#2d3436")
self.frame.pack_propagate(0)
self.frame.pack(fill='both', side='left', expand='True')
# Bottom frame
buttons_frame = tk.Frame(self.frame, background="#2d3436")
buttons_frame.pack(fill='both', side='bottom')
self.current_frame = tk.IntVar()
self.frame_slider = tk.Scale(
buttons_frame,
from_=0,
to=0,
orient='horizontal',
variable=self.current_frame,
command=self.slider_changed
)
self.frame_slider.pack(fill='both', side='left', expand='True')
self.test_button = tk.Button(buttons_frame, text="Test", bg="#f1c40f", relief="flat", width=15, borderwidth=0, highlightthickness=0)
self.test_button.pack( side='right', fill='y')
def init_slider(self, frames_count, change_handler):
self.frame_change = change_handler
self.frame_slider.configure(to=frames_count)
def slider_changed(self, event):
self.frame_change(self.frame_slider.get())
def set_test_handler(self, test_handler):
self.test_button.config(command = test_handler)
# Show the window
def show(self):
self.visible = True
self.window.deiconify()
# Hide the window
def hide(self):
self.visible = False
self.window.withdraw()
def update(self, frame):
if not self.visible:
return
img = Image.fromarray(frame)
img = img.resize((600, 650), Image.ANTIALIAS)
photo_img = ImageTk.PhotoImage(img)
img_frame = tk.Frame(self.frame)
img_frame.place(x=0, y=0)
img_label = tk.Label(img_frame, image=photo_img)
img_label.image = photo_img
img_label.pack(side='top')