Doing args the proper way
This commit is contained in:
henryruhs
parent
9eaeecaab5
commit
b001fb5a89
84
roop/core.py
84
roop/core.py
@ -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-vendor', help='choice your GPU vendor', dest='gpu_vendor', choices=['apple', 'amd', 'intel', 'nvidia'])
|
||||
|
||||
args = {}
|
||||
|
||||
for name, value in vars(parser.parse_args()).items():
|
||||
args[name] = value
|
||||
args = parser.parse_known_args()[0]
|
||||
|
||||
if 'all_faces' in args:
|
||||
roop.globals.all_faces = True
|
||||
|
||||
if args['cpu_cores']:
|
||||
roop.globals.cpu_cores = args['cpu_cores']
|
||||
if args.cpu_cores:
|
||||
roop.globals.cpu_cores = int(args.cpu_cores)
|
||||
|
||||
# cpu thread fix for mac
|
||||
if sys.platform == 'darwin':
|
||||
roop.globals.cpu_cores = 1
|
||||
|
||||
if args['gpu_threads']:
|
||||
roop.globals.gpu_threads = args['gpu_threads']
|
||||
if args.gpu_threads:
|
||||
roop.globals.gpu_threads = int(args.gpu_threads)
|
||||
|
||||
if args['gpu_vendor']:
|
||||
roop.globals.gpu_vendor = args['gpu_vendor']
|
||||
# gpu thread fix for amd
|
||||
if args.gpu_vendor == 'amd':
|
||||
roop.globals.gpu_threads = 1
|
||||
|
||||
if args.gpu_vendor:
|
||||
roop.globals.gpu_vendor = args.gpu_vendor
|
||||
else:
|
||||
roop.globals.providers = ['CPUExecutionProvider']
|
||||
|
||||
@ -62,12 +64,10 @@ sep = "/"
|
||||
if os.name == "nt":
|
||||
sep = "\\"
|
||||
|
||||
POOL = None
|
||||
|
||||
|
||||
def limit_resources():
|
||||
if args['max_memory']:
|
||||
memory = args['max_memory'] * 1024 * 1024 * 1024
|
||||
if args.max_memory:
|
||||
memory = args.max_memory * 1024 * 1024 * 1024
|
||||
if str(platform.system()).lower() == 'windows':
|
||||
import ctypes
|
||||
kernel32 = ctypes.windll.kernel32
|
||||
@ -156,27 +156,27 @@ def process_video_multi_cores(source_img, frame_paths):
|
||||
|
||||
|
||||
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.")
|
||||
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.")
|
||||
return
|
||||
if not args['output_file']:
|
||||
target_path = args['target_path']
|
||||
args['output_file'] = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path
|
||||
target_path = args['target_path']
|
||||
test_face = get_face_single(cv2.imread(args['source_img']))
|
||||
if not args.output_file:
|
||||
target_path = args.target_path
|
||||
args.output_file = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path
|
||||
target_path = args.target_path
|
||||
test_face = get_face_single(cv2.imread(args.source_img))
|
||||
if not test_face:
|
||||
print("\n[WARNING] No face detected in source image. Please try with another one.\n")
|
||||
return
|
||||
if is_img(target_path):
|
||||
if predict_image(target_path) > 0.85:
|
||||
quit()
|
||||
process_img(args['source_img'], target_path, args['output_file'])
|
||||
process_img(args.source_img, target_path, args.output_file)
|
||||
status("swap successful!")
|
||||
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):
|
||||
quit()
|
||||
video_name_full = target_path.split("/")[-1]
|
||||
@ -185,7 +185,7 @@ def start(preview_callback = None):
|
||||
Path(output_dir).mkdir(exist_ok=True)
|
||||
status("detecting video's FPS...")
|
||||
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"
|
||||
set_fps(target_path, 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)
|
||||
status("extracting frames...")
|
||||
extract_frames(target_path, output_dir)
|
||||
args['frame_paths'] = tuple(sorted(
|
||||
args.frame_paths = tuple(sorted(
|
||||
glob.glob(output_dir + "/*.png"),
|
||||
key=lambda x: int(x.split(sep)[-1].replace(".png", ""))
|
||||
))
|
||||
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
|
||||
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:
|
||||
process_video(args['source_img'], args["frame_paths"])
|
||||
process_video(args.source_img, args.frame_paths)
|
||||
status("creating video...")
|
||||
create_video(video_name, exact_fps, output_dir)
|
||||
status("adding audio...")
|
||||
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"
|
||||
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"
|
||||
print("\n\nVideo saved as:", save_path, "\n\n")
|
||||
status("swap successful!")
|
||||
|
||||
|
||||
def select_face_handler(path: str):
|
||||
args['source_img'] = path
|
||||
args.source_img = path
|
||||
|
||||
|
||||
def select_target_handler(path: str):
|
||||
args['target_path'] = path
|
||||
return preview_video(args['target_path'])
|
||||
args.target_path = path
|
||||
return preview_video(args.target_path)
|
||||
|
||||
|
||||
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):
|
||||
args['keep_fps'] = int(value != 1)
|
||||
args.keep_fps = int(value != 1)
|
||||
|
||||
|
||||
def toggle_keep_frames_handler(value: int):
|
||||
args['keep_frames'] = value
|
||||
args.keep_frames = value
|
||||
|
||||
|
||||
def save_file_handler(path: str):
|
||||
args['output_file'] = path
|
||||
args.output_file = path
|
||||
|
||||
|
||||
def create_test_preview(frame_number):
|
||||
return process_faces(
|
||||
get_face_single(cv2.imread(args['source_img'])),
|
||||
get_video_frame(args['target_path'], frame_number)
|
||||
get_face_single(cv2.imread(args.source_img)),
|
||||
get_video_frame(args.target_path, frame_number)
|
||||
)
|
||||
|
||||
|
||||
@ -250,16 +250,16 @@ def run():
|
||||
|
||||
pre_check()
|
||||
limit_resources()
|
||||
if args['source_img']:
|
||||
args['cli_mode'] = True
|
||||
if args.source_img:
|
||||
args.cli_mode = True
|
||||
start()
|
||||
quit()
|
||||
|
||||
window = ui.init(
|
||||
{
|
||||
'all_faces': roop.globals.all_faces,
|
||||
'keep_fps': args['keep_fps'],
|
||||
'keep_frames': args['keep_frames']
|
||||
'keep_fps': args.keep_fps,
|
||||
'keep_frames': args.keep_frames
|
||||
},
|
||||
select_face_handler,
|
||||
select_target_handler,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user