Skip to content

Commit

Permalink
use exact fps
Browse files Browse the repository at this point in the history
  • Loading branch information
s0md3v committed May 31, 2023
1 parent ab2fc1c commit a9571f6
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 19 deletions.
11 changes: 5 additions & 6 deletions core/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,10 @@ def detect_fps(input_path):
output = os.popen(f'ffprobe -v error -select_streams v -of default=noprint_wrappers=1:nokey=1 -show_entries stream=r_frame_rate "{input_path}"').read()
if "/" in output:
try:
return int(output.split("/")[0]) // int(output.split("/")[1])
return int(output.split("/")[0]) // int(output.split("/")[1]), output
except:
pass
return 60
return 30, 30


def set_fps(input_path, output_path, fps):
Expand All @@ -36,22 +36,21 @@ 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 7 -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):
input_path, output_dir = path(input_path), path(output_dir)
os.system(f'ffmpeg -i "{input_path}" "{output_dir}{sep}%04d.png"')


def add_audio(output_dir, target_path, keep_frames, output_file):
video = target_path.split("/")[-1]
def add_audio(output_dir, target_path, video, keep_frames, output_file):
video_name = video.split(".")[0]
save_to = output_file if output_file else output_dir + "/swapped-" + video_name + ".mp4"
save_to_ff, output_dir_ff = path(save_to), path(output_dir)
os.system(f'ffmpeg -i "{output_dir_ff}{sep}output.mp4" -i "{output_dir_ff}{sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to_ff}"')
if not os.path.isfile(save_to):
shutil.move(output_dir + "/output.mp4", save_to)
shutil.move(f'{output_dir_ff}{sep}output.mp4', save_to)
if not keep_frames:
shutil.rmtree(output_dir)

Expand Down
27 changes: 14 additions & 13 deletions run.py
100755 → 100644
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
import platform
import sys
import time
import torch
import shutil
import core.globals
import glob
Expand Down Expand Up @@ -35,7 +34,7 @@
parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False)
parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', 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('--max-memory', help='set max memory', type=int)
parser.add_argument('--max-cores', help='number of cores to use', dest='cores_count', type=int, default=max(psutil.cpu_count() - 2, 2))

for name, value in vars(parser.parse_args()).items():
Expand All @@ -47,7 +46,7 @@


def limit_resources():
if args['max_memory'] >= 1:
if args['max_memory']:
memory = args['max_memory'] * 1024 * 1024 * 1024
if str(platform.system()).lower() == 'windows':
import ctypes
Expand All @@ -67,10 +66,11 @@ def pre_check():
if not os.path.isfile(model_path):
quit('File "inswapper_128.onnx" does not exist!')
if '--gpu' in sys.argv:
CUDA_VERSION = torch.version.cuda
CUDNN_VERSION = torch.backends.cudnn.version()
NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider']
if len(list(set(core.globals.providers) - set(NVIDIA_PROVIDERS))) > 1:
if len(list(set(core.globals.providers) - set(NVIDIA_PROVIDERS))) == 1:
import torch
CUDA_VERSION = torch.version.cuda
CUDNN_VERSION = torch.backends.cudnn.version()
if not torch.cuda.is_available() or not CUDA_VERSION:
quit("You are using --gpu flag but CUDA isn't available or properly installed on your system.")
if CUDA_VERSION > '11.8':
Expand All @@ -96,8 +96,9 @@ def start_processing():
frame_paths = args["frame_paths"]
n = len(frame_paths)//(args['cores_count'])
for i in range(n):
continue
if dataset(random.choice(frame_paths)) > 0.7:
return
quit("[WARNING] Unable to determine location of the face in the target. Please make sure the target isn't wearing clothes matching to their skin.")
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],))
Expand Down Expand Up @@ -197,16 +198,16 @@ def start():
process_img(args['source_img'], target_path, args['output_file'])
status("swap successful!")
return
video_name = os.path.basename(target_path)
video_name = os.path.splitext(video_name)[0]
video_name_full = target_path.split("/")[-1]
video_name = os.path.splitext(video_name_full)[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)
fps, exact_fps = detect_fps(target_path)
if not args['keep_fps'] and fps > 30:
this_path = output_dir + "/" + video_name + ".mp4"
set_fps(target_path, this_path, 30)
target_path, fps = this_path, 30
target_path, exact_fps = this_path, 30
else:
shutil.copy(target_path, output_dir)
status("extracting frames...")
Expand All @@ -218,9 +219,9 @@ def start():
status("swapping in progress...")
start_processing()
status("creating video...")
create_video(video_name, fps, output_dir)
create_video(video_name, exact_fps, output_dir)
status("adding audio...")
add_audio(output_dir, target_path, 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"
print("\n\nVideo saved as:", save_path, "\n\n")
status("swap successful!")
Expand Down

0 comments on commit a9571f6

Please sign in to comment.