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

Doing args the proper way

Browse Source
This commit is contained in:
henryruhs 2023-06-04 18:37:55 +02:00
parent 9eaeecaab5
commit b001fb5a89
1 changed files with 42 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

84
roop/core.py
View File

@ -36,25 +36,27 @@ parser.add_argument('--cpu-cores', help='number of CPU cores to use', dest='cpu_
parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=4) parser.add_argument('--gpu-threads', help='number of threads to be use for the GPU', dest='gpu_threads', type=int, default=4)
parser.add_argument('--gpu-vendor', help='choice your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'intel', 'nvidia']) parser.add_argument('--gpu-vendor', help='choice your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'intel', 'nvidia'])
args = {} args = parser.parse_known_args()[0]
for name, value in vars(parser.parse_args()).items():
args[name] = value
if 'all_faces' in args: if 'all_faces' in args:
roop.globals.all_faces = True roop.globals.all_faces = True
if args['cpu_cores']: if args.cpu_cores:
roop.globals.cpu_cores = args['cpu_cores'] roop.globals.cpu_cores = int(args.cpu_cores)
# cpu thread fix for mac
if sys.platform == 'darwin': if sys.platform == 'darwin':
roop.globals.cpu_cores = 1 roop.globals.cpu_cores = 1
if args['gpu_threads']: if args.gpu_threads:
roop.globals.gpu_threads = args['gpu_threads'] roop.globals.gpu_threads = int(args.gpu_threads)
if args['gpu_vendor']: # gpu thread fix for amd
roop.globals.gpu_vendor = args['gpu_vendor'] if args.gpu_vendor == 'amd':
roop.globals.gpu_threads = 1
if args.gpu_vendor:
roop.globals.gpu_vendor = args.gpu_vendor
else: else:
roop.globals.providers = ['CPUExecutionProvider'] roop.globals.providers = ['CPUExecutionProvider']
@ -62,12 +64,10 @@ sep = "/"
if os.name == "nt": if os.name == "nt":
sep = "\\" sep = "\\"
POOL = None
def limit_resources(): def limit_resources():
if args['max_memory']: if args.max_memory:
memory = args['max_memory'] * 1024 * 1024 * 1024 memory = args.max_memory * 1024 * 1024 * 1024
if str(platform.system()).lower() == 'windows': if str(platform.system()).lower() == 'windows':
import ctypes import ctypes
kernel32 = ctypes.windll.kernel32 kernel32 = ctypes.windll.kernel32
@ -156,27 +156,27 @@ def process_video_multi_cores(source_img, frame_paths):
def start(preview_callback = None): def start(preview_callback = None):
if not args['source_img'] or not os.path.isfile(args['source_img']): if not args.source_img or not os.path.isfile(args.source_img):
print("\n[WARNING] Please select an image containing a face.") print("\n[WARNING] Please select an image containing a face.")
return return
elif not args['target_path'] or not os.path.isfile(args['target_path']): elif not args.target_path or not os.path.isfile(args.target_path):
print("\n[WARNING] Please select a video/image to swap face in.") print("\n[WARNING] Please select a video/image to swap face in.")
return return
if not args['output_file']: if not args.output_file:
target_path = args['target_path'] target_path = args.target_path
args['output_file'] = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path args.output_file = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path
target_path = args['target_path'] target_path = args.target_path
test_face = get_face_single(cv2.imread(args['source_img'])) test_face = get_face_single(cv2.imread(args.source_img))
if not test_face: if not test_face:
print("\n[WARNING] No face detected in source image. Please try with another one.\n") print("\n[WARNING] No face detected in source image. Please try with another one.\n")
return return
if is_img(target_path): if is_img(target_path):
if predict_image(target_path) > 0.85: if predict_image(target_path) > 0.85:
quit() quit()
process_img(args['source_img'], target_path, args['output_file']) process_img(args.source_img, target_path, args.output_file)
status("swap successful!") status("swap successful!")
return return
seconds, probabilities = predict_video_frames(video_path=args['target_path'], frame_interval=100) seconds, probabilities = predict_video_frames(video_path=args.target_path, frame_interval=100)
if any(probability > 0.85 for probability in probabilities): if any(probability > 0.85 for probability in probabilities):
quit() quit()
video_name_full = target_path.split("/")[-1] video_name_full = target_path.split("/")[-1]
@ -185,7 +185,7 @@ def start(preview_callback = None):
Path(output_dir).mkdir(exist_ok=True) Path(output_dir).mkdir(exist_ok=True)
status("detecting video's FPS...") status("detecting video's FPS...")
fps, exact_fps = detect_fps(target_path) fps, exact_fps = detect_fps(target_path)
if not args['keep_fps'] and fps > 30: if not args.keep_fps and fps > 30:
this_path = output_dir + "/" + video_name + ".mp4" this_path = output_dir + "/" + video_name + ".mp4"
set_fps(target_path, this_path, 30) set_fps(target_path, this_path, 30)
target_path, exact_fps = this_path, 30 target_path, exact_fps = this_path, 30
@ -193,33 +193,33 @@ def start(preview_callback = None):
shutil.copy(target_path, output_dir) shutil.copy(target_path, output_dir)
status("extracting frames...") status("extracting frames...")
extract_frames(target_path, output_dir) extract_frames(target_path, output_dir)
args['frame_paths'] = tuple(sorted( args.frame_paths = tuple(sorted(
glob.glob(output_dir + "/*.png"), glob.glob(output_dir + "/*.png"),
key=lambda x: int(x.split(sep)[-1].replace(".png", "")) key=lambda x: int(x.split(sep)[-1].replace(".png", ""))
)) ))
status("swapping in progress...") status("swapping in progress...")
if roop.globals.cpu_cores > 0 and roop.globals.gpu_vendor is None: if roop.globals.gpu_vendor is None and roop.globals.cpu_cores > 0:
global POOL global POOL
POOL = mp.Pool(roop.globals.cpu_cores) POOL = mp.Pool(roop.globals.cpu_cores)
process_video_multi_cores(args['source_img'], args['frame_paths']) process_video_multi_cores(args.source_img, args.frame_paths)
else: else:
process_video(args['source_img'], args["frame_paths"]) process_video(args.source_img, args.frame_paths)
status("creating video...") status("creating video...")
create_video(video_name, exact_fps, output_dir) create_video(video_name, exact_fps, output_dir)
status("adding audio...") status("adding audio...")
add_audio(output_dir, target_path, video_name_full, args['keep_frames'], args['output_file']) add_audio(output_dir, target_path, video_name_full, args.keep_frames, args.output_file)
save_path = args['output_file'] if args['output_file'] else output_dir + "/" + video_name + ".mp4" save_path = args.output_file if args.output_file else output_dir + "/" + video_name + ".mp4"
print("\n\nVideo saved as:", save_path, "\n\n") print("\n\nVideo saved as:", save_path, "\n\n")
status("swap successful!") status("swap successful!")
def select_face_handler(path: str): def select_face_handler(path: str):
args['source_img'] = path args.source_img = path
def select_target_handler(path: str): def select_target_handler(path: str):
args['target_path'] = path args.target_path = path
return preview_video(args['target_path']) return preview_video(args.target_path)
def toggle_all_faces_handler(value: int): def toggle_all_faces_handler(value: int):
@ -227,21 +227,21 @@ def toggle_all_faces_handler(value: int):
def toggle_fps_limit_handler(value: int): def toggle_fps_limit_handler(value: int):
args['keep_fps'] = int(value != 1) args.keep_fps = int(value != 1)
def toggle_keep_frames_handler(value: int): def toggle_keep_frames_handler(value: int):
args['keep_frames'] = value args.keep_frames = value
def save_file_handler(path: str): def save_file_handler(path: str):
args['output_file'] = path args.output_file = path
def create_test_preview(frame_number): def create_test_preview(frame_number):
return process_faces( return process_faces(
get_face_single(cv2.imread(args['source_img'])), get_face_single(cv2.imread(args.source_img)),
get_video_frame(args['target_path'], frame_number) get_video_frame(args.target_path, frame_number)
) )
@ -250,16 +250,16 @@ def run():
pre_check() pre_check()
limit_resources() limit_resources()
if args['source_img']: if args.source_img:
args['cli_mode'] = True args.cli_mode = True
start() start()
quit() quit()
window = ui.init( window = ui.init(
{ {
'all_faces': roop.globals.all_faces, 'all_faces': roop.globals.all_faces,
'keep_fps': args['keep_fps'], 'keep_fps': args.keep_fps,
'keep_frames': args['keep_frames'] 'keep_frames': args.keep_frames
}, },
select_face_handler, select_face_handler,
select_target_handler, select_target_handler,