2.2.0 main file rewritten, various of enhancements

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.
+MODELS_AVAILABLE = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo',
-MEM_PER_THREAD = 2.5
+                    'anime_style_art_rgb', 'photo', 'anime_style_art_y']
-
-
-class ArgumentError(Exception):
-    def __init__(self, message):
-        super().__init__(message)
 
 
 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',
+        if self.model_type not in MODELS_AVAILABLE:
-                            'anime_style_art_rgb', 'photo', 'anime_style_art_y']
+            raise InvalidModelType('Specified model type not available')
-        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)
 
     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.output_width, self.output_height), self.upscaled_frames)
-            fm.convert_video(framerate, '{}x{}'.format(self.width, self.height), self.upscaled_frames)
         Avalon.info('Conversion completed')
 
         # Extract and press audio in

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())))
+# Each thread might take up to 2.5 GB during initialization.
-FRAMES = '{}\\frames'.format(EXEC_PATH)  # Folder containing extracted frames
+# (system memory, not to be confused with GPU memory)
-UPSCALED = '{}\\upscaled'.format(EXEC_PATH)  # Folder containing enlarges frames
+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('-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', 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('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb')
+    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
+    if os.path.isfile(args.input):
-    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)
+        """ 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: