diff --git a/roop/core.py b/roop/core.py index 845ea45..191791b 100755 --- a/roop/core.py +++ b/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: diff --git a/roop/ui.py b/roop/ui.py index 5c84d64..b1a9a43 100644 --- a/roop/ui.py +++ b/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') - start() + 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() + diff --git a/roop/utilities.py b/roop/utilities.py index 38cc888..e05e323 100644 --- a/roop/utilities.py +++ b/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: - try: - if video_path and os.path.isfile(video_path): - run_ffmpeg(['-v', 'error', '-i', video_path, '-f', 'null', '-']) - return True - except subprocess.CalledProcessError: - pass + if os.path.isfile(video_path): + try: + capture = cv2.VideoCapture(video_path) + if capture.isOpened(): + is_video, _ = capture.read() + capture.release() + return is_video + except Exception: + pass return False