Update roop/ui.py

roop/ui.py

@@ -31,7 +31,6 @@ source_label = None
 target_label = None
 status_label = None
 
-
 def init(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.CTk:
     global ROOT, PREVIEW
 
@@ -40,7 +39,6 @@ def init(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.CTk:
 
     return ROOT
 
-
 def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.CTk:
     global source_label, target_label, status_label
 
@@ -101,7 +99,6 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C
 
     return root
 
-
 def create_preview(parent: ctk.CTkToplevel) -> ctk.CTkToplevel:
     global preview_label, preview_slider
 
@@ -119,17 +116,15 @@ def create_preview(parent: ctk.CTkToplevel) -> ctk.CTkToplevel:
 
     return preview
 
-
 def update_status(text: str) -> None:
     status_label.configure(text=text)
     ROOT.update()
 
-
 def select_source_path() -> None:
     global RECENT_DIRECTORY_SOURCE
 
     PREVIEW.withdraw()
-    source_path = ctk.filedialog.askopenfilename(title='select an source image', initialdir=RECENT_DIRECTORY_SOURCE)
+    source_path = ctk.filedialog.askopenfilename(title='select a source image', initialdir=RECENT_DIRECTORY_SOURCE)
     if is_image(source_path):
         roop.globals.source_path = source_path
         RECENT_DIRECTORY_SOURCE = os.path.dirname(roop.globals.source_path)
@@ -139,12 +134,11 @@ def select_source_path() -> None:
         roop.globals.source_path = None
         source_label.configure(image=None)
 
-
 def select_target_path() -> None:
     global RECENT_DIRECTORY_TARGET
 
     PREVIEW.withdraw()
-    target_path = ctk.filedialog.askopenfilename(title='select an target image or video', initialdir=RECENT_DIRECTORY_TARGET)
+    target_path = ctk.filedialog.askopenfilename(title='select a target image or video', initialdir=RECENT_DIRECTORY_TARGET)
     if is_image(target_path):
         roop.globals.target_path = target_path
         RECENT_DIRECTORY_TARGET = os.path.dirname(roop.globals.target_path)
@@ -159,7 +153,6 @@ def select_target_path() -> None:
         roop.globals.target_path = None
         target_label.configure(image=None)
 
-
 def select_output_path(start: Callable[[], None]) -> None:
     global RECENT_DIRECTORY_OUTPUT
 
@@ -174,58 +167,39 @@ def select_output_path(start: Callable[[], None]) -> None:
         RECENT_DIRECTORY_OUTPUT = os.path.dirname(roop.globals.output_path)
         start()
 
-
 def render_image_preview(image_path: str, size: Tuple[int, int]) -> ctk.CTkImage:
     image = Image.open(image_path)
     if size:
         image = ImageOps.fit(image, size, Image.LANCZOS)
     return ctk.CTkImage(image, size=image.size)
 
-
 def render_video_preview(video_path: str, size: Tuple[int, int], frame_number: int = 0) -> ctk.CTkImage:
     capture = cv2.VideoCapture(video_path)
     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 size:
-            image = ImageOps.fit(image, size, Image.LANCZOS)
-        return ctk.CTkImage(image, size=image.size)
-    capture.release()
-    cv2.destroyAllWindows()
+    if not has_frame:
+        return None
+    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    image = Image.fromarray(frame)
+    if size:
+        image = ImageOps.fit(image, size, Image.LANCZOS)
+    return ctk.CTkImage(image, size=image.size)
 
+def update_preview(frame_value: float) -> None:
+    frame_number = int(frame_value)
+    preview_image = None
+    if is_video(roop.globals.target_path):
+        preview_image = render_video_preview(roop.globals.target_path, (PREVIEW_MAX_WIDTH, PREVIEW_MAX_HEIGHT), frame_number)
+    if preview_image:
+        preview_label.configure(image=preview_image)
 
 def toggle_preview() -> None:
-    if PREVIEW.state() == 'normal':
+    if PREVIEW.winfo_ismapped():
         PREVIEW.withdraw()
-    elif roop.globals.source_path and roop.globals.target_path:
-        init_preview()
-        update_preview()
+    else:
+        if is_video(roop.globals.target_path):
+            total_frames = get_video_frame_total(roop.globals.target_path)
+            preview_slider.configure(from_=0, to=total_frames, command=lambda frame_value: update_preview(frame_value))
         PREVIEW.deiconify()
-
-
-def init_preview() -> None:
-    if is_image(roop.globals.target_path):
-        preview_slider.pack_forget()
-    if is_video(roop.globals.target_path):
-        video_frame_total = get_video_frame_total(roop.globals.target_path)
-        preview_slider.configure(to=video_frame_total)
-        preview_slider.pack(fill='x')
-        preview_slider.set(0)
-
-
-def update_preview(frame_number: int = 0) -> None:
-    if roop.globals.source_path and roop.globals.target_path:
-        temp_frame = get_video_frame(roop.globals.target_path, frame_number)
-        if predict_frame(temp_frame):
-            quit()
-        for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
-            temp_frame = frame_processor.process_frame(
-                get_one_face(cv2.imread(roop.globals.source_path)),
-                temp_frame
-            )
-        image = Image.fromarray(cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB))
-        image = ImageOps.contain(image, (PREVIEW_MAX_WIDTH, PREVIEW_MAX_HEIGHT), Image.LANCZOS)
-        image = ctk.CTkImage(image, size=image.size)
-        preview_label.configure(image=image)
+        PREVIEW.update()