Merge branch 'main' into main

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
+-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
+--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.

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

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)