From 9c66fd9712685b4342946aab8aa4a12dcf0aa2bc Mon Sep 17 00:00:00 2001
From: henryruhs <info@henryruhs.com>
Date: Wed, 7 Jun 2023 15:40:31 +0200
Subject: [PATCH] Get the preview working

---
 roop/capturer.py  | 12 ++++++++++++
 roop/ui.py        | 36 ++++++++++++++++++++++++++----------
 roop/utilities.py |  7 ++++---
 3 files changed, 42 insertions(+), 13 deletions(-)
 create mode 100644 roop/capturer.py

diff --git a/roop/capturer.py b/roop/capturer.py
new file mode 100644
index 0000000..ecebf49
--- /dev/null
+++ b/roop/capturer.py
@@ -0,0 +1,12 @@
+import cv2
+
+
+def get_video_frame(video_path: str, frame_number: int = 1):
+    capture = cv2.VideoCapture(video_path)
+    frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT)
+    capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1))
+    has_frame, frame = capture.read()
+    capture.release()
+    if has_frame:
+        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    return None
diff --git a/roop/ui.py b/roop/ui.py
index c0d68c8..b7aa848 100644
--- a/roop/ui.py
+++ b/roop/ui.py
@@ -6,6 +6,9 @@ from typing import Callable, Any, Tuple
 import cv2
 from PIL import Image, ImageTk, ImageOps
 import roop.globals
+from roop.analyser import get_one_face
+from roop.capturer import get_video_frame
+from roop.swapper import process_faces
 from roop.utilities import is_image, is_video
 
 PRIMARY_COLOR = '#2d3436'
@@ -81,7 +84,7 @@ def create_root(start: Callable, destroy: Callable) -> tk.Tk:
 
 
 def create_preview(parent) -> tk.Toplevel:
-    global preview_label
+    global preview_label, preview_scale
 
     preview = tk.Toplevel(parent)
     preview.withdraw()
@@ -94,9 +97,8 @@ def create_preview(parent) -> tk.Toplevel:
     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')
+    preview_scale = tk.Scale(preview, orient='horizontal', command=lambda frame_value: update_preview(int(frame_value)))
+    preview_scale.pack(fill='x')
 
     return preview
 
@@ -183,7 +185,7 @@ def select_target_path():
 
 def select_output_path(start):
     output_path = filedialog.asksaveasfilename(title='Save to output file', initialfile='output.mp4')
-    if output_path and os.path.isfile(output_path):
+    if output_path:
         roop.globals.output_path = output_path
         start()
 
@@ -195,15 +197,15 @@ def render_image_preview(image_path: str, dimensions: Tuple[int, int] = None) ->
     return ImageTk.PhotoImage(image)
 
 
-def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, frame: int = 1) -> ImageTk.PhotoImage:
+def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, frame_number: int = 1) -> ImageTk.PhotoImage:
     capture = cv2.VideoCapture(video_path)
-    if frame:
-        capture.set(cv2.CAP_PROP_POS_FRAMES, frame)
+    if frame_number:
+        capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
     has_frame, frame = capture.read()
     if has_frame:
+        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
         if dimensions:
-            image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-        image = ImageOps.fit(image, dimensions, Image.LANCZOS)
+            image = ImageOps.fit(image, dimensions, Image.LANCZOS)
         return ImageTk.PhotoImage(image)
     capture.release()
     cv2.destroyAllWindows()
@@ -213,4 +215,18 @@ def toggle_preview() -> None:
     if PREVIEW.state() == 'normal':
         PREVIEW.withdraw()
     else:
+        update_preview(1)
         PREVIEW.deiconify()
+
+
+def update_preview(frame_number: int) -> None:
+    if roop.globals.source_path and roop.globals.target_path and frame_number:
+        video_frame = process_faces(
+            get_one_face(cv2.imread(roop.globals.source_path)),
+            get_video_frame(roop.globals.target_path, frame_number)
+        )
+        img = Image.fromarray(video_frame)
+        img = ImageOps.contain(img, (PREVIEW_WIDTH, PREVIEW_HEIGHT), Image.LANCZOS)
+        img = ImageTk.PhotoImage(img)
+        preview_label.configure(image=img)
+        preview_label.image = img
diff --git a/roop/utilities.py b/roop/utilities.py
index 27afbe6..8746a75 100644
--- a/roop/utilities.py
+++ b/roop/utilities.py
@@ -35,7 +35,7 @@ def extract_frames(target_path: str) -> None:
 
 
 def create_video(target_path: str, fps: int) -> None:
-    run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', roop.globals.video_quality, '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)])
+    run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)])
 
 
 def restore_audio(target_path: str, output_path: str) -> None:
@@ -67,8 +67,9 @@ def move_temp(target_path: str, output_path: str) -> None:
 
 
 def clean_temp(target_path: str) -> None:
-    if not roop.globals.keep_frames:
-        shutil.rmtree(get_temp_directory_path(target_path))
+    temp_directory_path = get_temp_directory_path(target_path)
+    if not roop.globals.keep_frames and  os.path.isdir(temp_directory_path):
+        shutil.rmtree(temp_directory_path)
 
 
 def has_image_extention(image_path: str) -> bool: