diff --git a/roop/swapper.py b/roop/swapper.py index ddc4404..8bc060b 100644 --- a/roop/swapper.py +++ b/roop/swapper.py @@ -1,4 +1,6 @@ import os +from typing import Any + from tqdm import tqdm import cv2 import insightface @@ -10,7 +12,7 @@ FACE_SWAPPER = None THREAD_LOCK = threading.Lock() -def get_face_swapper(): +def get_face_swapper() -> None: global FACE_SWAPPER with THREAD_LOCK: if FACE_SWAPPER is None: @@ -19,13 +21,13 @@ def get_face_swapper(): return FACE_SWAPPER -def swap_face_in_frame(source_face, target_face, frame): +def swap_face_in_frame(source_face: Any, target_face: Any, frame: Any) -> None: if target_face: return get_face_swapper().get(frame, target_face, source_face, paste_back=True) return frame -def process_faces(source_face, target_frame): +def process_faces(source_face: Any, target_frame: Any) -> Any: if roop.globals.many_faces: many_faces = get_many_faces(target_frame) if many_faces: @@ -38,8 +40,8 @@ def process_faces(source_face, target_frame): return target_frame -def process_frames(source_img, frame_paths, progress=None): - source_face = get_one_face(cv2.imread(source_img)) +def process_frames(source_path: str, frame_paths: [str], progress=None) -> None: + source_face = get_one_face(cv2.imread(source_path)) for frame_path in frame_paths: frame = cv2.imread(frame_path) try: @@ -52,15 +54,14 @@ def process_frames(source_img, frame_paths, progress=None): progress.update(1) -def multi_process_frame(source_img, frame_paths, progress): +def multi_process_frame(source_img, frame_paths, progress) -> None: threads = [] - num_threads = roop.globals.gpu_threads - num_frames_per_thread = len(frame_paths) // num_threads - remaining_frames = len(frame_paths) % num_threads + frames_per_thread = len(frame_paths) // roop.globals.gpu_threads + remaining_frames = len(frame_paths) % roop.globals.gpu_threads start_index = 0 # create threads by frames - for _ in range(num_threads): - end_index = start_index + num_frames_per_thread + for _ in range(roop.globals.gpu_threads): + end_index = start_index + frames_per_thread if remaining_frames > 0: end_index += 1 remaining_frames -= 1 @@ -74,15 +75,15 @@ def multi_process_frame(source_img, frame_paths, progress): thread.join() -def process_image(source_img, target_path, output_file): +def process_image(source_path: str, target_path: str, output_file) -> None: frame = cv2.imread(target_path) target_frame = get_one_face(frame) - source_face = get_one_face(cv2.imread(source_img)) + source_face = get_one_face(cv2.imread(source_path)) result = get_face_swapper().get(frame, target_frame, source_face, paste_back=True) cv2.imwrite(output_file, result) -def process_video(source_path, frame_paths, mode: str): +def process_video(source_path: str, frame_paths: [str], mode: str) -> None: progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' total = len(frame_paths) with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress: