mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2024-11-20 23:10:08 -08:00
123bccf01a
3.optional) denoise data_dst images.bat Apply it if dst video is very sharp. Denoise dst images before face extraction. This technique helps neural network not to learn the noise. The result is less pixel shake of the predicted face.
271 lines
8.9 KiB
Python
271 lines
8.9 KiB
Python
import subprocess
|
|
import numpy as np
|
|
import ffmpeg
|
|
from pathlib import Path
|
|
from core import pathex
|
|
from core.interact import interact as io
|
|
|
|
def extract_video(input_file, output_dir, output_ext=None, fps=None):
|
|
input_file_path = Path(input_file)
|
|
output_path = Path(output_dir)
|
|
|
|
if not output_path.exists():
|
|
output_path.mkdir(exist_ok=True)
|
|
|
|
|
|
if input_file_path.suffix == '.*':
|
|
input_file_path = pathex.get_first_file_by_stem (input_file_path.parent, input_file_path.stem)
|
|
else:
|
|
if not input_file_path.exists():
|
|
input_file_path = None
|
|
|
|
if input_file_path is None:
|
|
io.log_err("input_file not found.")
|
|
return
|
|
|
|
if fps is None:
|
|
fps = io.input_int ("Enter FPS", 0, help_message="How many frames of every second of the video will be extracted. 0 - full fps")
|
|
|
|
if output_ext is None:
|
|
output_ext = io.input_str ("Output image format", "png", ["png","jpg"], help_message="png is lossless, but extraction is x10 slower for HDD, requires x10 more disk space than jpg.")
|
|
|
|
for filename in pathex.get_image_paths (output_path, ['.'+output_ext]):
|
|
Path(filename).unlink()
|
|
|
|
job = ffmpeg.input(str(input_file_path))
|
|
|
|
kwargs = {'pix_fmt': 'rgb24'}
|
|
if fps != 0:
|
|
kwargs.update ({'r':str(fps)})
|
|
|
|
if output_ext == 'jpg':
|
|
kwargs.update ({'q:v':'2'}) #highest quality for jpg
|
|
|
|
job = job.output( str (output_path / ('%5d.'+output_ext)), **kwargs )
|
|
|
|
try:
|
|
job = job.run()
|
|
except:
|
|
io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )
|
|
|
|
def cut_video ( input_file, from_time=None, to_time=None, audio_track_id=None, bitrate=None):
|
|
input_file_path = Path(input_file)
|
|
if input_file_path is None:
|
|
io.log_err("input_file not found.")
|
|
return
|
|
|
|
output_file_path = input_file_path.parent / (input_file_path.stem + "_cut" + input_file_path.suffix)
|
|
|
|
if from_time is None:
|
|
from_time = io.input_str ("From time", "00:00:00.000")
|
|
|
|
if to_time is None:
|
|
to_time = io.input_str ("To time", "00:00:00.000")
|
|
|
|
if audio_track_id is None:
|
|
audio_track_id = io.input_int ("Specify audio track id.", 0)
|
|
|
|
if bitrate is None:
|
|
bitrate = max (1, io.input_int ("Bitrate of output file in MB/s", 25) )
|
|
|
|
kwargs = {"c:v": "libx264",
|
|
"b:v": "%dM" %(bitrate),
|
|
"pix_fmt": "yuv420p",
|
|
}
|
|
|
|
job = ffmpeg.input(str(input_file_path), ss=from_time, to=to_time)
|
|
|
|
job_v = job['v:0']
|
|
job_a = job['a:' + str(audio_track_id) + '?' ]
|
|
|
|
job = ffmpeg.output(job_v, job_a, str(output_file_path), **kwargs).overwrite_output()
|
|
|
|
try:
|
|
job = job.run()
|
|
except:
|
|
io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )
|
|
|
|
def denoise_image_sequence( input_dir, ext=None, factor=None ):
|
|
input_path = Path(input_dir)
|
|
|
|
if not input_path.exists():
|
|
io.log_err("input_dir not found.")
|
|
return
|
|
|
|
image_paths = [ Path(filepath) for filepath in pathex.get_image_paths(input_path) ]
|
|
|
|
# Check extension of all images
|
|
image_paths_suffix = None
|
|
for filepath in image_paths:
|
|
if image_paths_suffix is None:
|
|
image_paths_suffix = filepath.suffix
|
|
else:
|
|
if filepath.suffix != image_paths_suffix:
|
|
io.log_err(f"All images in {input_path.name} should be with the same extension.")
|
|
return
|
|
|
|
if factor is None:
|
|
factor = np.clip ( io.input_int ("Denoise factor?", 7, add_info="1-20"), 1, 20 )
|
|
|
|
# Rename to temporary filenames
|
|
for i,filepath in io.progress_bar_generator( enumerate(image_paths), "Renaming", leave=False):
|
|
src = filepath
|
|
dst = filepath.parent / ( f'{i+1:06}_{filepath.name}' )
|
|
try:
|
|
src.rename (dst)
|
|
except:
|
|
io.log_error ('fail to rename %s' % (src.name) )
|
|
return
|
|
|
|
# Rename to sequental filenames
|
|
for i,filepath in io.progress_bar_generator( enumerate(image_paths), "Renaming", leave=False):
|
|
|
|
src = filepath.parent / ( f'{i+1:06}_{filepath.name}' )
|
|
dst = filepath.parent / ( f'{i+1:06}{filepath.suffix}' )
|
|
try:
|
|
src.rename (dst)
|
|
except:
|
|
io.log_error ('fail to rename %s' % (src.name) )
|
|
return
|
|
|
|
# Process image sequence in ffmpeg
|
|
kwargs = {}
|
|
if image_paths_suffix == '.jpg':
|
|
kwargs.update ({'q:v':'2'})
|
|
|
|
job = ( ffmpeg
|
|
.input(str ( input_path / ('%6d'+image_paths_suffix) ) )
|
|
.filter("hqdn3d", factor, factor, 5,5)
|
|
.output(str ( input_path / ('%6d'+image_paths_suffix) ), **kwargs )
|
|
)
|
|
|
|
try:
|
|
job = job.run()
|
|
except:
|
|
io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )
|
|
|
|
# Rename to temporary filenames
|
|
for i,filepath in io.progress_bar_generator( enumerate(image_paths), "Renaming", leave=False):
|
|
src = filepath.parent / ( f'{i+1:06}{filepath.suffix}' )
|
|
dst = filepath.parent / ( f'{i+1:06}_{filepath.name}' )
|
|
try:
|
|
src.rename (dst)
|
|
except:
|
|
io.log_error ('fail to rename %s' % (src.name) )
|
|
return
|
|
|
|
# Rename to initial filenames
|
|
for i,filepath in io.progress_bar_generator( enumerate(image_paths), "Renaming", leave=False):
|
|
src = filepath.parent / ( f'{i+1:06}_{filepath.name}' )
|
|
dst = filepath
|
|
|
|
try:
|
|
src.rename (dst)
|
|
except:
|
|
io.log_error ('fail to rename %s' % (src.name) )
|
|
return
|
|
|
|
def video_from_sequence( input_dir, output_file, reference_file=None, ext=None, fps=None, bitrate=None, include_audio=False, lossless=None ):
|
|
input_path = Path(input_dir)
|
|
output_file_path = Path(output_file)
|
|
reference_file_path = Path(reference_file) if reference_file is not None else None
|
|
|
|
if not input_path.exists():
|
|
io.log_err("input_dir not found.")
|
|
return
|
|
|
|
if not output_file_path.parent.exists():
|
|
output_file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
return
|
|
|
|
out_ext = output_file_path.suffix
|
|
|
|
if ext is None:
|
|
ext = io.input_str ("Input image format (extension)", "png")
|
|
|
|
if lossless is None:
|
|
lossless = io.input_bool ("Use lossless codec", False)
|
|
|
|
video_id = None
|
|
audio_id = None
|
|
ref_in_a = None
|
|
if reference_file_path is not None:
|
|
if reference_file_path.suffix == '.*':
|
|
reference_file_path = pathex.get_first_file_by_stem (reference_file_path.parent, reference_file_path.stem)
|
|
else:
|
|
if not reference_file_path.exists():
|
|
reference_file_path = None
|
|
|
|
if reference_file_path is None:
|
|
io.log_err("reference_file not found.")
|
|
return
|
|
|
|
#probing reference file
|
|
probe = ffmpeg.probe (str(reference_file_path))
|
|
|
|
#getting first video and audio streams id with fps
|
|
for stream in probe['streams']:
|
|
if video_id is None and stream['codec_type'] == 'video':
|
|
video_id = stream['index']
|
|
fps = stream['r_frame_rate']
|
|
|
|
if audio_id is None and stream['codec_type'] == 'audio':
|
|
audio_id = stream['index']
|
|
|
|
if audio_id is not None:
|
|
#has audio track
|
|
ref_in_a = ffmpeg.input (str(reference_file_path))[str(audio_id)]
|
|
|
|
if fps is None:
|
|
#if fps not specified and not overwritten by reference-file
|
|
fps = max (1, io.input_int ("Enter FPS", 25) )
|
|
|
|
if not lossless and bitrate is None:
|
|
bitrate = max (1, io.input_int ("Bitrate of output file in MB/s", 16) )
|
|
|
|
input_image_paths = pathex.get_image_paths(input_path)
|
|
|
|
i_in = ffmpeg.input('pipe:', format='image2pipe', r=fps)
|
|
|
|
output_args = [i_in]
|
|
|
|
if include_audio and ref_in_a is not None:
|
|
output_args += [ref_in_a]
|
|
|
|
output_args += [str (output_file_path)]
|
|
|
|
output_kwargs = {}
|
|
|
|
if lossless:
|
|
output_kwargs.update ({"c:v": "libx264",
|
|
"crf": "0",
|
|
"pix_fmt": "yuv420p",
|
|
})
|
|
else:
|
|
output_kwargs.update ({"c:v": "libx264",
|
|
"b:v": "%dM" %(bitrate),
|
|
"pix_fmt": "yuv420p",
|
|
})
|
|
|
|
if include_audio and ref_in_a is not None:
|
|
output_kwargs.update ({"c:a": "aac",
|
|
"b:a": "192k",
|
|
"ar" : "48000",
|
|
"strict": "experimental"
|
|
})
|
|
|
|
job = ( ffmpeg.output(*output_args, **output_kwargs).overwrite_output() )
|
|
|
|
try:
|
|
job_run = job.run_async(pipe_stdin=True)
|
|
|
|
for image_path in input_image_paths:
|
|
with open (image_path, "rb") as f:
|
|
image_bytes = f.read()
|
|
job_run.stdin.write (image_bytes)
|
|
|
|
job_run.stdin.close()
|
|
job_run.wait()
|
|
except:
|
|
io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )
|