diff --git a/roop/core.py b/roop/core.py index 191791b..ab8dd07 100755 --- a/roop/core.py +++ b/roop/core.py @@ -42,6 +42,7 @@ def parse_args() -> None: parser.add_argument('--keep-audio', help='maintain original audio', dest='keep_audio', action='store_true', default=True) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) parser.add_argument('--many-faces', help='swap every face in the frame', dest='many_faces', action='store_true', default=False) + parser.add_argument('--video-quality', help='adjust video quality of output file', dest='video_quality', type=int, default=10) parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_cores', type=int, default=max(psutil.cpu_count() / 2, 1)) parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=8) @@ -57,6 +58,7 @@ def parse_args() -> None: roop.globals.keep_audio = args.keep_audio roop.globals.keep_frames = args.keep_frames roop.globals.many_faces = args.many_faces + roop.globals.video_quality = args.video_quality if args.cpu_cores: roop.globals.cpu_cores = int(args.cpu_cores) diff --git a/roop/globals.py b/roop/globals.py index 6680e63..f280566 100644 --- a/roop/globals.py +++ b/roop/globals.py @@ -7,6 +7,7 @@ keep_fps = None keep_audio = None keep_frames = None many_faces = None +video_quality = None cpu_cores = None gpu_threads = None gpu_vendor = None diff --git a/roop/ui.py b/roop/ui.py index ca9e78a..c0d68c8 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -35,6 +35,7 @@ def create_root(start: Callable, destroy: Callable) -> tk.Tk: root.title('roop') root.configure(bg=PRIMARY_COLOR) root.option_add('*Font', ('Arial', 11)) + root.protocol('WM_DELETE_WINDOW', lambda: destroy()) source_label = tk.Label(root, bg=PRIMARY_COLOR) source_label.place(relx=0.1, rely=0.1, relwidth=0.3, relheight=0.25) @@ -73,7 +74,7 @@ def create_root(start: Callable, destroy: Callable) -> tk.Tk: preview_button = create_secondary_button(root, 'Preview', lambda: toggle_preview()) preview_button.place(relx=0.65, rely=0.75, relwidth=0.2, relheight=0.05) - status_label = tk.Label(root, justify='center', text='Status: UI under heavy development, more features will soon be (re)added', fg=ACCENT_COLOR, bg=PRIMARY_COLOR) + status_label = tk.Label(root, justify='center', text='Status: None', fg=ACCENT_COLOR, bg=PRIMARY_COLOR) status_label.place(relx=0.1, rely=0.9) return root @@ -88,6 +89,7 @@ def create_preview(parent) -> tk.Toplevel: preview.configure(bg=PRIMARY_COLOR) preview.option_add('*Font', ('Arial', 11)) preview.minsize(PREVIEW_WIDTH, PREVIEW_HEIGHT) + preview.protocol('WM_DELETE_WINDOW', lambda: toggle_preview()) preview_label = tk.Label(preview, bg=PRIMARY_COLOR) preview_label.pack(fill='both', expand=True) @@ -180,8 +182,8 @@ def select_target_path(): def select_output_path(start): - output_path = filedialog.askopenfilename(title='Save to output file') - if os.path.isfile(output_path): + output_path = filedialog.asksaveasfilename(title='Save to output file', initialfile='output.mp4') + if output_path and os.path.isfile(output_path): roop.globals.output_path = output_path start() @@ -207,7 +209,7 @@ def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, fr cv2.destroyAllWindows() -def toggle_preview(): +def toggle_preview() -> None: if PREVIEW.state() == 'normal': PREVIEW.withdraw() else: diff --git a/roop/utilities.py b/roop/utilities.py index e05e323..27afbe6 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', '7', '-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', roop.globals.video_quality, '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) def restore_audio(target_path: str, output_path: str) -> None: @@ -76,7 +76,7 @@ def has_image_extention(image_path: str) -> bool: def is_image(image_path: str) -> bool: - if os.path.isfile(image_path): + if image_path and os.path.isfile(image_path): try: image = Image.open(image_path) image.verify() @@ -87,7 +87,7 @@ def is_image(image_path: str) -> bool: def is_video(video_path: str) -> bool: - if os.path.isfile(video_path): + if video_path and os.path.isfile(video_path): try: capture = cv2.VideoCapture(video_path) if capture.isOpened():