diff --git a/README.md b/README.md index 49f7cb7..d117967 100644 --- a/README.md +++ b/README.md @@ -31,13 +31,12 @@ options: use this face -t TARGET_PATH, --target TARGET_PATH replace this face - -o OUTPUT_FILE, --output OUTPUT_FILE - save output to this file - --keep-fps maintain original fps - --gpu use gpu - --keep-frames keep frames directory - --max-memory MAX_MEMORY - set max memory +-o OUTPUT_FILE, --output OUTPUT_FILE + save output to this file +--keep-fps keep original fps +--gpu use gpu +--keep-frames don't delete frames directory +--cores number of cores to use ``` Looking for a CLI mode? Using the -f/--face argument will make the program in cli mode. diff --git a/core/utils.py b/core/utils.py index 2d12fa6..66dab65 100644 --- a/core/utils.py +++ b/core/utils.py @@ -36,7 +36,7 @@ def set_fps(input_path, output_path, fps): def create_video(video_name, fps, output_dir): output_dir = path(output_dir) - os.system(f'ffmpeg -framerate {fps} -i "{output_dir}{sep}%04d.png" -c:v libx264 -crf 32 -pix_fmt yuv420p -y "{output_dir}{sep}output.mp4"') + os.system(f'ffmpeg -framerate {fps} -i "{output_dir}{sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{sep}output.mp4"') def extract_frames(input_path, output_dir): diff --git a/run.py b/run.py index ed4f833..7fe4e78 100644 --- a/run.py +++ b/run.py @@ -34,10 +34,14 @@ parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) parser.add_argument('--max-memory', help='set max memory', default=16, type=int) +parser.add_argument('--cores', help='number of cores to use', dest='cores_count', type=int) for name, value in vars(parser.parse_args()).items(): args[name] = value +if not args['cores_count']: + args['cores_count'] = psutil.cpu_count()-1 + sep = "/" if os.name == "nt": sep = "\\" @@ -90,7 +94,7 @@ def start_processing(): print(f"Processing time: {end_time - start_time:.2f} seconds", flush=True) return frame_paths = args["frame_paths"] - n = len(frame_paths)//(psutil.cpu_count()-1) + n = len(frame_paths)//(args['cores_count']) processes = [] for i in range(0, len(frame_paths), n): p = pool.apply_async(process_video, args=(args['source_img'], frame_paths[i:i+n],)) @@ -179,7 +183,7 @@ def start(): if not args['output_file']: args['output_file'] = rreplace(args['target_path'], "/", "/swapped-", 1) if "/" in target_path else "swapped-"+target_path global pool - pool = mp.Pool(psutil.cpu_count()-1) + pool = mp.Pool(args['cores_count']) target_path = args['target_path'] test_face = get_face(cv2.imread(args['source_img'])) if not test_face: @@ -189,8 +193,9 @@ def start(): process_img(args['source_img'], target_path, args['output_file']) status("swap successful!") return - video_name = target_path.split("/")[-1].split(".")[0] - output_dir = target_path.replace(target_path.split("/")[-1], "").rstrip("/") + "/" + video_name + video_name = os.path.basename(target_path) + video_name = os.path.splitext(video_name)[0] + output_dir = os.path.join(os.path.dirname(target_path),video_name) Path(output_dir).mkdir(exist_ok=True) status("detecting video's FPS...") fps = detect_fps(target_path)