Introduce render_video_preview()

roop/core.py

@@ -35,8 +35,8 @@ warnings.simplefilter(action='ignore', category=FutureWarning)
 def parse_args() -> None:
     signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
     parser = argparse.ArgumentParser()
-    parser.add_argument('-f', '--face', help='use this face', dest='source_path')
-    parser.add_argument('-t', '--target', help='replace this face', dest='target_path')
+    parser.add_argument('-f', '--face', help='use a face image', dest='source_path')
+    parser.add_argument('-t', '--target', help='replace image or video with face', dest='target_path')
     parser.add_argument('-o', '--output', help='save output to this file', dest='output_path')
     parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False)
     parser.add_argument('--keep-audio', help='maintain original audio', dest='keep_audio', action='store_true', default=True)
@@ -149,10 +149,10 @@ def update_status(message: str):
 
 def start() -> None:
     if not roop.globals.source_path or not os.path.isfile(roop.globals.source_path):
-        update_status('Please select an image containing a face.')
+        update_status('Select an image that contains a face.')
         return
     elif not roop.globals.target_path or not os.path.isfile(roop.globals.target_path):
-        update_status('Please select a video/image target!')
+        update_status('Select an image or video target!')
         return
     test_face = get_one_face(cv2.imread(roop.globals.source_path))
     if not test_face:

roop/ui.py

@@ -1,12 +1,12 @@
+import os
 import tkinter as tk
 from tkinter import filedialog
 from typing import Callable, Any
 
 import cv2
 from PIL import Image, ImageTk, ImageOps
-
 import roop.globals
-from roop.utilities import is_image
+from roop.utilities import is_image, is_video
 
 PRIMARY_COLOR = '#2d3436'
 SECONDARY_COLOR = '#74b9ff'
@@ -120,10 +120,10 @@ def update_status(text: str) -> None:
 
 
 def select_source_path():
-    path = filedialog.askopenfilename(title='Select a face')
-    if is_image(path):
-        roop.globals.source_path = path
-        image = render_frame_image(roop.globals.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)
         source_label.configure(image=image)
         source_label.image = image
     else:
@@ -133,12 +133,17 @@ def select_source_path():
 
 
 def select_target_path():
-    path = filedialog.askopenfilename(title='Select a target')
-    if is_image(path):
-        roop.globals.target_path = path
-        image = render_frame_image(roop.globals.target_path)
+    target_path = filedialog.askopenfilename(title='Select an image or video target')
+    if is_image(target_path):
+        roop.globals.target_path = target_path
+        image = render_image_preview(roop.globals.target_path)
         target_label.configure(image=image)
         target_label.image = image
+    elif is_video(target_path):
+        roop.globals.target_path = target_path
+        video_frame = render_video_preview(target_path)
+        target_label.configure(image=video_frame)
+        target_label.image = video_frame
     else:
         roop.globals.target_path = None
         target_label.configure(image=None)
@@ -146,12 +151,27 @@ def select_target_path():
 
 
 def select_output_path(start):
-    roop.globals.output_path = filedialog.askdirectory(title='Select a target')
+    output_path = filedialog.askopenfilename(title='Save to output file')
+    if os.path.isfile(output_path):
+        roop.globals.output_path = output_path
         start()
 
 
-def render_frame_image(image_path: str) -> ImageTk.PhotoImage:
+def render_image_preview(image_path: str) -> ImageTk.PhotoImage:
     image = Image.open(image_path)
     image = ImageOps.fit(image, (200, 200), 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)
+    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)
+        return ImageTk.PhotoImage(image)
+    capture.release()
+    cv2.destroyAllWindows()
+

roop/utilities.py

@@ -5,6 +5,8 @@ import subprocess
 from pathlib import Path
 from typing import List, Any
 
+import cv2
+
 import roop.globals
 from PIL import Image
 
@@ -74,7 +76,7 @@ def has_image_extention(image_path: str) -> bool:
 
 
 def is_image(image_path: str) -> bool:
-    if image_path and os.path.isfile(image_path):
+    if os.path.isfile(image_path):
         try:
             image = Image.open(image_path)
             image.verify()
@@ -85,10 +87,13 @@ def is_image(image_path: str) -> bool:
 
 
 def is_video(video_path: str) -> bool:
+    if os.path.isfile(video_path):
         try:
-        if video_path and os.path.isfile(video_path):
-            run_ffmpeg(['-v', 'error', '-i', video_path, '-f', 'null', '-'])
-            return True
-    except subprocess.CalledProcessError:
+            capture = cv2.VideoCapture(video_path)
+            if capture.isOpened():
+                is_video, _ = capture.read()
+                capture.release()
+                return is_video
+        except Exception:
             pass
     return False