From 03d6e3f3a2ed081204043f0bde16fedc36586f86 Mon Sep 17 00:00:00 2001 From: K4YT3X Date: Sat, 19 May 2018 00:54:06 -0400 Subject: [PATCH] v2.0 beta: better CUI + enhancements + bug fixes --- video2x.py | 162 +++++++++++++++++++++++++++++++++++------------------ 1 file changed, 108 insertions(+), 54 deletions(-) diff --git a/video2x.py b/video2x.py index 209da43..00ed823 100644 --- a/video2x.py +++ b/video2x.py @@ -13,37 +13,41 @@ __ __ _ _ ___ __ __ Name: Video2x Controller Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: Feb 25, 2018 +Last Modified: May 19, 2018 Licensed under the GNU General Public License Version 3 (GNU GPL v3), available at: https://www.gnu.org/licenses/gpl-3.0.txt -(C) 2016 - 2017 K4YT3X +(C) 2018 K4YT3X Description: Video2X is an automation software based on waifu2x image enlarging engine. It extracts frames from a video, enlarge it by a number of times without losing any details or quality, keeping lines smooth and edges sharp. - -Version 1.1 alpha """ - from ffmpeg import FFMPEG from fractions import Fraction +from tqdm import tqdm from waifu2x import WAIFU2X -import avalon_framework as avalon import argparse +import avalon_framework as avalon +import inspect import json import os +import shutil +import subprocess import traceback -# FFMPEG bin folder. Mind that "/" at the end -FFMPEG_PATH = "C:/Program Files (x86)/ffmpeg/bin/" -# waifu2x executable path. Mind all the forward slashes -WAIFU2X_PATH = "\"C:/Program Files (x86)/waifu2x-caffe/waifu2x-caffe-cui.exe\"" +VERSION = '2.0 beta' -FOLDERIN = "frames" # Folder containing extracted frames -FOLDEROUT = "upscaled" # Folder contaning enlarges frames +EXEC_PATH = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) +FRAMES = '{}\\frames'.format(EXEC_PATH) # Folder containing extracted frames +UPSCALED = '{}\\upscaled'.format(EXEC_PATH) # Folder containing enlarges frames + +# FFMPEG bin folder. Mind that '/' at the end +FFMPEG_PATH = 'C:/Program Files (x86)/ffmpeg/bin/' +# waifu2x executable path. Mind all the forward slashes +WAIFU2X_PATH = '\"C:/Program Files (x86)/waifu2x-caffe/waifu2x-caffe-cui.exe\"' def processArguments(): @@ -53,17 +57,29 @@ def processArguments(): This allows users to customize options for the output video. """ - global args parser = argparse.ArgumentParser() control_group = parser.add_argument_group('Controls') - control_group.add_argument("-f", "--factor", help="Factor to enlarge video by", action="store", default=2) - control_group.add_argument("-v", "--video", help="Specify video file", action="store", default=False) - control_group.add_argument("-o", "--output", help="Specify output file", action="store", default=False) - control_group.add_argument("-y", "--model_type", help="Specify model to use", action="store", default="anime_style_art_rgb") - control_group.add_argument("--cpu", help="Use CPU for enlarging", action="store_true", default=False) - control_group.add_argument("--gpu", help="Use GPU for enlarging", action="store_true", default=False) - control_group.add_argument("--cudnn", help="Use CUDNN for enlarging", action="store_true", default=False) - args = parser.parse_args() + control_group.add_argument('-f', '--factor', help='Factor to enlarge video by', action='store', default=2) + control_group.add_argument('-v', '--video', help='Specify video file', action='store', default=False) + control_group.add_argument('-o', '--output', help='Specify output file', action='store', default=False) + control_group.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb') + control_group.add_argument('--cpu', help='Use CPU for enlarging', action='store_true', default=False) + control_group.add_argument('--gpu', help='Use GPU for enlarging', action='store_true', default=False) + control_group.add_argument('--cudnn', help='Use CUDNN for enlarging', action='store_true', default=False) + return parser.parse_args() + + +def print_logo(): + print('__ __ _ _ ___ __ __') + print('\\ \\ / / (_) | | |__ \\ \\ \\ / /') + print(' \\ \\ / / _ __| | ___ ___ ) | \\ V /') + print(' \\ \\/ / | | / _` | / _ \\ / _ \\ / / > <') + print(' \\ / | | | (_| | | __/ | (_) | / /_ / . \\') + print(' \\/ |_| \\__,_| \\___| \\___/ |____| /_/ \\_\\') + print('\n Video2X Video Enlarger') + spaces = ((44 - len("Version " + VERSION)) // 2) * " " + print(avalon.FM.BD + "\n" + spaces + + ' Version ' + VERSION + '\n' + avalon.FM.RST) def get_vid_info(): @@ -77,19 +93,21 @@ def get_vid_info(): Returns: dictionary -- original video information """ - os.system("{} -v quiet -print_format json -show_format -show_streams {} > info.json".format("\"" + FFMPEG_PATH + "ffprobe.exe\"", args.video)) - json_file = open('info.json', 'r') - json_str = json_file.read() - print(json.loads(json_str)) + json_str = subprocess.check_output( + '{} -v quiet -print_format json -show_format -show_streams {}'.format('\"' + FFMPEG_PATH + 'ffprobe.exe\"', args.video)) return json.loads(json_str) def check_model_type(args): - models_available = ["upconv_7_anime_style_art_rgb", "upconv_7_photo", - "anime_style_art_rgb", "photo", "anime_style_art_y"] + """ + Check if the model demanded from cli + argument is legal. + """ + models_available = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo', + 'anime_style_art_rgb', 'photo', 'anime_style_art_y'] if args.model_type not in models_available: avalon.error('Specified model type not found!') - avalon.info("Available models:") + avalon.info('Available models:') for model in models_available: print(model) exit(1) @@ -105,57 +123,93 @@ def video2x(): check_model_type(args) if args.cpu: - method = "cpu" + method = 'cpu' elif args.gpu: - method = "gpu" + method = 'gpu' elif args.cudnn: - method = "cudnn" + method = 'cudnn' - fm = FFMPEG("\"" + FFMPEG_PATH + "ffmpeg.exe\"", args.output) - w2 = WAIFU2X(WAIFU2X_PATH, method) + fm = FFMPEG('\"' + FFMPEG_PATH + 'ffmpeg.exe\"', args.output) + w2 = WAIFU2X(WAIFU2X_PATH, method, args.model_type) - # Extract Frames - if not os.path.isdir(FOLDERIN): - os.mkdir(FOLDERIN) - fm.extract_frames(args.video, FOLDERIN) + # Clear and create directories + if os.path.isdir(FRAMES): + shutil.rmtree(FRAMES) + if os.path.isdir(UPSCALED): + shutil.rmtree(UPSCALED) + os.mkdir(FRAMES) + os.mkdir(UPSCALED) + + # Extract frames from video + fm.extract_frames(args.video, FRAMES) info = get_vid_info() - # Framerate is read as fraction from the json dictionary - width, height, framerate = info["streams"][0]["width"], info["streams"][0]["height"], float(Fraction(info["streams"][0]["avg_frame_rate"])) - print("Framerate: ", framerate) - final_resolution = str(width * int(args.factor)) + "x" + str(height * int(args.factor)) + # Analyze original video with ffprobe and retrieve framerate + width, height, framerate = info['streams'][0]['width'], info['streams'][0]['height'], float( + Fraction(info['streams'][0]['avg_frame_rate'])) + avalon.info('Framerate: {}'.format(framerate)) + final_resolution = str(width * int(args.factor)) + \ + 'x' + str(height * int(args.factor)) - # Upscale Frames - if not os.path.isdir(FOLDEROUT): - os.mkdir(FOLDEROUT) - w2.upscale(FOLDERIN, FOLDEROUT, int(args.factor) * width, int(args.factor) * height, args.model_type) + # Upscale images one by one using waifu2x + avalon.info('Starting to upscale extracted images') + for (dirpath, dirnames, filenames) in os.walk(FRAMES): + file_list = tqdm(filenames, ascii=True) + for file in file_list: + if file[-4:].lower() == '.png': + image_path = '{}\\{}'.format(dirpath, file) + file_list.set_description('Upscaling: {}'.format(file)) + # avalon.dbgInfo('Upscaling: {}'.format(image_path)) + w2.upscale(image_path, UPSCALED, int(args.factor) * + width, int(args.factor) * height) + avalon.info('Extraction complete') # Frames to Video - fm.to_vid(framerate, final_resolution, FOLDEROUT) + avalon.info('Converting extracted frames into video') + fm.to_vid(framerate, final_resolution, UPSCALED) # Extract and press audio in - fm.extract_audio(args.video, FOLDEROUT) - fm.insert_audio_track("output.mp4", FOLDEROUT) + avalon.info('Stripping audio track from original video') + fm.extract_audio(args.video, UPSCALED) + avalon.info('Inserting audio track into new video') + fm.insert_audio_track(UPSCALED) -processArguments() +# /////////////////// Execution /////////////////// # + +args = processArguments() +# Convert paths to absolute paths +args.video = os.path.abspath(args.video) +args.output = os.path.abspath(args.output) +print_logo() + +if not os.path.isdir(FFMPEG_PATH): + avalon.error('FFMPEG binaries not found') + avalon.error('Please specify FFMPEG binaries location in source code') + print('Current value: {}\n'.format(FFMPEG_PATH)) + raise FileNotFoundError('FFMPEG binaries not found') +if not os.path.isfile(WAIFU2X_PATH.strip('\"')): + avalon.error('Waifu2x CUI executable not found') + avalon.error('Please specify Waifu2x CUI location in source code') + print('Current value: {}\n'.format(WAIFU2X_PATH)) + raise FileNotFoundError('Waifu2x CUI executable not found') + # Check if arguments are valid / all necessary argument # values are specified if not args.video: - print("Error: You need to specify the video to process") + avalon.error('You need to specify the video to process') exit(1) elif not args.output: - print("Error: You need to specify the output video name") + avalon.error('You need to specify the output video name') exit(1) elif not args.cpu and not args.gpu and not args.cudnn: - print("Error: You need to specify the enlarging processing unit") + avalon.error('You need to specify the enlarging processing unit') exit(1) if __name__ == '__main__': try: video2x() except Exception as e: - # This code block is reserved for future - # fail-safe handlers + avalon.error('An exception occurred') traceback.print_exc()