upagge/roop
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add typing to swapper

Browse Source
This commit is contained in:
henryruhs 2023-06-09 12:39:50 +02:00
parent 27a76eea60
commit 1a7db694cf
1 changed files with 15 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
roop/swapper.py
View File

@ -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: