diff --git a/bin/upscaler.py b/bin/upscaler.py old mode 100644 new mode 100755 index b7254ba..5b83101 --- a/bin/upscaler.py +++ b/bin/upscaler.py @@ -4,7 +4,7 @@ Name: Video2X Upscaler Author: K4YT3X Date Created: December 10, 2018 -Last Modified: December 10, 2018 +Last Modified: December 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 @@ -13,24 +13,19 @@ Licensed under the GNU General Public License Version 3 (GNU GPL v3), """ from avalon_framework import Avalon +from exceptions import * from ffmpeg import Ffmpeg from fractions import Fraction from waifu2x import Waifu2x import json import os -import psutil import shutil import subprocess import tempfile import threading -# Each thread might take up to 2.5 GB during initialization. -MEM_PER_THREAD = 2.5 - - -class ArgumentError(Exception): - def __init__(self, message): - super().__init__(message) +MODELS_AVAILABLE = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo', + 'anime_style_art_rgb', 'photo', 'anime_style_art_y'] class Upscaler: @@ -66,14 +61,8 @@ class Upscaler: def _check_model_type(self, args): """ Validate upscaling model """ - models_available = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo', - 'anime_style_art_rgb', 'photo', 'anime_style_art_y'] - if self.model_type not in models_available: - Avalon.error('Specified model type not found!') - Avalon.info('Available models:') - for model in models_available: - print(model) - exit(1) + if self.model_type not in MODELS_AVAILABLE: + raise InvalidModelType('Specified model type not available') def _check_arguments(self): # Check if arguments are valid / all necessary argument @@ -82,36 +71,11 @@ class Upscaler: raise ArgumentError('You need to specify the video to process') elif (not self.output_width or not self.output_height) and not self.factor: raise ArgumentError('You must specify output video width and height or upscale factor') - elif not self.video_output: + elif not self.output_video: raise ArgumentError('You need to specify the output video name') elif not self.method: raise ArgumentError('You need to specify the enlarging processing unit') - def _check_memory(self): - """ Check usable memory - Warn the user if insufficient memory is available for - the number of threads that the user have chosen. - """ - memory_available = psutil.virtual_memory().available / (1024 ** 3) - - # If user doesn't even have enough memory to run even one thread - if memory_available < MEM_PER_THREAD: - Avalon.warning('You might have an insufficient amount of memory available to run this program ({} GB)'.format(memory_available)) - Avalon.warning('Proceed with caution') - # If memory available is less than needed, warn the user - elif memory_available < (MEM_PER_THREAD * self.threads): - Avalon.warning('Each waifu2x-caffe thread will require up to 2.5 GB during initialization') - Avalon.warning('You demanded {} threads to be created, but you only have {} GB memory available'.format(self.threads, round(memory_available, 4))) - Avalon.warning('{} GB of memory is recommended for {} threads'.format(MEM_PER_THREAD * self.threads, self.threads)) - Avalon.warning('With your current amount of memory available, {} threads is recommended'.format(int(memory_available // MEM_PER_THREAD))) - - # Ask the user if he / she wants to change to the recommended - # number of threads - if Avalon.ask('Change to the recommended value?', True): - self.threads = int(memory_available // MEM_PER_THREAD) - else: - Avalon.warning('Proceed with caution') - def _upscale_frames(self, w2): """ Upscale video frames with waifu2x-caffe @@ -159,7 +123,7 @@ class Upscaler: # Create threads and start them for thread_info in thread_pool: # Create thread - thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.width, self.height)) + thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.output_width, self.output_height)) thread.name = thread_info[1] # Add threads into the pool @@ -184,7 +148,6 @@ class Upscaler: # Check argument sanity self._check_model_type(self.model_type) self._check_arguments() - self._check_memory() # Convert paths to absolute paths self.input_video = os.path.abspath(self.input_video) @@ -202,7 +165,7 @@ class Upscaler: raise FileNotFoundError(self.waifu2x_path) # Initialize objects for ffmpeg and waifu2x-caffe - fm = Ffmpeg(self.ffmpeg_path, self.video_output, self.ffmpeg_arguments) + fm = Ffmpeg(self.ffmpeg_path, self.output_video, self.ffmpeg_arguments) w2 = Waifu2x(self.waifu2x_path, self.method, self.model_type) # Extract frames from video @@ -228,6 +191,13 @@ class Upscaler: framerate = float(Fraction(info['streams'][video_stream_index]['avg_frame_rate'])) Avalon.info('Framerate: {}'.format(framerate)) + # Width/height will be coded width/height x upscale factor + if self.factor: + coded_width = info['streams'][video_stream_index]['coded_width'] + coded_height = info['streams'][video_stream_index]['coded_height'] + self.output_width = self.factor * coded_width + self.output_height = self.factor * coded_height + # Upscale images one by one using waifu2x Avalon.info('Starting to upscale extracted images') self._upscale_frames(w2) @@ -236,15 +206,8 @@ class Upscaler: # Frames to Video Avalon.info('Converting extracted frames into video') - # Width/height will be coded width/height x upscale factor - if self.factor: - coded_width = info['streams'][video_stream_index]['coded_width'] - coded_height = info['streams'][video_stream_index]['coded_height'] - fm.convert_video(framerate, '{}x{}'.format(self.factor * coded_width, self.factor * coded_height), self.upscaled_frames) - # Use user defined output size - else: - fm.convert_video(framerate, '{}x{}'.format(self.width, self.height), self.upscaled_frames) + fm.convert_video(framerate, '{}x{}'.format(self.output_width, self.output_height), self.upscaled_frames) Avalon.info('Conversion completed') # Extract and press audio in diff --git a/bin/video2x.py b/bin/video2x.py index 6961783..f0a7719 100755 --- a/bin/video2x.py +++ b/bin/video2x.py @@ -10,10 +10,10 @@ __ __ _ _ ___ __ __ \/ |_| \__,_| \___| \___/ |____| /_/ \_\ -Name: Video2x Controller +Name: Video2X Controller Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: December 10, 2018 +Last Modified: December 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 @@ -27,18 +27,19 @@ details or quality, keeping lines smooth and edges sharp. """ from avalon_framework import Avalon from upscaler import Upscaler +from upscaler import MODELS_AVAILABLE import argparse -import inspect import json import os +import psutil import time import traceback -VERSION = '2.1.7' +VERSION = '2.2.0' -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 +# Each thread might take up to 2.5 GB during initialization. +# (system memory, not to be confused with GPU memory) +MEM_PER_THREAD = 2.5 def process_arguments(): @@ -52,24 +53,20 @@ def process_arguments(): # Video options basic_options = parser.add_argument_group('Basic Options') - basic_options.add_argument('-i', '--input', help='Specify source video file', action='store', default=False, required=True) - basic_options.add_argument('-o', '--output', help='Specify output video file', action='store', default=False, required=True) - basic_options.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb') + basic_options.add_argument('-i', '--input', help='Specify source video file/directory', action='store', default=False, required=True) + basic_options.add_argument('-o', '--output', help='Specify output video file/directory', action='store', default=False, required=True) + basic_options.add_argument('-m', '--method', help='Specify upscaling method', action='store', default='gpu', choices=['cpu', 'gpu', 'cudnn'], required=True) + basic_options.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb', choices=MODELS_AVAILABLE) basic_options.add_argument('-t', '--threads', help='Specify number of threads to use for upscaling', action='store', type=int, default=5) basic_options.add_argument('-c', '--config', help='Manually specify config file', action='store', default='video2x.json') # Scaling options - scaling_options = parser.add_argument_group('Scaling Options', required=True) + # scaling_options = parser.add_argument_group('Scaling Options', required=True) # TODO: (width & height) || (factor) + scaling_options = parser.add_argument_group('Scaling Options') # TODO: (width & height) || (factor) scaling_options.add_argument('--width', help='Output video width', action='store', type=int, default=False) scaling_options.add_argument('--height', help='Output video height', action='store', type=int, default=False) scaling_options.add_argument('-f', '--factor', help='Factor to upscale the videos by', action='store', type=int, default=False) - # Render drivers, at least one option must be specified - driver_group = parser.add_argument_group('Render Drivers', required=True) - driver_group.add_argument('--cpu', help='Use CPU for enlarging', action='store_true', default=False) - driver_group.add_argument('--gpu', help='Use GPU for enlarging', action='store_true', default=False) - driver_group.add_argument('--cudnn', help='Use CUDNN for enlarging', action='store_true', default=False) - # Parse arguments return parser.parse_args() @@ -82,14 +79,40 @@ def print_logo(): print(' \\ / | | | (_| | | __/ | (_) | / /_ / . \\') print(' \\/ |_| \\__,_| \\___| \\___/ |____| /_/ \\_\\') print('\n Video2X Video Enlarger') - spaces = ((44 - len("Version " + VERSION)) // 2) * " " + spaces = ((44 - len("Version {}".format(VERSION))) // 2) * " " print('{}\n{} Version {}\n{}'.format(Avalon.FM.BD, spaces, VERSION, Avalon.FM.RST)) +def check_system_memory(): + """ Check usable system memory + Warn the user if insufficient memory is available for + the number of threads that the user have chosen. + """ + memory_available = psutil.virtual_memory().available / (1024 ** 3) + + # If user doesn't even have enough memory to run even one thread + if memory_available < MEM_PER_THREAD: + Avalon.warning('You might have an insufficient amount of memory available to run this program ({} GB)'.format(memory_available)) + Avalon.warning('Proceed with caution') + # If memory available is less than needed, warn the user + elif memory_available < (MEM_PER_THREAD * args.threads): + Avalon.warning('Each waifu2x-caffe thread will require up to 2.5 GB during initialization') + Avalon.warning('You demanded {} threads to be created, but you only have {} GB memory available'.format(args.threads, round(memory_available, 4))) + Avalon.warning('{} GB of memory is recommended for {} threads'.format(MEM_PER_THREAD * args.threads, args.threads)) + Avalon.warning('With your current amount of memory available, {} threads is recommended'.format(int(memory_available // MEM_PER_THREAD))) + + # Ask the user if he / she wants to change to the recommended + # number of threads + if Avalon.ask('Change to the recommended value?', True): + args.threads = int(memory_available // MEM_PER_THREAD) + else: + Avalon.warning('Proceed with caution') + + def read_config(config_file): """ Reads configuration file - Returns a dictionary read by json. + Returns a dictionary read by JSON. """ with open(config_file, 'r') as raw_config: config = json.load(raw_config) @@ -101,38 +124,43 @@ def read_config(config_file): # This is not a library if __name__ != '__main__': Avalon.error('This file cannot be imported') - exit(1) + raise ImportError('{} cannot be imported'.format(__file__)) print_logo() -# Process cli arguments +# Process CLI arguments args = process_arguments() +# Check system available memory +check_system_memory() + # Read configurations from JSON -config = read_config() +config = read_config(args.config) waifu2x_path = config['waifu2x_path'] ffmpeg_path = config['ffmpeg_path'] ffmpeg_arguments = config['ffmpeg_arguments'] ffmpeg_hwaccel = config['ffmpeg_hwaccel'] -# Parse arguments for waifu2x -if args.cpu: - method = 'cpu' -elif args.gpu: - method = 'gpu' - ffmpeg_arguments.append('-hwaccel {}'.format(ffmpeg_hwaccel)) -elif args.cudnn: - method = 'cudnn' - ffmpeg_arguments.append('-hwaccel {}'.format(ffmpeg_hwaccel)) - # Start execution try: # Start timer begin_time = time.time() - # Initialize and run upscaling - upscaler = Upscaler(input_video=args.video, output_video=args.output, method=method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, factor=args.factor, model_type=args.model_type, threads=args.threads) - upscaler.run() + if os.path.isfile(args.input): + """ Upscale single video file """ + Avalon.info('Upscaling single video file: {}'.format(args.input)) + upscaler = Upscaler(input_video=args.input, output_video=args.output, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, factor=args.factor, model_type=args.model_type, threads=args.threads) + upscaler.run() + elif os.path.isdir(args.input): + """ Upscale videos in a folder/directory """ + Avalon.info('Upscaling videos in folder: {}'.format(args.input)) + for input_video in [f for f in os.listdir(args.input) if os.path.isfile(os.path.join(args.input, f))]: + output_video = '{}\\{}'.format(args.output, input_video) + upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, factor=args.factor, model_type=args.model_type, threads=args.threads) + upscaler.run() + else: + Avalon.error('Input path is neither a file nor a folder/directory') + raise FileNotFoundError('{} is neither file nor folder/directory'.format(args.input)) Avalon.info('Program completed, taking {} seconds'.format(round((time.time() - begin_time), 5))) except Exception: diff --git a/bin/video2x_setup.py b/bin/video2x_setup.py old mode 100644 new mode 100755