diff --git a/main.py b/main.py index 1158957..873e32e 100644 --- a/main.py +++ b/main.py @@ -120,9 +120,65 @@ if __name__ == "__main__": convert_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.") convert_parser.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") convert_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Convert on CPU.") - convert_parser.set_defaults(func=process_convert) - + + videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers() + + def process_videoed_extract_video(arguments): + from mainscripts import VideoEd + VideoEd.extract_video (arguments.input_file, arguments.output_dir, arguments.output_ext, arguments.fps) + p = videoed_parser.add_parser( "extract-video", help="Extract images from video file.") + p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.") + p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted images will be stored.") + p.add_argument('--ouptut-ext', dest="output_ext", default='png', help="Image format (extension) of output files.") + p.add_argument('--fps', type=int, dest="fps", default=None, help="How many frames of every second of the video will be extracted. 0 - full fps.") + p.set_defaults(func=process_videoed_extract_video) + + def process_videoed_cut_video(arguments): + from mainscripts import VideoEd + VideoEd.cut_video (arguments.input_file, + arguments.from_time, + arguments.to_time, + arguments.audio_track_id, + arguments.bitrate) + p = videoed_parser.add_parser( "cut-video", help="Cut video file.") + p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.") + p.add_argument('--from-time', dest="from_time", default=None, help="From time, for example 00:00:00.000") + p.add_argument('--to-time', dest="to_time", default=None, help="To time, for example 00:00:00.000") + p.add_argument('--audio-track-id', type=int, dest="audio_track_id", default=None, help="Specify audio track id.") + p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.") + p.set_defaults(func=process_videoed_cut_video) + + def process_videoed_denoise_image_sequence(arguments): + from mainscripts import VideoEd + VideoEd.denoise_image_sequence (arguments.input_dir, arguments.ext, arguments.factor) + p = videoed_parser.add_parser( "denoise-image-sequence", help="Denoise sequence of images, keeping sharp edges. This allows you to make the final fake more believable, since the neural network is not able to make a detailed skin texture, but it makes the edges quite clear. Therefore, if the whole frame is more `blurred`, then a fake will seem more believable. Especially true for scenes of the film, which are usually very clear.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input file to be processed. Specify .*-extension to find first file.") + p.add_argument('--ext', dest="ext", default='png', help="Image format (extension) of input files.") + p.add_argument('--factor', type=int, dest="factor", default=None, help="Denoise factor (1-20).") + p.set_defaults(func=process_videoed_denoise_image_sequence) + + def process_videoed_video_from_sequence(arguments): + from mainscripts import VideoEd + VideoEd.video_from_sequence (arguments.input_dir, + arguments.output_file, + arguments.reference_file, + arguments.ext, + arguments.fps, + arguments.bitrate, + arguments.lossless) + + p = videoed_parser.add_parser( "video-from-sequence", help="Make video from image sequence.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input file to be processed. Specify .*-extension to find first file.") + p.add_argument('--output-file', required=True, action=fixPathAction, dest="output_file", help="Input file to be processed. Specify .*-extension to find first file.") + p.add_argument('--reference-file', action=fixPathAction, dest="reference_file", help="Reference file used to determine proper FPS and transfer audio from it. Specify .*-extension to find first file.") + p.add_argument('--ext', dest="ext", default='png', help="Image format (extension) of input files.") + p.add_argument('--fps', type=int, dest="fps", default=None, help="FPS of output file. Overwritten by reference-file.") + p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.") + p.add_argument('--lossless', action="store_true", dest="lossless", default=False, help="PNG codec.") + + p.set_defaults(func=process_videoed_video_from_sequence) + def bad_args(arguments): parser.print_help() exit(0) diff --git a/mainscripts/VideoEd.py b/mainscripts/VideoEd.py new file mode 100644 index 0000000..8849c32 --- /dev/null +++ b/mainscripts/VideoEd.py @@ -0,0 +1,199 @@ +import subprocess +import numpy as np +import ffmpeg +from pathlib import Path +from utils import Path_utils +from 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 = Path_utils.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 output_ext is None: + output_ext = io.input_str ("Output image format (extension)? ( default:png ) : ", "png") + + if fps is None: + fps = io.input_int ("Enter FPS ( ?:help skip:fullfps ) : ", 0, help_message="How many frames of every second of the video will be extracted.") + + for filename in Path_utils.get_image_paths (output_path, ['.'+output_ext]): + Path(filename).unlink() + + job = ffmpeg.input(str(input_file_path)) + + kwargs = {} + if fps != 0: + kwargs.update ({'r':str(fps)}) + + 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 (skip: 00:00:00.000) : ", "00:00:00.000") + + if to_time is None: + to_time = io.input_str ("To time (skip: 00:00:00.000) : ", "00:00:00.000") + + if audio_track_id is None: + audio_track_id = io.input_int ("Specify audio track id. ( skip:0 ) : ", 0) + + if bitrate is None: + bitrate = max (1, io.input_int ("Bitrate of output file in MB/s ? (default:25) : ", 25) ) + + from_time = "0:0:04.34" + to_time = "0:0:10.000" + + 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) + 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 + + if ext is None: + ext = io.input_str ("Input image format (extension)? ( default:png ) : ", "png") + + if factor is None: + factor = np.clip ( io.input_int ("Denoise factor? (1-20 default:5) : ", 5), 1, 20 ) + + job = ( ffmpeg + .input(str ( input_path / ('%5d.'+ext) ) ) + .filter("hqdn3d", factor, factor, 5,5) + .output(str ( input_path / ('%5d.'+ext) ) ) + ) + + try: + job = job.run() + except: + io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) ) + +def video_from_sequence( input_dir, output_file, reference_file=None, ext=None, fps=None, bitrate=None, 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)? ( default:png ) : ", "png") + + if lossless is None: + lossless = io.input_bool ("Use lossless codec ? ( default:no ) : ", 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 = Path_utils.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 ("FPS ? (default:25) : ", 25) ) + + if not lossless and bitrate is None: + bitrate = max (1, io.input_int ("Bitrate of output file in MB/s ? (default:16) : ", 16) ) + + i_in = ffmpeg.input(str (input_path / ('%5d.'+ext)), r=fps) + + output_args = [i_in] + + if 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": "png" + }) + else: + output_kwargs.update ({"c:v": "libx264", + "b:v": "%dM" %(bitrate), + "pix_fmt": "yuv420p", + }) + + output_kwargs.update ({"c:a": "aac", + "b:a": "192k", + "ar" : "48000" + }) + + job = ( ffmpeg.output(*output_args, **output_kwargs).overwrite_output() ) + try: + job = job.run() + except: + io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) ) + diff --git a/requirements-cpu.txt b/requirements-cpu.txt index 94a1f3c..bfb1cc1 100644 --- a/requirements-cpu.txt +++ b/requirements-cpu.txt @@ -8,4 +8,5 @@ tensorflow==1.13.1 scikit-image dlib==19.16.0 tqdm +ffmpeg-python==0.1.17 git+https://www.github.com/keras-team/keras-contrib.git diff --git a/requirements-cuda10.0-cudnn7.4.1.txt b/requirements-cuda10.0-cudnn7.4.1.txt index 1122902..8ab28c1 100644 --- a/requirements-cuda10.0-cudnn7.4.1.txt +++ b/requirements-cuda10.0-cudnn7.4.1.txt @@ -9,4 +9,5 @@ plaidml-keras==0.5.0 scikit-image dlib==19.16.0 tqdm +ffmpeg-python==0.1.17 git+https://www.github.com/keras-team/keras-contrib.git diff --git a/requirements-opencl.txt b/requirements-opencl.txt index 4d85acf..7f51c18 100644 --- a/requirements-opencl.txt +++ b/requirements-opencl.txt @@ -8,4 +8,5 @@ tensorflow==1.13.1 plaidml-keras==0.5.0 scikit-image tqdm +ffmpeg-python==0.1.17 git+https://www.github.com/keras-team/keras-contrib.git diff --git a/utils/Path_utils.py b/utils/Path_utils.py index 6002bbb..5305956 100644 --- a/utils/Path_utils.py +++ b/utils/Path_utils.py @@ -3,7 +3,7 @@ from scandir import scandir image_extensions = [".jpg", ".jpeg", ".png", ".tif", ".tiff"] -def get_image_paths(dir_path): +def get_image_paths(dir_path, image_extensions=image_extensions): dir_path = Path (dir_path) result = [] @@ -38,3 +38,17 @@ def get_all_dir_names_startswith (dir_path, startswith): if x.name.lower().startswith(startswith): result.append ( x.name[len(startswith):] ) return result + +def get_first_file_by_stem (dir_path, stem, exts=None): + dir_path = Path (dir_path) + stem = stem.lower() + + if dir_path.exists(): + for x in list(scandir(str(dir_path))): + if not x.is_file(): + continue + xp = Path(x.path) + if xp.stem.lower() == stem and (exts is None or xp.suffix.lower() in exts): + return xp + + return None \ No newline at end of file