Merge remote-tracking branch 'origin/master'

added alpha version of video2x gui

README.md

@@ -6,9 +6,9 @@
 
 You can go to the [releases page](https://github.com/k4yt3x/video2x/releases) to download the latest builds of `Video2X`. The exe files will require no Python or Python module installation.
 
-The **`full`** package provides all packages that will possibly be needed by `Video2X`, including `FFmpeg`, `waifu2x-caffe` and `waifu2x-converter-cpp`. The config file (`video2x.json`) is also already configured for the environment. All you need to do is just to launch `video2x.exe`.
+The **`full`** package provides all packages that will possibly be needed by `Video2X`, including `FFmpeg`, `waifu2x-caffe`, `waifu2x-converter-cpp`, and `waifu2x-ncnn-vulkan`. The config file (`video2x.json`) is also already configured for the environment. All you need to do is just to launch `video2x.exe`.
 
-The **`light`** package provides only the most basic functions of `Video2X`. Only `video2x.exe`, `video2x_setup.exe` and `video2x.json` are included. To setup `FFmpeg`, `waifu2x-caffe` or `waifu2x-converter-cpp` automatically, simply launch `video2x_setup.exe`.
+The **`light`** package provides only the most basic functions of `Video2X`. Only `video2x.exe`, `video2x_setup.exe` and `video2x.json` are included. To setup dependencies (e.g. `FFmpeg` and `Waifu2X`) automatically, simply launch `video2x_setup.exe`.
 
 ## Prerequisites
 
@@ -18,24 +18,32 @@ Component names that are **bolded** can be automatically downloaded and configur
 2. AMD GPU / Nvidia GPU
 3. AMD GPU driver / Nvidia GPU driver / Nvidia CUDNN
 4. [**FFmpeg**](https://ffmpeg.zeranoe.com/builds/)
-5. [**waifu2x-caffe**](https://github.com/lltcggie/waifu2x-caffe/releases) / [**waifu2x-converter-cpp**](https://github.com/DeadSix27/waifu2x-converter-cpp/releases)
+5. [**waifu2x-caffe**](https://github.com/lltcggie/waifu2x-caffe/releases) / [**waifu2x-converter-cpp**](https://github.com/DeadSix27/waifu2x-converter-cpp/releases) / [**waifu2x-ncnn-vulkan**](https://github.com/nihui/waifu2x-ncnn-vulkan)
 
 ## Recent Changes
 
+### 2.9.0 (July 27, 2019)
+
+- Changed file handling method from `os` to `pathlib`
+- Removed f_string dependency and support for legacy versions of Python
+- Organized file import statements
+
+### 2.8.1 (July 9, 2019)
+
+- Added automatic pixel format detection
+- Added automatic color bit depth detection
+
+### 2.8.0 (June 25, 2019)
+
+- Added support for [waifu2x-ncnn-vulkan](https://github.com/nihui/waifu2x-ncnn-vulkan)
+
 ### 2.7.1 (April 18, 2019)
 
 - Fixed video2x custom temp folder bug found by @cr08 .
 
-### 2.7.0 (March 30, 2019)
+### Setup Script 1.3.0 (June 25, 2019)
 
-- Added support for different extracted image formats.
-- Redesigned FFmpeg wrapper, FFmpeg settings are now customizable in the `video2x.json` config file.
-- Other minor enhancements and adjustments (e.g. argument -> method variable)
-
-### Setup Script 1.2.0 (March 26, 2019)
-
-- `video2x_setup.py` script can now automatically download and configure `waifu2x-converter-cpp`.
-- replaced old progress indicator with progress bar.
+- Added automatic installation support for `waifu2x-ncnn-vulkan`
 
 ## Description
 
@@ -101,7 +109,7 @@ cd video2x/bin
 
 Then you may run the `video2x_setup.py` script to install and configure the dependencies automatically. This script is designed and tested on Windows 10.
 
-This script will install the newest version of `ffmpeg`, either or both of `waifu2x-caffe` and `waifu2x-converter-cpp` to `%LOCALAPPDATA%\\video2x` and all required python libraries.
+This script will install the newest version of `ffmpeg`, any one or all `waifu2x-caffe`, `waifu2x-converter-cpp`, and `waifu2x-ncnn-vulkan` to `%LOCALAPPDATA%\\video2x` and all required python libraries.
 
 ```shell
 python video2x_setup.py
@@ -152,6 +160,12 @@ Enlarge the video by 2 times using OpenCL. Note that `waifu2x-converter-cpp` doe
 python video2x.py -i sample_input.mp4 -o sample_output.mp4 -m gpu -r 2 -d waifu2x_converter
 ```
 
+### AMD or Nvidia (waifu2x-ncnn-vulkan Vulkan)
+
+```shell
+python video2x.py -i sample_input.mp4 -o sample_output.mp4 -m gpu -r 2 -d waifu2x_ncnn_vulkan
+```
+
 ### CPU
 
 Enlarge the video to 1920x1080 using the CPU. You may also use the `-r/--ratio` option. This is potentially much slower than using a GPU. The configuration file for this method is similar to the previous methods.
@@ -226,13 +240,15 @@ This project relies on the following software and projects.
 - [FFmpeg]('https://www.ffmpeg.org/')
 - [waifu2x-caffe](https://github.com/lltcggie/waifu2x-caffe)
 - [waifu2x-converter-cpp](https://github.com/DeadSix27/waifu2x-converter-cpp)
+- [waifu2x-ncnn-vulkan](https://github.com/nihui/waifu2x-ncnn-vulkan)
 
 ## Special Thanks
 
-Appreciations given to the following contributors:
+Appreciations given to the following code contributors:
 
 - @BrianPetkovsek
+- @SAT3LL
 
 ## Related Resources
 
-- [Dandere2x](https://github.com/CardinalPanda/dandere2x): `Dandere2x` is a lossy video upscaler also built around `waifu2x`, but with video compression techniques to shorten the time needed to process a video.
+- [Dandere2x](https://github.com/CardinalPanda/dandere2x): `Dandere2x` is a lossy video upscaler also built around `waifu2x`, but with video compression techniques to shorten the time needed to process a video.

bin/exceptions.py

@@ -1,13 +1,28 @@
 #!/usr/bin/env python3
-# -*- code:utf-8 -*-
+# -*- coding: utf-8 -*-
 """
 Name: Video2X Exceptions
 Dev: K4YT3X
 Date Created: December 13, 2018
-Last Modified: March 19, 2019
+Last Modified: July 27, 2019
 """
 
 
 class ArgumentError(Exception):
     def __init__(self, message):
         super().__init__(message)
+
+
+class StreamNotFoundError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+
+
+class UnrecognizedDriverError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+
+
+class UnsupportedPixelError(Exception):
+    def __init__(self, message):
+        super().__init__(message)

bin/ffmpeg.py

@@ -1,23 +1,28 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
-Name: FFMPEG Class
+Name: Video2X FFmpeg Controller
 Author: K4YT3X
 Date Created: Feb 24, 2018
-Last Modified: May 4, 2019
+Last Modified: July 27, 2019
 
-Description: This class handles all FFMPEG related
-operations.
+Description: This class handles all FFmpeg related operations.
 """
-from avalon_framework import Avalon
+
+# built-in imports
 import json
+import os
+import pathlib
 import subprocess
 
+# third-party imports
+from avalon_framework import Avalon
+
 
 class Ffmpeg:
-    """This class communicates with ffmpeg
+    """This class communicates with FFmpeg
 
-    This class deals with ffmpeg. It handles extracitng
+    This class deals with FFmpeg. It handles extracting
     frames, stripping audio, converting images into videos
     and inserting audio tracks to videos.
     """
@@ -25,21 +30,53 @@ class Ffmpeg:
     def __init__(self, ffmpeg_settings, image_format):
         self.ffmpeg_settings = ffmpeg_settings
 
-        self.ffmpeg_path = self.ffmpeg_settings['ffmpeg_path']
-        # add a forward slash to directory if not present
-        # otherwise there will be a format error
-        if self.ffmpeg_path[-1] != '/' and self.ffmpeg_path[-1] != '\\':
-            self.ffmpeg_path = f'{self.ffmpeg_path}\\'
-
-        self.ffmpeg_binary = f'{self.ffmpeg_path}ffmpeg.exe'
-        self.ffmpeg_probe_binary = f'{self.ffmpeg_path}ffprobe.exe'
+        self.ffmpeg_path = pathlib.Path(self.ffmpeg_settings['ffmpeg_path'])
+        self.ffmpeg_binary = self.ffmpeg_path / 'ffmpeg.exe'
+        self.ffmpeg_probe_binary = self.ffmpeg_path / 'ffprobe.exe'
         self.image_format = image_format
+        self.pixel_format = None
+
+    def get_pixel_formats(self):
+        """ Get a dictionary of supported pixel formats
+
+        List all supported pixel formats and their
+        corresponding bit depth.
+
+        Returns:
+            dictionary -- JSON dict of all pixel formats to bit depth
+        """
+        execute = [
+            self.ffmpeg_probe_binary,
+            '-v',
+            'quiet',
+            '-pix_fmts'
+        ]
+
+        # turn elements into str
+        execute = [str(e) for e in execute]
+
+        Avalon.debug_info(f'Executing: {" ".join(execute)}')
+
+        # initialize dictionary to store pixel formats
+        pixel_formats = {}
+
+        # record all pixel formats into dictionary
+        for line in subprocess.run(execute, check=True, stdout=subprocess.PIPE).stdout.decode().split('\n'):
+            try:
+                pixel_formats[' '.join(line.split()).split()[1]] = int(' '.join(line.split()).split()[3])
+            except (IndexError, ValueError):
+                pass
+
+        # print pixel formats for debugging
+        Avalon.debug_info(pixel_formats)
+
+        return pixel_formats
 
     def get_video_info(self, input_video):
         """ Gets input video information
 
         This method reads input video information
-        using ffprobe in dictionary.
+        using ffprobe in dictionary
 
         Arguments:
             input_video {string} -- input video file path
@@ -62,6 +99,9 @@ class Ffmpeg:
             input_video
         ]
 
+        # turn elements into str
+        execute = [str(e) for e in execute]
+
         Avalon.debug_info(f'Executing: {" ".join(execute)}')
         json_str = subprocess.run(execute, check=True, stdout=subprocess.PIPE).stdout
         return json.loads(json_str.decode('utf-8'))
@@ -69,8 +109,7 @@ class Ffmpeg:
     def extract_frames(self, input_video, extracted_frames):
         """Extract every frame from original videos
 
-        This method extracts every frame from videoin
-        using ffmpeg
+        This method extracts every frame from input video using FFmpeg
 
         Arguments:
             input_video {string} -- input video path
@@ -80,21 +119,25 @@ class Ffmpeg:
             self.ffmpeg_binary
         ]
 
-        execute.extend(self._read_configuration(phase='video_to_frames'))
-
         execute.extend([
             '-i',
-            input_video,
-            f'{extracted_frames}\\extracted_%0d.{self.image_format}'
+            input_video
         ])
 
+        execute.extend(self._read_configuration(phase='video_to_frames', section='output_options'))
+
+        execute.extend([
+            extracted_frames / f'extracted_%0d.{self.image_format}'
+        ])
+
+        execute.extend(self._read_configuration(phase='video_to_frames'))
+
         self._execute(execute)
 
     def convert_video(self, framerate, resolution, upscaled_frames):
         """Converts images into videos
 
-        This method converts a set of images into a
-        video.
+        This method converts a set of images into a video
 
         Arguments:
             framerate {float} -- target video framerate
@@ -112,10 +155,20 @@ class Ffmpeg:
         # read FFmpeg input options
         execute.extend(self._read_configuration(phase='frames_to_video', section='input_options'))
 
+        # WORKAROUND FOR WAIFU2X-NCNN-VULKAN
+        # Dev: SAT3LL
+        # rename all .png.png suffixes to .png
+        import re
+        import shutil
+        regex = re.compile(r'\.png\.png$')
+        for frame_name in upscaled_frames.iterdir():
+            (upscaled_frames / frame_name).rename(upscaled_frames / regex.sub('.png', str(frame_name)))
+        # END WORKAROUND
+
         # append input frames path into command
         execute.extend([
             '-i',
-            f'{upscaled_frames}\\extracted_%d.{self.image_format}'
+            upscaled_frames / f'extracted_%d.{self.image_format}'
         ])
 
         # read FFmpeg output options
@@ -126,7 +179,7 @@ class Ffmpeg:
 
         # specify output file location
         execute.extend([
-            f'{upscaled_frames}\\no_audio.mp4'
+            upscaled_frames / 'no_audio.mp4'
         ])
 
         self._execute(execute)
@@ -142,12 +195,12 @@ class Ffmpeg:
         execute = [
             self.ffmpeg_binary,
             '-i',
-            f'{upscaled_frames}\\no_audio.mp4',
+            upscaled_frames / 'no_audio.mp4',
             '-i',
             input_video
         ]
 
-        execute.extend(self._read_configuration(phase='frames_to_video', section='output_options'))
+        execute.extend(self._read_configuration(phase='migrating_tracks', section='output_options'))
 
         execute.extend([
             output_video
@@ -175,6 +228,13 @@ class Ffmpeg:
         # from only that section
         if section:
             source = self.ffmpeg_settings[phase][section].keys()
+
+            # if pixel format is not specified, use the source pixel format
+            try:
+                if self.ffmpeg_settings[phase][section].get('-pix_fmt') is None:
+                    self.ffmpeg_settings[phase][section]['-pix_fmt'] = self.pixel_format
+            except KeyError:
+                pass
         else:
             source = self.ffmpeg_settings[phase].keys()
 
@@ -186,8 +246,18 @@ class Ffmpeg:
                 value = self.ffmpeg_settings[phase][key]
 
             # null or None means that leave this option out (keep default)
-            if value is None or value is False or isinstance(value, list) or isinstance(value, dict):
+            if value is None or value is False or isinstance(value, dict):
                 continue
+
+            # if the value is a list, append the same argument and all values
+            elif isinstance(value, list):
+
+                for subvalue in value:
+                    configuration.append(key)
+                    if value is not True:
+                        configuration.append(str(subvalue))
+
+            # otherwise the value is typical
             else:
                 configuration.append(key)
 
@@ -209,4 +279,8 @@ class Ffmpeg:
             int -- execution return code
         """
         Avalon.debug_info(f'Executing: {execute}')
+
+        # turn all list elements into string to avoid errors
+        execute = [str(e) for e in execute]
+
         return subprocess.run(execute, shell=True, check=True).returncode

bin/image_cleaner.py

@@ -5,12 +5,13 @@ Name: Video2X Image Cleaner
 Author: BrianPetkovsek
 Author: K4YT3X
 Date Created: March 24, 2019
-Last Modified: April 28, 2019
+Last Modified: July 27, 2019
 
 Description: This class is to remove the extracted frames
 that have already been upscaled.
 """
 
+# built-in imports
 import os
 import threading
 import time
@@ -58,19 +59,19 @@ class ImageCleaner(threading.Thread):
         """
 
         # list all images in the extracted frames
-        output_frames = [f for f in os.listdir(self.output_directory) if os.path.isfile(os.path.join(self.output_directory, f))]
+        output_frames = [f for f in self.output_directory.iterdir() if f.is_file()]
 
         # compare and remove frames downscaled images that finished being upscaled
         # within each thread's  extracted frames directory
-        for i in range(self.threads):
-            dir_path = os.path.join(self.input_directory, str(i))
+        for thread_id in range(self.threads):
+            dir_path = self.input_directory / str(thread_id)
 
             # for each file within all the directories
-            for f in os.listdir(dir_path):
-                file_path = os.path.join(dir_path, f)
+            for file in dir_path.iterdir():
+                file_path = dir_path / file
 
                 # if file also exists in the output directory, then the file
                 # has already been processed, thus not needed anymore
-                if os.path.isfile(file_path) and f in output_frames:
-                    os.remove(file_path)
-                    output_frames.remove(f)
+                if file_path.is_file() and file in output_frames:
+                    file_path.unlink(file)
+                    output_frames.remove(file)

bin/upscaler.py

@@ -4,7 +4,9 @@
 Name: Video2X Upscaler
 Author: K4YT3X
 Date Created: December 10, 2018
-Last Modified: April 28, 2019
+Last Modified: July 27, 2019
+
+Dev: SAT3LL
 
 Licensed under the GNU General Public License Version 3 (GNU GPL v3),
     available at: https://www.gnu.org/licenses/gpl-3.0.txt
@@ -12,20 +14,28 @@ Licensed under the GNU General Public License Version 3 (GNU GPL v3),
 (C) 2018-2019 K4YT3X
 """
 
-from avalon_framework import Avalon
+# local imports
 from exceptions import *
 from ffmpeg import Ffmpeg
-from fractions import Fraction
 from image_cleaner import ImageCleaner
-from tqdm import tqdm
 from waifu2x_caffe import Waifu2xCaffe
 from waifu2x_converter import Waifu2xConverter
-import os
+from waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan
+
+# built-in imports
+from fractions import Fraction
+import copy
+import pathlib
 import re
 import shutil
 import tempfile
 import threading
 import time
+import traceback
+
+# third-party imports
+from avalon_framework import Avalon
+from tqdm import tqdm
 
 
 class Upscaler:
@@ -52,16 +62,16 @@ class Upscaler:
         self.scale_ratio = None
         self.model_dir = None
         self.threads = 5
-        self.video2x_cache_directory = f'{tempfile.gettempdir()}\\video2x'
+        self.video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'
         self.image_format = 'png'
         self.preserve_frames = False
 
     def create_temp_directories(self):
         """create temporary directory
         """
-        self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_directory)
+        self.extracted_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory))
         Avalon.debug_info(f'Extracted frames are being saved to: {self.extracted_frames}')
-        self.upscaled_frames = tempfile.mkdtemp(dir=self.video2x_cache_directory)
+        self.upscaled_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory))
         Avalon.debug_info(f'Upscaled frames are being saved to: {self.upscaled_frames}')
 
     def cleanup_temp_directories(self):
@@ -75,7 +85,8 @@ class Upscaler:
                     print(f'Cleaning up cache directory: {directory}')
                     shutil.rmtree(directory)
                 except (OSError, FileNotFoundError):
-                    pass
+                    print(f'Unable to delete: {directory}')
+                    traceback.print_exc()
 
     def _check_arguments(self):
         # check if arguments are valid / all necessary argument
@@ -97,12 +108,13 @@ class Upscaler:
         and the output directory. This is originally
         suggested by @ArmandBernard.
         """
-        # get number of extracted frames
-        total_frames = 0
-        for directory in extracted_frames_directories:
-            total_frames += len([f for f in os.listdir(directory) if f[-4:] == '.png'])
 
-        with tqdm(total=total_frames, ascii=True, desc='Upscaling Progress') as progress_bar:
+        # get number of extracted frames
+        self.total_frames = 0
+        for directory in extracted_frames_directories:
+            self.total_frames += len([f for f in directory.iterdir() if str(f)[-4:] == f'.{self.image_format}'])
+
+        with tqdm(total=self.total_frames, ascii=True, desc='Upscaling Progress') as progress_bar:
 
             # tqdm update method adds the value to the progress
             # bar instead of setting the value. Therefore, a delta
@@ -111,12 +123,12 @@ class Upscaler:
             while not self.progress_bar_exit_signal:
 
                 try:
-                    total_frames_upscaled = len([f for f in os.listdir(self.upscaled_frames) if f[-4:] == '.png'])
-                    delta = total_frames_upscaled - previous_cycle_frames
-                    previous_cycle_frames = total_frames_upscaled
+                    self.total_frames_upscaled = len([f for f in self.upscaled_frames.iterdir() if str(f)[-4:] == f'.{self.image_format}'])
+                    delta = self.total_frames_upscaled - previous_cycle_frames
+                    previous_cycle_frames = self.total_frames_upscaled
 
                     # if upscaling is finished
-                    if total_frames_upscaled >= total_frames:
+                    if self.total_frames_upscaled >= self.total_frames:
                         return
 
                     # adds the detla into the progress bar
@@ -126,7 +138,7 @@ class Upscaler:
 
                 time.sleep(1)
 
-    def _upscale_frames(self, w2):
+    def _upscale_frames(self):
         """ Upscale video frames with waifu2x-caffe
 
         This function upscales all the frames extracted
@@ -143,95 +155,129 @@ class Upscaler:
         # if this thread is not empty, then an exception has occured
         self.upscaler_exceptions = []
 
+        # initialize waifu2x driver
+        drivers = ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan']
+        if self.waifu2x_driver not in drivers:
+            raise UnrecognizedDriverError(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
+
         # it's easier to do multi-threading with waifu2x_converter
         # the number of threads can be passed directly to waifu2x_converter
         if self.waifu2x_driver == 'waifu2x_converter':
+            w2 = Waifu2xConverter(self.waifu2x_settings, self.model_dir)
 
             progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],))
             progress_bar.start()
 
             w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.image_format, self.upscaler_exceptions)
-            for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]:
+            for image in [f for f in self.upscaled_frames.iterdir() if f.is_file()]:
                 renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', image)
-                shutil.move(f'{self.upscaled_frames}\\{image}', f'{self.upscaled_frames}\\{renamed}')
+                (self.upscaled_frames / image).rename(self.upscaled_frames / renamed)
 
             self.progress_bar_exit_signal = True
             progress_bar.join()
             return
+        else:
+            # create a container for all upscaler threads
+            upscaler_threads = []
 
-        # create a container for all upscaler threads
-        upscaler_threads = []
+            # list all images in the extracted frames
+            frames = [(self.extracted_frames / f) for f in self.extracted_frames.iterdir() if f.is_file]
 
-        # list all images in the extracted frames
-        frames = [os.path.join(self.extracted_frames, f) for f in os.listdir(self.extracted_frames) if os.path.isfile(os.path.join(self.extracted_frames, f))]
+            # if we have less images than threads,
+            # create only the threads necessary
+            if len(frames) < self.threads:
+                self.threads = len(frames)
 
-        # if we have less images than threads,
-        # create only the threads necessary
-        if len(frames) < self.threads:
-            self.threads = len(frames)
+            # create a directory for each thread and append directory
+            # name into a list
 
-        # create a directory for each thread and append directory
-        # name into a list
+            thread_pool = []
+            thread_directories = []
+            for thread_id in range(self.threads):
+                thread_directory = self.extracted_frames / str(thread_id)
+                thread_directories.append(thread_directory)
 
-        thread_pool = []
-        thread_directories = []
-        for thread_id in range(self.threads):
-            thread_directory = f'{self.extracted_frames}\\{str(thread_id)}'
-            thread_directories.append(thread_directory)
+                # delete old directories and create new directories
+                if thread_directory.is_dir():
+                    shutil.rmtree(thread_directory)
+                thread_directory.mkdir(parents=True, exist_ok=True)
 
-            # delete old directories and create new directories
-            if os.path.isdir(thread_directory):
-                shutil.rmtree(thread_directory)
-            os.mkdir(thread_directory)
+                # append directory path into list
+                thread_pool.append((thread_directory, thread_id))
 
-            # append directory path into list
-            thread_pool.append((thread_directory, thread_id))
+            # evenly distribute images into each directory
+            # until there is none left in the directory
+            for image in frames:
+                # move image
+                image.rename(thread_pool[0][0] / image.name)
+                # rotate list
+                thread_pool = thread_pool[-1:] + thread_pool[:-1]
 
-        # evenly distribute images into each directory
-        # until there is none left in the directory
-        for image in frames:
-            # move image
-            shutil.move(image, thread_pool[0][0])
-            # rotate list
-            thread_pool = thread_pool[-1:] + thread_pool[:-1]
+            # create threads and start them
+            for thread_info in thread_pool:
 
-        # create threads and start them
-        for thread_info in thread_pool:
-            # create thread
-            if self.scale_ratio:
-                thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.image_format, self.upscaler_exceptions))
-            else:
-                thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, self.image_format, self.upscaler_exceptions))
-            thread.name = thread_info[1]
+                # create a separate w2 instance for each thread
+                if self.waifu2x_driver == 'waifu2x_caffe':
+                    w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir, self.bit_depth)
+                    if self.scale_ratio:
+                        thread = threading.Thread(target=w2.upscale,
+                                                  args=(thread_info[0],
+                                                        self.upscaled_frames,
+                                                        self.scale_ratio,
+                                                        False,
+                                                        False,
+                                                        self.image_format,
+                                                        self.upscaler_exceptions))
+                    else:
+                        thread = threading.Thread(target=w2.upscale,
+                                                  args=(thread_info[0],
+                                                        self.upscaled_frames,
+                                                        False,
+                                                        self.scale_width,
+                                                        self.scale_height,
+                                                        self.image_format,
+                                                        self.upscaler_exceptions))
 
-            # add threads into the pool
-            upscaler_threads.append(thread)
+                # if the driver being used is waifu2x_ncnn_vulkan
+                elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan':
+                    w2 = Waifu2xNcnnVulkan(copy.deepcopy(self.waifu2x_settings))
+                    thread = threading.Thread(target=w2.upscale,
+                                              args=(thread_info[0],
+                                                    self.upscaled_frames,
+                                                    self.scale_ratio,
+                                                    self.upscaler_exceptions))
 
-        # start progress bar in a different thread
-        progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,))
-        progress_bar.start()
+                # create thread
+                thread.name = thread_info[1]
 
-        # create the clearer and start it
-        Avalon.debug_info('Starting upscaled image cleaner')
-        image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
-        image_cleaner.start()
+                # add threads into the pool
+                upscaler_threads.append(thread)
 
-        # start all threads
-        for thread in upscaler_threads:
-            thread.start()
+            # start progress bar in a different thread
+            progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,))
+            progress_bar.start()
 
-        # wait for threads to finish
-        for thread in upscaler_threads:
-            thread.join()
+            # create the clearer and start it
+            Avalon.debug_info('Starting upscaled image cleaner')
+            image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
+            image_cleaner.start()
 
-        # upscaling done, kill the clearer
-        Avalon.debug_info('Killing upscaled image cleaner')
-        image_cleaner.stop()
+            # start all threads
+            for thread in upscaler_threads:
+                thread.start()
 
-        self.progress_bar_exit_signal = True
+            # wait for threads to finish
+            for thread in upscaler_threads:
+                thread.join()
 
-        if len(self.upscaler_exceptions) != 0:
-            raise(self.upscaler_exceptions[0])
+            # upscaling done, kill the clearer
+            Avalon.debug_info('Killing upscaled image cleaner')
+            image_cleaner.stop()
+
+            self.progress_bar_exit_signal = True
+
+            if len(self.upscaler_exceptions) != 0:
+                raise(self.upscaler_exceptions[0])
 
     def run(self):
         """Main controller for Video2X
@@ -245,20 +291,12 @@ class Upscaler:
         self._check_arguments()
 
         # convert paths to absolute paths
-        self.input_video = os.path.abspath(self.input_video)
-        self.output_video = os.path.abspath(self.output_video)
+        self.input_video = self.input_video.absolute()
+        self.output_video = self.output_video.absolute()
 
         # initialize objects for ffmpeg and waifu2x-caffe
         fm = Ffmpeg(self.ffmpeg_settings, self.image_format)
 
-        # initialize waifu2x driver
-        if self.waifu2x_driver == 'waifu2x_caffe':
-            w2 = Waifu2xCaffe(self.waifu2x_settings, self.method, self.model_dir)
-        elif self.waifu2x_driver == 'waifu2x_converter':
-            w2 = Waifu2xConverter(self.waifu2x_settings, self.model_dir)
-        else:
-            raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
-
         # extract frames from video
         fm.extract_frames(self.input_video, self.extracted_frames)
 
@@ -277,10 +315,21 @@ class Upscaler:
         # exit if no video stream found
         if video_stream_index is None:
             Avalon.error('Aborting: No video stream found')
-            exit(1)
+            raise StreamNotFoundError('no video stream found')
 
         # get average frame rate of video stream
         framerate = float(Fraction(video_info['streams'][video_stream_index]['avg_frame_rate']))
+        fm.pixel_format = video_info['streams'][video_stream_index]['pix_fmt']
+
+        # get a dict of all pixel formats and corresponding bit depth
+        pixel_formats = fm.get_pixel_formats()
+
+        try:
+            self.bit_depth = pixel_formats[fm.pixel_format]
+        except KeyError:
+            Avalon.error(f'Unsupported pixel format: {fm.pixel_format}')
+            raise UnsupportedPixelError(f'unsupported pixel format {fm.pixel_format}')
+
         Avalon.info(f'Framerate: {framerate}')
 
         # width/height will be coded width/height x upscale factor
@@ -292,7 +341,7 @@ class Upscaler:
 
         # upscale images one by one using waifu2x
         Avalon.info('Starting to upscale extracted images')
-        self._upscale_frames(w2)
+        self._upscale_frames()
         Avalon.info('Upscaling completed')
 
         # frames to Video

bin/video2x.json

@@ -39,36 +39,53 @@
     "output": null,
     "input": null
   },
+  "waifu2x_ncnn_vulkan": {
+    "waifu2x_ncnn_vulkan_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\waifu2x-ncnn-vulkan\\waifu2x-ncnn-vulkan.exe",
+    "input": null,
+    "output": null,
+    "noise-level": 2,
+    "scale-ratio": null,
+    "tile-size": 200,
+    "model-path": null,
+    "gpu": 0,
+    "load-proc-save_threads": null,
+    "verbose": null
+  },
   "ffmpeg": {
     "ffmpeg_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\ffmpeg-latest-win64-static\\bin",
     "video_to_frames": {
+      "output_options": {
         "-qscale:v": null,
-        "-hwaccel": "auto",
-        "-y": true
+        "-pix_fmt": "rgba64be"
+      },
+      "-hwaccel": "auto",
+      "-y": true
     },
     "frames_to_video": {
-        "input_options":{
-            "-qscale:v": null,
-            "-qscale:a": null,
-            "-f": "image2"
-        },
-        "output_options":{
-            "-vcodec": "libx264",
-            "-crf": 17,
-            "-b:v": null,
-            "-pix_fmt": "yuv420p"
-        },
-        "-y": true
+      "input_options": {
+        "-qscale:v": null,
+        "-qscale:a": null,
+        "-f": "image2"
+      },
+      "output_options": {
+        "-vcodec": "libx264",
+        "-crf": 17,
+        "-b:v": null,
+        "-pix_fmt": null
+      },
+      "-y": true
     },
     "migrating_tracks": {
-        "output_options":{
-            "-map": "0:v:0?",
-            "-map": "1?",
-            "-c": "copy",
-            "-map": "-1:v?",
-            "-pix_fmt": "yuv420p"
-        },
-        "-y": true
+      "output_options": {
+        "-map": [
+          "0:v:0?",
+          "1?",
+          "-1:v?"
+        ],
+        "-c": "copy",
+        "-pix_fmt": null
+      },
+      "-y": true
     }
   },
   "video2x": {

bin/video2x.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-"""
+r"""
 
 __      __  _       _                  ___   __   __
 \ \    / / (_)     | |                |__ \  \ \ / /
@@ -13,7 +13,10 @@ __      __  _       _                  ___   __   __
 Name: Video2X Controller
 Author: K4YT3X
 Date Created: Feb 24, 2018
-Last Modified: April 28, 2019
+Last Modified: July 27, 2019
+
+Dev: BrianPetkovsek
+Dev: SAT3LL
 
 Licensed under the GNU General Public License Version 3 (GNU GPL v3),
     available at: https://www.gnu.org/licenses/gpl-3.0.txt
@@ -38,13 +41,17 @@ 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.
 """
-from avalon_framework import Avalon
+
+# local imports
+from exceptions import *
 from upscaler import Upscaler
+
+# built-in imports
 import argparse
-import GPUtil
+import glob
 import json
 import os
-import psutil
+import pathlib
 import re
 import shutil
 import sys
@@ -52,7 +59,12 @@ import tempfile
 import time
 import traceback
 
-VERSION = '2.7.1'
+# third-party imports
+from avalon_framework import Avalon
+import GPUtil
+import psutil
+
+VERSION = '2.9.0'
 
 # each thread might take up to 2.5 GB during initialization.
 # (system memory, not to be confused with GPU memory)
@@ -71,35 +83,43 @@ def process_arguments():
 
     # video options
     file_options = parser.add_argument_group('File Options')
-    file_options.add_argument('-i', '--input', help='Source video file/directory', action='store', required=True)
-    file_options.add_argument('-o', '--output', help='Output video file/directory', action='store', required=True)
+    file_options.add_argument('-i', '--input', type=pathlib.Path, help='source video file/directory', action='store')
+    file_options.add_argument('-o', '--output', type=pathlib.Path, help='output video file/directory', action='store')
 
     # upscaler options
     upscaler_options = parser.add_argument_group('Upscaler Options')
-    upscaler_options.add_argument('-m', '--method', help='Upscaling method', action='store', default='gpu', choices=['cpu', 'gpu', 'cudnn'], required=True)
-    upscaler_options.add_argument('-d', '--driver', help='Waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter'])
-    upscaler_options.add_argument('-y', '--model_dir', help='directory containing model JSON files', action='store')
-    upscaler_options.add_argument('-t', '--threads', help='Number of threads to use for upscaling', action='store', type=int, default=1)
-    upscaler_options.add_argument('-c', '--config', help='Video2X config file location', action='store', default=f'{os.path.dirname(os.path.abspath(sys.argv[0]))}\\video2x.json')
-    upscaler_options.add_argument('-b', '--batch', help='Enable batch mode (select all default values to questions)', action='store_true')
+    upscaler_options.add_argument('-m', '--method', help='upscaling method', action='store', default='gpu', choices=['cpu', 'gpu', 'cudnn'])
+    upscaler_options.add_argument('-d', '--driver', help='waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan'])
+    upscaler_options.add_argument('-y', '--model_dir', type=pathlib.Path, help='directory containing model JSON files', action='store')
+    upscaler_options.add_argument('-t', '--threads', help='number of threads to use for upscaling', action='store', type=int, default=1)
+    upscaler_options.add_argument('-c', '--config', type=pathlib.Path, help='video2x config file location', action='store', default=pathlib.Path(sys.argv[0]).parent.absolute() / 'video2x.json')
+    upscaler_options.add_argument('-b', '--batch', help='enable batch mode (select all default values to questions)', action='store_true')
 
     # scaling options
     scaling_options = parser.add_argument_group('Scaling Options')
-    scaling_options.add_argument('--width', help='Output video width', action='store', type=int)
-    scaling_options.add_argument('--height', help='Output video height', action='store', type=int)
-    scaling_options.add_argument('-r', '--ratio', help='Scaling ratio', action='store', type=float)
+    scaling_options.add_argument('--width', help='output video width', action='store', type=int)
+    scaling_options.add_argument('--height', help='output video height', action='store', type=int)
+    scaling_options.add_argument('-r', '--ratio', help='scaling ratio', action='store', type=float)
+
+    # extra options
+    extra_options = parser.add_argument_group('Extra Options')
+    extra_options.add_argument('-v', '--version', help='display version, lawful information and exit', action='store_true')
 
     # parse arguments
     return parser.parse_args()
 
 
 def print_logo():
-    print('__      __  _       _                  ___   __   __')
-    print('\\ \\    / / (_)     | |                |__ \\  \\ \\ / /')
-    print(' \\ \\  / /   _    __| |   ___    ___      ) |  \\ V /')
-    print('  \\ \\/ /   | |  / _` |  / _ \\  / _ \\    / /    > <')
-    print('   \\  /    | | | (_| | |  __/ | (_) |  / /_   / . \\')
-    print('    \\/     |_|  \\__,_|  \\___|  \\___/  |____| /_/ \\_\\')
+    """print video2x logo"""
+    logo = r'''
+    __      __  _       _                  ___   __   __
+    \ \    / / (_)     | |                |__ \  \ \ / /
+     \ \  / /   _    __| |   ___    ___      ) |  \ V /
+      \ \/ /   | |  / _` |  / _ \  / _ \    / /    > <
+       \  /    | | | (_| | |  __/ | (_) |  / /_   / . \
+        \/     |_|  \__,_|  \___|  \___/  |____| /_/ \_\
+    '''
+    print(logo)
     print('\n               Video2X Video Enlarger')
     spaces = ((44 - len(f'Version {VERSION}')) // 2) * ' '
     print(f'{Avalon.FM.BD}\n{spaces}    Version {VERSION}\n{Avalon.FM.RST}')
@@ -119,7 +139,8 @@ def check_memory():
     # check if Nvidia-smi is available
     # GPUtil requires nvidia-smi.exe to interact with GPU
     if args.method == 'gpu' or args.method == 'cudnn':
-        if not os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe'):
+        if not (shutil.which('nvidia-smi') or
+                pathlib.Path(r'C:\Program Files\NVIDIA Corporation\NVSMI\nvidia-smi.exe').is_file()):
             # Nvidia System Management Interface not available
             Avalon.warning('Nvidia-smi not available, skipping available memory check')
             Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using')
@@ -183,24 +204,28 @@ def absolutify_paths(config):
     Returns:
         dict -- configuration file dictionary
     """
-    current_directory = os.path.dirname(os.path.abspath(sys.argv[0]))
+    current_directory = pathlib.Path(sys.argv[0]).parent.absolute()
 
     # check waifu2x-caffe path
     if not re.match('^[a-z]:', config['waifu2x_caffe']['waifu2x_caffe_path'], re.IGNORECASE):
-        config['waifu2x_caffe']['waifu2x_caffe_path'] = f'{current_directory}\\{config["waifu2x_caffe"]["waifu2x_caffe_path"]}'
+        config['waifu2x_caffe']['waifu2x_caffe_path'] = current_directory / config['waifu2x_caffe']['waifu2x_caffe_path']
 
     # check waifu2x-converter-cpp path
     if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE):
-        config['waifu2x_converter']['waifu2x_converter_path'] = f'{current_directory}\\{config["waifu2x_converter"]["waifu2x_converter_path"]}'
+        config['waifu2x_converter']['waifu2x_converter_path'] = current_directory / config['waifu2x_converter']['waifu2x_converter_path']
+
+    # check waifu2x_ncnn_vulkan path
+    if not re.match('^[a-z]:', config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'], re.IGNORECASE):
+        config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = current_directory / config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path']
 
     # check ffmpeg path
     if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE):
-        config['ffmpeg']['ffmpeg_path'] = f'{current_directory}\\{config["ffmpeg"]["ffmpeg_path"]}'
+        config['ffmpeg']['ffmpeg_path'] = current_directory / config['ffmpeg']['ffmpeg_path']
 
     # check video2x cache path
     if config['video2x']['video2x_cache_directory']:
         if not re.match('^[a-z]:', config['video2x']['video2x_cache_directory'], re.IGNORECASE):
-            config['video2x']['video2x_cache_directory'] = f'{current_directory}\\{config["video2x"]["video2x_cache_directory"]}'
+            config['video2x']['video2x_cache_directory'] = current_directory / config['video2x']['video2x_cache_directory']
 
     return config
 
@@ -212,21 +237,40 @@ if __name__ != '__main__':
     Avalon.error('This file cannot be imported')
     raise ImportError(f'{__file__} cannot be imported')
 
+# print video2x logo
 print_logo()
 
 # process CLI arguments
 args = process_arguments()
 
+# display version and lawful informaition
+if args.version:
+    print(f'Video2X Version: {VERSION}')
+    print('Author: K4YT3X')
+    print('License: GNU GPL v3')
+    print('Github Page: https://github.com/k4yt3x/video2x')
+    print('Contact: k4yt3x@k4yt3x.com')
+    exit(0)
+
 # arguments sanity check
-if args.driver == 'waifu2x_converter' and args.width and args.height:
-    Avalon.error('Waifu2x Converter CPP accepts only scaling ratio')
-    exit(1)
+if not args.input:
+    Avalon.error('You must specify input video file/directory path')
+    raise ArgumentError('input video path not specified')
+if not args.output:
+    Avalon.error('You must specify output video file/directory path')
+    raise ArgumentError('output video path not specified')
+if (args.driver == 'waifu2x_converter' or args.driver == 'waifu2x_ncnn_vulkan') and args.width and args.height:
+    Avalon.error('Waifu2x Converter CPP/NCNN accepts only scaling ratio')
+    raise ArgumentError('waifu2x-converter supports only scaling ratio')
+if args.driver == 'waifu2x_ncnn_vulkan' and (args.ratio > 2 or not args.ratio.is_integer()):
+    Avalon.error('Scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
+    raise ArgumentError('scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
 if (args.width or args.height) and args.ratio:
     Avalon.error('You can only specify either scaling ratio or output width and height')
-    exit(1)
+    raise ArgumentError('both scaling ration and width/height specified')
 if (args.width and not args.height) or (not args.width and args.height):
     Avalon.error('You must specify both width and height')
-    exit(1)
+    raise ArgumentError('only one of width or height is specified')
 
 # check available memory
 check_memory()
@@ -238,46 +282,59 @@ config = absolutify_paths(config)
 # load waifu2x configuration
 if args.driver == 'waifu2x_caffe':
     waifu2x_settings = config['waifu2x_caffe']
-    if not os.path.isfile(waifu2x_settings['waifu2x_caffe_path']):
+    if not pathlib.Path(waifu2x_settings['waifu2x_caffe_path']).is_file():
         Avalon.error('Specified waifu2x-caffe directory doesn\'t exist')
         Avalon.error('Please check the configuration file settings')
         raise FileNotFoundError(waifu2x_settings['waifu2x_caffe_path'])
 elif args.driver == 'waifu2x_converter':
     waifu2x_settings = config['waifu2x_converter']
-    if not os.path.isdir(waifu2x_settings['waifu2x_converter_path']):
-        Avalon.error('Specified waifu2x-conver-cpp directory doesn\'t exist')
+    if not pathlib.Path(waifu2x_settings['waifu2x_converter_path']).is_dir():
+        Avalon.error('Specified waifu2x-converter-cpp directory doesn\'t exist')
         Avalon.error('Please check the configuration file settings')
         raise FileNotFoundError(waifu2x_settings['waifu2x_converter_path'])
+elif args.driver == 'waifu2x_ncnn_vulkan':
+    waifu2x_settings = config['waifu2x_ncnn_vulkan']
+    if not pathlib.Path(waifu2x_settings['waifu2x_ncnn_vulkan_path']).is_file():
+        Avalon.error('Specified waifu2x_ncnn_vulkan directory doesn\'t exist')
+        Avalon.error('Please check the configuration file settings')
+        raise FileNotFoundError(waifu2x_settings['waifu2x_ncnn_vulkan_path'])
 
-# check if waifu2x path is valid
-
-
-# read FFMPEG configuration
+# read FFmpeg configuration
 ffmpeg_settings = config['ffmpeg']
 
 # load video2x settings
-video2x_cache_directory = config['video2x']['video2x_cache_directory']
 image_format = config['video2x']['image_format'].lower()
 preserve_frames = config['video2x']['preserve_frames']
 
-# create temp directories if they don't exist
-if not video2x_cache_directory:
-    video2x_cache_directory = f'{tempfile.gettempdir()}\\video2x'
+# load cache directory
+if isinstance(config['video2x']['video2x_cache_directory'], str):
+    video2x_cache_directory = pathlib.Path(config['video2x']['video2x_cache_directory'])
+else:
+    video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'
 
-if video2x_cache_directory and not os.path.isdir(video2x_cache_directory):
-    if not os.path.isfile(video2x_cache_directory) and not os.path.islink(video2x_cache_directory):
-        Avalon.warning(f'Specified cache directory {video2x_cache_directory} does not exist')
-        if Avalon.ask('Create directory?', default=True, batch=args.batch):
-            if os.mkdir(video2x_cache_directory) is None:
-                Avalon.info(f'{video2x_cache_directory} created')
-            else:
-                Avalon.error(f'Unable to create {video2x_cache_directory}')
-                Avalon.error('Aborting...')
-                exit(1)
+if video2x_cache_directory.exists() and not video2x_cache_directory.is_dir():
+    Avalon.error('Specified cache directory is a file/link')
+    raise FileExistsError('Specified cache directory is a file/link')
+
+elif not video2x_cache_directory.exists():
+    # if destination file is a file or a symbolic link
+    Avalon.warning(f'Specified cache directory {video2x_cache_directory} does not exist')
+
+    # try creating the cache directory
+    if Avalon.ask('Create directory?', default=True, batch=args.batch):
+        try:
+            video2x_cache_directory.mkdir(parents=True, exist_ok=True)
+            Avalon.info(f'{video2x_cache_directory} created')
+
+        # there can be a number of exceptions here
+        # PermissionError, FileExistsError, etc.
+        # therefore, we put a catch-them-all here
+        except Exception as e:
+            Avalon.error(f'Unable to create {video2x_cache_directory}')
+            Avalon.error('Aborting...')
+            raise e
     else:
-        Avalon.error('Specified cache directory is a file/link')
-        Avalon.error('Unable to continue, exiting...')
-        exit(1)
+        raise FileNotFoundError('Could not create cache directory')
 
 
 # start execution
@@ -286,16 +343,17 @@ try:
     begin_time = time.time()
 
     # if input specified is a single file
-    if os.path.isfile(args.input):
-        """ Upscale single video file """
+    if args.input.is_file():
+
+        # upscale single video file
         Avalon.info(f'Upscaling single video file: {args.input}')
 
         # check for input output format mismatch
-        if os.path.isdir(args.output):
+        if args.output.is_dir():
             Avalon.error('Input and output path type mismatch')
             Avalon.error('Input is single file but output is directory')
             raise Exception('input output path type mismatch')
-        if not re.search('.*\..*$', args.output):
+        if not re.search('.*\..*$', str(args.output)):
             Avalon.error('No suffix found in output file path')
             Avalon.error('Suffix must be specified for FFmpeg')
             raise Exception('No suffix specified')
@@ -319,12 +377,12 @@ try:
         upscaler.cleanup_temp_directories()
 
     # if input specified is a directory
-    elif os.path.isdir(args.input):
-        """ Upscale videos in a directory """
+    elif args.input.is_dir():
+        # upscale videos in a directory
         Avalon.info(f'Upscaling videos in directory: {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 = f'{args.output}\\{input_video}'
-            upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings)
+        for input_video in [f for f in args.input.iterdir() if f.is_file()]:
+            output_video = args.output / input_video
+            upscaler = Upscaler(input_video=args.input / input_video, output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings)
 
             # set optional options
             upscaler.waifu2x_driver = args.driver
@@ -350,7 +408,7 @@ except Exception:
     Avalon.error('An exception has occurred')
     traceback.print_exc()
 finally:
-    # remove Video2X Cache directory
+    # remove Video2X cache directory
     try:
         if not preserve_frames:
             shutil.rmtree(video2x_cache_directory)

bin/video2x_gui.py

@@ -0,0 +1,374 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Name: Video2x GUI
+Author: K4YT3X
+Date Created: July 27, 2019
+Last Modified: July 27, 2019
+
+Description: GUI for Video2X
+"""
+
+# local imports
+from exceptions import *
+from upscaler import Upscaler
+
+# built-in imports
+from tkinter import *
+from tkinter import messagebox
+from tkinter import ttk
+from tkinter.filedialog import *
+import json
+import pathlib
+import tempfile
+import threading
+import time
+import tkinter as tk
+
+VERSION = '1.0.0'
+
+
+# global static variables
+
+AVAILABLE_METHODS = {'GPU': 'gpu',
+    'CUDNN': 'cudnn',
+    'CPU': 'cpu'
+}
+
+AVAILABLE_DRIVERS = {
+    'Waifu2X Caffe': 'waifu2x_caffe',
+    'Waifu2X Converter CPP': 'waifu2x_converter',
+    'Waifu2x NCNN Vulkan': 'waifu2x_ncnn_vulkan'
+}
+
+IMAGE_FORMATS = {'PNG', 'JPG'}
+
+
+class Video2xGui():
+
+    def __init__(self):
+
+        self.running = False
+
+        # create main window
+        self.main_window = Tk()
+        self.main_window.title('Video2X GUI')
+        self.main_frame = Frame()
+        self.main_frame.pack(fill=BOTH, expand=True)
+
+        # file frame
+        self.file_frame = Frame(self.main_frame)
+        self.file_frame.pack(fill=X, padx=5, pady=5, expand=True)
+
+        # input file
+        self.input_file = StringVar()
+        label_text = StringVar()
+        label_text.set('Input File')
+        Label(self.file_frame, textvariable=label_text, relief=RIDGE, width=10).grid(row=0, column=0, padx=5, pady=5, sticky=W)
+        Entry(self.file_frame, textvariable=self.input_file, width=60).grid(row=0, column=1, padx=5, pady=5, sticky=W)
+        Button(self.file_frame, text='Select', command=self._select_input).grid(row=0, column=2, padx=5, pady=5, sticky=W)
+
+        # output file
+        self.output_file = StringVar()
+        label_text = StringVar()
+        label_text.set('Output File')
+        Label(self.file_frame, textvariable=label_text, relief=RIDGE, width=10).grid(row=1, column=0, padx=5, pady=5, sticky=W)
+        Entry(self.file_frame, textvariable=self.output_file, width=60).grid(row=1, column=1, padx=5, pady=5, sticky=W)
+        Button(self.file_frame, text='Select', command=self._select_output).grid(row=1, column=2, padx=5, pady=5, sticky=W)
+
+        # options
+        self.options_frame = Frame()
+        # self.options_left.pack(fill=X, padx=5, pady=5, expand=True)
+        self.options_frame.pack(fill=X, padx=5, pady=5, expand=True)
+
+        self.options_left = Frame(self.options_frame)
+        # self.options_left.pack(fill=X, padx=5, pady=5, expand=True)
+        self.options_left.grid(row=0, column=0, padx=5, pady=5, sticky=N)
+
+        # width
+        self.width = IntVar()
+        # self.width.set(1920)
+        Label(self.options_left, text='Width', relief=RIDGE, width=15).grid(row=0, column=0, padx=2, pady=3)
+        width_field = Entry(self.options_left, textvariable=self.width)
+        width_field.grid(row=0, column=1, padx=2, pady=3, sticky=W)
+
+        # height
+        self.height = IntVar()
+        # self.height.set(1080)
+        Label(self.options_left, text='Height', relief=RIDGE, width=15).grid(row=1, column=0, padx=2, pady=3)
+        height_field = Entry(self.options_left, textvariable=self.height)
+        height_field.grid(row=1, column=1, padx=2, pady=3, sticky=W)
+
+        # scale ratio
+        self.scale_ratio = DoubleVar()
+        # self.scale_ratio.set(2.0)
+        Label(self.options_left, text='Scale Ratio', relief=RIDGE, width=15).grid(row=2, column=0, padx=2, pady=3)
+        scale_ratio_field = Entry(self.options_left, textvariable=self.scale_ratio)
+        scale_ratio_field.grid(row=2, column=1, padx=2, pady=3, sticky=W)
+
+        # image format
+        self.image_format = StringVar(self.options_left)
+        self.image_format.set('PNG')
+        Label(self.options_left, text='Image Format', relief=RIDGE, width=15).grid(row=3, column=0, padx=2, pady=3)
+        image_format_menu = OptionMenu(self.options_left, self.image_format, *IMAGE_FORMATS)
+        image_format_menu.grid(row=3, column=1, padx=2, pady=3, sticky=W)
+
+        # options
+        self.options_right = Frame(self.options_frame)
+        # self.options_left.pack(fill=X, padx=5, pady=5, expand=True)
+        self.options_right.grid(row=0, column=1, padx=5, pady=5, sticky=N)
+
+        # threads
+        self.threads = IntVar()
+        self.threads.set(1)
+        Label(self.options_right, text='Threads', relief=RIDGE, width=15).grid(row=0, column=0, padx=2, pady=3)
+        threads_field = Entry(self.options_right, textvariable=self.threads)
+        threads_field.grid(row=0, column=1, padx=2, pady=3, sticky=W)
+
+        # method
+        self.method = StringVar(self.options_left)
+        self.method.set('GPU')
+        Label(self.options_right, text='Method', relief=RIDGE, width=15).grid(row=1, column=0, padx=2, pady=3)
+        method_menu = OptionMenu(self.options_right, self.method, *AVAILABLE_METHODS)
+        method_menu.grid(row=1, column=1, padx=2, pady=3, sticky=W)
+
+        # driver
+        self.driver = StringVar(self.options_left)
+        self.driver.set('Waifu2X Caffe')
+        Label(self.options_right, text='Driver', relief=RIDGE, width=15).grid(row=2, column=0, padx=2, pady=3)
+        driver_menu = OptionMenu(self.options_right, self.driver, *AVAILABLE_DRIVERS)
+        driver_menu.grid(row=2, column=1, padx=2, pady=3, sticky=W)
+
+        # preserve frames
+        self.preserve_frames = BooleanVar(self.options_left)
+        self.preserve_frames.set(True)
+        Label(self.options_right, text='Preserve Frames', relief=RIDGE, width=15).grid(row=3, column=0, padx=2, pady=3)
+        preserve_frames_menu = OptionMenu(self.options_right, self.preserve_frames, *{True, False})
+        preserve_frames_menu.grid(row=3, column=1, padx=2, pady=3, sticky=W)
+
+        # progress bar
+        self.progress_bar_frame = Frame()
+        self.progress_bar_frame.pack(fill=X, padx=5, pady=5, expand=True)
+
+        self.progress_bar = ttk.Progressbar(self.progress_bar_frame, orient='horizontal', length=100, mode='determinate')
+        self.progress_bar.pack(fill=X)
+
+        # start button frame
+        self.start_frame = Frame()
+        self.start_frame.pack(fill=X, padx=5, pady=5, expand=True)
+
+        # start button
+        self.start_button_text = StringVar()
+        self.start_button_text.set('Start')
+        Button(self.start_frame, textvariable=self.start_button_text, command=self._launch_upscaling, width=20).pack(side=RIGHT)
+
+        self.main_frame.mainloop()
+    
+    def _launch_upscaling(self):
+
+        # prevent launching multiple instances
+        if self.running:
+            messagebox.showerror('Error', 'Video2X is already running')
+            return
+
+        # arguments sanity check
+        if self.input_file.get() == '':
+            messagebox.showerror('Error', 'You must specify input video file/directory path')
+            return
+        if self.output_file.get() == '':
+            messagebox.showerror('Error', 'You must specify output video file/directory path')
+            return
+        if (self.driver.get() == 'waifu2x_converter' or self.driver.get() == 'waifu2x_ncnn_vulkan') and self.width.get() and self.height.get():
+            messagebox.showerror('Error', 'Waifu2x Converter CPP/NCNN accepts only scaling ratio')
+            return
+        if self.driver.get() == 'waifu2x_ncnn_vulkan' and (self.scale_ratio.get() > 2 or not self.scale_ratio.get().is_integer()):
+            messagebox.showerror('Error', 'Scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
+            return
+        if (self.width.get() or self.height.get()) and self.scale_ratio.get():
+            messagebox.showerror('Error', 'You can only specify either scaling ratio or output width and height')
+            return
+        if (self.width.get() and not self.height.get()) or (not self.width.get() and self.height.get()):
+            messagebox.showerror('Error', 'You must specify both width and height')
+            return
+        if (not self.width.get() or not self.height.get()) and not self.scale_ratio.get():
+            messagebox.showerror('Error', 'You must specify either output dimensions or scaling ratio')
+            return
+
+        upscale = threading.Thread(target=self._upscale)
+        upscale.start()
+        self.running = True
+        self.start_button_text.set('Running')
+
+    def _upscale(self):
+
+        # read configuration file
+        config = read_config('video2x.json')
+        config = absolutify_paths(config)
+
+        input_file = pathlib.Path(self.input_file.get())
+        output_file = pathlib.Path(self.output_file.get())
+        driver = AVAILABLE_DRIVERS[self.driver.get()]
+
+        if driver == 'waifu2x_caffe':
+            waifu2x_settings = config['waifu2x_caffe']
+            if not pathlib.Path(waifu2x_settings['waifu2x_caffe_path']).is_file():
+                messagebox.showerror('Error', 'Specified waifu2x-caffe directory doesn\'t exist\nPlease check the configuration file settings')
+                raise FileNotFoundError(waifu2x_settings['waifu2x_caffe_path'])
+        elif driver == 'waifu2x_converter':
+            waifu2x_settings = config['waifu2x_converter']
+            if not pathlib.Path(waifu2x_settings['waifu2x_converter_path']).is_dir():
+                messagebox.showerror('Error', 'Specified waifu2x-converter-cpp directory doesn\'t exist\nPlease check the configuration file settings')
+        elif driver == 'waifu2x_ncnn_vulkan':
+            waifu2x_settings = config['waifu2x_ncnn_vulkan']
+            if not pathlib.Path(waifu2x_settings['waifu2x_ncnn_vulkan_path']).is_file():
+                messagebox.showerror('Error', 'Specified waifu2x_ncnn_vulkan directory doesn\'t exist\nPlease check the configuration file settings')
+                raise FileNotFoundError(waifu2x_settings['waifu2x_ncnn_vulkan_path'])
+
+        # read FFmpeg configuration
+        ffmpeg_settings = config['ffmpeg']
+        # load video2x settings
+        image_format = config['video2x']['image_format'].lower()
+        preserve_frames = config['video2x']['preserve_frames']
+
+        # load cache directory
+        if isinstance(config['video2x']['video2x_cache_directory'], str):
+            video2x_cache_directory = pathlib.Path(config['video2x']['video2x_cache_directory'])
+        else:
+            video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'
+
+        if video2x_cache_directory.exists() and not video2x_cache_directory.is_dir():
+            messagebox.showerror('Error', 'Specified cache directory is a file/link')
+            raise FileExistsError('Specified cache directory is a file/link')
+
+        elif not video2x_cache_directory.exists():
+            # try creating the cache directory
+            if messagebox.askyesno('Question', f'Specified cache directory {video2x_cache_directory} does not exist\nCreate directory?'):
+                try:
+                    video2x_cache_directory.mkdir(parents=True, exist_ok=True)
+
+                # there can be a number of exceptions here
+                # PermissionError, FileExistsError, etc.
+                # therefore, we put a catch-them-all here
+                except Exception as e:
+                    messagebox.showerror('Error', f'Unable to create {video2x_cache_directory}\nAborting...')
+                    raise e
+            else:
+                raise FileNotFoundError('Could not create cache directory')
+
+        # load more settings from gui
+        width = self.width.get()
+        height = self.height.get()
+        scale_ratio = self.scale_ratio.get()
+        image_format = self.image_format.get()
+        threads = self.threads.get()
+        method = AVAILABLE_METHODS[self.method.get()]
+        preserve_frames = self.preserve_frames.get()
+
+        self.upscaler = Upscaler(input_video=input_file, output_video=output_file, method=method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings)
+
+        # set optional options
+        self.upscaler.waifu2x_driver = driver
+        self.upscaler.scale_width = width
+        self.upscaler.scale_height = height
+        self.upscaler.scale_ratio = scale_ratio
+        self.upscaler.model_dir = None
+        self.upscaler.threads = threads
+        self.upscaler.video2x_cache_directory = video2x_cache_directory
+        self.upscaler.image_format = image_format
+        self.upscaler.preserve_frames = preserve_frames
+
+        # run upscaler
+        self.upscaler.create_temp_directories()
+
+        # start progress bar
+        progress_bar = threading.Thread(target=self._progress_bar)
+        progress_bar.start()
+
+        # start upscaling
+        self.upscaler.run()
+        self.upscaler.cleanup_temp_directories()
+
+        # show message when upscaling completes
+        messagebox.showinfo('Info', 'Upscaling Completed')
+        self.progress_bar['value'] = 100
+        self.running = False
+        self.start_button_text.set('Start')
+
+
+    def _progress_bar(self):
+        """ This method prints a progress bar
+
+        This method prints a progress bar by keeping track
+        of the amount of frames in the input directory
+        and the output directory. This is originally
+        suggested by @ArmandBernard.
+        """
+        # initialize variables early
+        self.upscaler.progress_bar_exit_signal = False
+        self.upscaler.total_frames_upscaled = 0
+        self.upscaler.total_frames = 1
+
+        # initialize progress bar values
+        self.progress_bar['value'] = 0
+
+        while not self.upscaler.progress_bar_exit_signal:
+            self.progress_bar['value'] = int(100 * self.upscaler.total_frames_upscaled / self.upscaler.total_frames)
+            time.sleep(1)
+
+    def _select_input(self):
+        self.input_file.set(askopenfilename(title='Select Input File'))
+        self.output_file.set(f'{self.input_file.get()}_output.mp4')
+
+    def _select_output(self):
+        self.output_file.set(asksaveasfilename(title='Select Output File'))
+
+
+def read_config(config_file):
+    """ Reads configuration file
+
+    Returns a dictionary read by JSON.
+    """
+    with open(config_file, 'r') as raw_config:
+        config = json.load(raw_config)
+        return config
+
+
+def absolutify_paths(config):
+    """ Check to see if paths to binaries are absolute
+
+    This function checks if paths to binary files are absolute.
+    If not, then absolutify the path.
+
+    Arguments:
+        config {dict} -- configuration file dictionary
+
+    Returns:
+        dict -- configuration file dictionary
+    """
+    current_directory = pathlib.Path(sys.argv[0]).parent.absolute()
+
+    # check waifu2x-caffe path
+    if not re.match('^[a-z]:', config['waifu2x_caffe']['waifu2x_caffe_path'], re.IGNORECASE):
+        config['waifu2x_caffe']['waifu2x_caffe_path'] = current_directory / config['waifu2x_caffe']['waifu2x_caffe_path']
+
+    # check waifu2x-converter-cpp path
+    if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE):
+        config['waifu2x_converter']['waifu2x_converter_path'] = current_directory / config['waifu2x_converter']['waifu2x_converter_path']
+
+    # check waifu2x_ncnn_vulkan path
+    if not re.match('^[a-z]:', config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'], re.IGNORECASE):
+        config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = current_directory / config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path']
+
+    # check ffmpeg path
+    if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE):
+        config['ffmpeg']['ffmpeg_path'] = current_directory / config['ffmpeg']['ffmpeg_path']
+
+    # check video2x cache path
+    if config['video2x']['video2x_cache_directory']:
+        if not re.match('^[a-z]:', config['video2x']['video2x_cache_directory'], re.IGNORECASE):
+            config['video2x']['video2x_cache_directory'] = current_directory / config['video2x']['video2x_cache_directory']
+
+    return config
+
+
+video2x_gui = Video2xGui()

bin/video2x_setup.py

@@ -3,8 +3,11 @@
 """
 Name: Video2X Setup Script
 Author: K4YT3X
+Author: BrianPetkovsek
 Date Created: November 28, 2018
-Last Modified: April 28, 2019
+Last Modified: July 30, 2019
+
+Dev: SAT3LL
 
 Licensed under the GNU General Public License Version 3 (GNU GPL v3),
     available at: https://www.gnu.org/licenses/gpl-3.0.txt
@@ -17,15 +20,22 @@ and generates a configuration for it.
 Installation Details:
 - ffmpeg: %LOCALAPPDATA%\\video2x\\ffmpeg
 - waifu2x-caffe: %LOCALAPPDATA%\\video2x\\waifu2x-caffe
-
+- waifu2x-cpp-converter: %LOCALAPPDATA%\\video2x\\waifu2x-converter-cpp
+- waifu2x_ncnn_vulkan: %LOCALAPPDATA%\\video2x\\waifu2x-ncnn-vulkan
 """
+
+# built-in imports
 import argparse
 import json
 import os
+import pathlib
+import re
+import shutil
 import subprocess
 import sys
 import tempfile
 import traceback
+import urllib
 import zipfile
 
 # Requests doesn't come with windows, therefore
@@ -33,7 +43,10 @@ import zipfile
 # later in the script.
 # import requests
 
-VERSION = '1.2.1'
+VERSION = '1.4.0'
+
+# global static variables
+LOCALAPPDATA = pathlib.Path(os.getenv('localappdata'))
 
 
 def process_arguments():
@@ -43,7 +56,7 @@ def process_arguments():
 
     # video options
     general_options = parser.add_argument_group('General Options')
-    general_options.add_argument('-d', '--driver', help='driver to download and configure', action='store', choices=['all', 'waifu2x_caffe', 'waifu2x_converter'], default='all')
+    general_options.add_argument('-d', '--driver', help='driver to download and configure', action='store', choices=['all', 'waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan'], default='all')
 
     # parse arguments
     return parser.parse_args()
@@ -72,10 +85,13 @@ class Video2xSetup:
         if self.driver == 'all':
             self._install_waifu2x_caffe()
             self._install_waifu2x_converter_cpp()
+            self._install_waifu2x_ncnn_vulkan()
         elif self.driver == 'waifu2x_caffe':
             self._install_waifu2x_caffe()
         elif self.driver == 'waifu2x_converter':
             self._install_waifu2x_converter_cpp()
+        elif self.driver == 'waifu2x_ncnn_vulkan':
+            self._install_waifu2x_ncnn_vulkan()
 
         print('\nGenerating Video2X configuration file')
         self._generate_config()
@@ -86,20 +102,22 @@ class Video2xSetup:
     def _install_python_requirements(self):
         """ Read requirements.txt and return its content
         """
-        with open('requirements.txt', 'r') as req:
-            for line in req:
-                package = line.split('==')[0]
-                pip_install(package)
+        pip_install('requirements.txt')
 
     def _cleanup(self):
         """ Cleanup all the temp files downloaded
         """
         for file in self.trash:
             try:
-                print(f'Deleting: {file}')
-                os.remove(file)
-            except FileNotFoundError:
-                pass
+                if file.is_dir():
+                    print(f'Deleting directory: {file}')
+                    shutil.rmtree(file)
+                else:
+                    print(f'Deleting file: {file}')
+                    file.unlink()
+            except Exception:
+                print(f'Error deleting: {file}')
+                traceback.print_exc()
 
     def _install_ffmpeg(self):
         """ Install FFMPEG
@@ -110,7 +128,7 @@ class Video2xSetup:
         self.trash.append(ffmpeg_zip)
 
         with zipfile.ZipFile(ffmpeg_zip) as zipf:
-            zipf.extractall(f'{os.getenv("localappdata")}\\video2x')
+            zipf.extractall(LOCALAPPDATA / 'video2x')
 
     def _install_waifu2x_caffe(self):
         """ Install waifu2x_caffe
@@ -119,7 +137,7 @@ class Video2xSetup:
         import requests
 
         # Get latest release of waifu2x-caffe via GitHub API
-        latest_release = json.loads(requests.get('https://api.github.com/repos/lltcggie/waifu2x-caffe/releases/latest').content)
+        latest_release = json.loads(requests.get('https://api.github.com/repos/lltcggie/waifu2x-caffe/releases/latest').content.decode('utf-8'))
 
         for a in latest_release['assets']:
             if 'waifu2x-caffe.zip' in a['browser_download_url']:
@@ -127,17 +145,16 @@ class Video2xSetup:
                 self.trash.append(waifu2x_caffe_zip)
 
         with zipfile.ZipFile(waifu2x_caffe_zip) as zipf:
-            zipf.extractall(f'{os.getenv("localappdata")}\\video2x')
+            zipf.extractall(LOCALAPPDATA / 'video2x')
 
     def _install_waifu2x_converter_cpp(self):
         """ Install waifu2x_caffe
         """
         print('\nInstalling waifu2x-converter-cpp')
-        import re
         import requests
 
         # Get latest release of waifu2x-caffe via GitHub API
-        latest_release = json.loads(requests.get('https://api.github.com/repos/DeadSix27/waifu2x-converter-cpp/releases/latest').content)
+        latest_release = json.loads(requests.get('https://api.github.com/repos/DeadSix27/waifu2x-converter-cpp/releases/latest').content.decode('utf-8'))
 
         for a in latest_release['assets']:
             if re.search(r'waifu2x-DeadSix27-win64_v[0-9]*\.zip', a['browser_download_url']):
@@ -145,7 +162,33 @@ class Video2xSetup:
                 self.trash.append(waifu2x_converter_cpp_zip)
 
         with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf:
-            zipf.extractall(f'{os.getenv("localappdata")}\\video2x\\waifu2x-converter-cpp')
+            zipf.extractall(LOCALAPPDATA / 'video2x' / 'waifu2x-converter-cpp')
+
+    def _install_waifu2x_ncnn_vulkan(self):
+        """ Install waifu2x-ncnn-vulkan
+        """
+        print('\nInstalling waifu2x-ncnn-vulkan')
+        import requests
+
+        # Get latest release of waifu2x-ncnn-vulkan via Github API
+        latest_release = json.loads(requests.get('https://api.github.com/repos/nihui/waifu2x-ncnn-vulkan/releases/latest').content.decode('utf-8'))
+
+        for a in latest_release['assets']:
+            if re.search(r'waifu2x-ncnn-vulkan-\d*\.zip', a['browser_download_url']):
+                waifu2x_ncnn_vulkan_zip = download(a['browser_download_url'], tempfile.gettempdir())
+                self.trash.append(waifu2x_ncnn_vulkan_zip)
+
+        # extract and rename
+        waifu2x_ncnn_vulkan_directory = LOCALAPPDATA / 'video2x' / 'waifu2x-ncnn-vulkan'
+        with zipfile.ZipFile(waifu2x_ncnn_vulkan_zip) as zipf:
+            zipf.extractall(LOCALAPPDATA / 'video2x')
+
+            # if directory already exists, remove it
+            if waifu2x_ncnn_vulkan_directory.exists():
+                shutil.rmtree(waifu2x_ncnn_vulkan_directory)
+
+            # rename the newly extracted directory
+            (LOCALAPPDATA / 'video2x' / zipf.namelist()[0]).rename(waifu2x_ncnn_vulkan_directory)
 
     def _generate_config(self):
         """ Generate video2x config
@@ -155,18 +198,19 @@ class Video2xSetup:
             template_dict = json.load(template)
             template.close()
 
-        local_app_data = os.getenv('localappdata')
-
         # configure only the specified drivers
         if self.driver == 'all':
-            template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = f'{local_app_data}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe'
-            template_dict['waifu2x_converter']['waifu2x_converter_path'] = f'{local_app_data}\\video2x\\waifu2x-converter-cpp'
+            template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = str(LOCALAPPDATA / 'video2x' / 'waifu2x-caffe' / 'waifu2x-caffe-cui.exe')
+            template_dict['waifu2x_converter']['waifu2x_converter_path'] = str(LOCALAPPDATA / 'video2x' / 'waifu2x-converter-cpp')
+            template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = str(LOCALAPPDATA / 'video2x' / 'waifu2x-ncnn-vulkan' / 'waifu2x-ncnn-vulkan.exe')
         elif self.driver == 'waifu2x_caffe':
-            template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = f'{local_app_data}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe'
+            template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = str(LOCALAPPDATA / 'video2x' / 'waifu2x-caffe' / 'waifu2x-caffe-cui.exe')
         elif self.driver == 'waifu2x_converter':
-            template_dict['waifu2x_converter']['waifu2x_converter_path'] = f'{local_app_data}\\video2x\\waifu2x-converter-cpp'
+            template_dict['waifu2x_converter']['waifu2x_converter_path'] = str(LOCALAPPDATA / 'video2x' / 'waifu2x-converter-cpp')
+        elif self.driver == 'waifu2x_ncnn_vulkan':
+            template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = str(LOCALAPPDATA / 'video2x' / 'waifu2x-ncnn-vulkan' / 'waifu2x-ncnn-vulkan.exe')
 
-        template_dict['ffmpeg']['ffmpeg_path'] = f'{local_app_data}\\video2x\\ffmpeg-latest-win64-static\\bin'
+        template_dict['ffmpeg']['ffmpeg_path'] = str(LOCALAPPDATA / 'video2x' / 'ffmpeg-latest-win64-static' / 'bin')
         template_dict['video2x']['video2x_cache_directory'] = None
         template_dict['video2x']['preserve_frames'] = False
 
@@ -182,14 +226,43 @@ def download(url, save_path, chunk_size=4096):
     from tqdm import tqdm
     import requests
 
-    output_file = f'{save_path}\\{url.split("/")[-1]}'
+    save_path = pathlib.Path(save_path)
+
+    # create target folder if it doesn't exist
+    save_path.mkdir(parents=True, exist_ok=True)
+
+    # create requests stream for steaming file
+    stream = requests.get(url, stream=True, allow_redirects=True)
+
+    # get file name
+    file_name = None
+    if 'content-disposition' in stream.headers:
+        disposition = stream.headers['content-disposition']
+        try:
+            file_name = re.findall("filename=(.+)", disposition)[0].strip('"')
+        except IndexError:
+            pass
+
+    if file_name is None:
+        # output_file = f'{save_path}\\{stream.url.split("/")[-1]}'
+        output_file = save_path / stream.url.split('/')[-1]
+    else:
+        output_file = save_path / file_name
+
+    # decode url encoding
+    output_file = pathlib.Path(urllib.parse.unquote(str(output_file)))
+
+    # get total size for progress bar if provided in headers
+    total_size = 0
+    if 'content-length' in stream.headers:
+        total_size = int(stream.headers['content-length'])
+
+    # print download information summary
     print(f'Downloading: {url}')
+    print(f'Total size: {total_size}')
     print(f'Chunk size: {chunk_size}')
     print(f'Saving to: {output_file}')
 
-    stream = requests.get(url, stream=True)
-    total_size = int(stream.headers['content-length'])
-
     # Write content into file
     with open(output_file, 'wb') as output:
         with tqdm(total=total_size, ascii=True) as progress_bar:
@@ -198,19 +271,20 @@ def download(url, save_path, chunk_size=4096):
                     output.write(chunk)
                     progress_bar.update(len(chunk))
 
+    # return the full path of saved file
     return output_file
 
 
-def pip_install(package):
+def pip_install(file):
     """ Install python package via python pip module
 
     pip.main() is not available after pip 9.0.1, thus
     pip module is not used in this case.
     """
-    return subprocess.run(['python', '-m', 'pip', 'install', '-U', package]).returncode
+    return subprocess.run([sys.executable, '-m', 'pip', 'install', '-U', '-r', file]).returncode
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     try:
         args = process_arguments()
         print('Video2X Setup Script')
@@ -227,7 +301,16 @@ if __name__ == "__main__":
         setup.run()
         print('\nScript finished successfully')
     except Exception:
+
         traceback.print_exc()
         print('An error has occurred')
         print('Video2X Automatic Setup has failed')
+
+        # in case of a failure, try cleaning up temp files
+        try:
+            setup._cleanup()
+        except Exception:
+            traceback.print_exc()
+            print('An error occurred while trying to cleanup files')
+
         exit(1)

bin/waifu2x_caffe.py

@@ -4,15 +4,19 @@
 Name: Waifu2x Caffe Driver
 Author: K4YT3X
 Date Created: Feb 24, 2018
-Last Modified: April 28, 2019
+Last Modified: July 27, 2019
 
 Description: This class is a high-level wrapper
 for waifu2x-caffe.
 """
-from avalon_framework import Avalon
+
+# built-in imports
 import subprocess
 import threading
 
+# third-party imports
+from avalon_framework import Avalon
+
 
 class Waifu2xCaffe:
     """This class communicates with waifu2x cui engine
@@ -23,10 +27,11 @@ class Waifu2xCaffe:
     the upscale function.
     """
 
-    def __init__(self, waifu2x_settings, process, model_dir):
+    def __init__(self, waifu2x_settings, process, model_dir, bit_depth):
         self.waifu2x_settings = waifu2x_settings
         self.waifu2x_settings['process'] = process
         self.waifu2x_settings['model_dir'] = model_dir
+        self.waifu2x_settings['output_depth'] = bit_depth
 
         # arguments passed through command line overwrites config file values
         self.process = process

bin/waifu2x_converter.py

@@ -4,15 +4,20 @@
 Name: Waifu2x Converter CPP Driver
 Author: K4YT3X
 Date Created: February 8, 2019
-Last Modified: April 28, 2019
+Last Modified: July 27, 2019
 
 Description: This class is a high-level wrapper
 for waifu2x-converter-cpp.
 """
-from avalon_framework import Avalon
+
+# built-in imports
+import pathlib
 import subprocess
 import threading
 
+# third-party imports
+from avalon_framework import Avalon
+
 
 class Waifu2xConverter:
     """This class communicates with waifu2x cui engine
@@ -59,7 +64,7 @@ class Waifu2xConverter:
             # models_rgb must be specified manually for waifu2x-converter-cpp
             # if it's not specified in the arguments, create automatically
             if self.waifu2x_settings['model-dir'] is None:
-                self.waifu2x_settings['model-dir'] = f'{self.waifu2x_settings["waifu2x_converter_path"]}\\models_rgb'
+                self.waifu2x_settings['model-dir'] = pathlib.Path(self.waifu2x_settings['waifu2x_converter_path']) / 'models_rgb'
 
             # print thread start message
             self.print_lock.acquire()
@@ -75,7 +80,7 @@ class Waifu2xConverter:
 
                 # the key doesn't need to be passed in this case
                 if key == 'waifu2x_converter_path':
-                    execute.append(f'{str(value)}\\waifu2x-converter-cpp.exe')
+                    execute.append(pathlib.Path(str(value)) / 'waifu2x-converter-cpp.exe')
 
                 # null or None means that leave this option out (keep default)
                 elif value is None or value is False:

bin/waifu2x_ncnn_vulkan.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Name: Waifu2x NCNN Vulkan Driver
+Author: SAT3LL
+Date Created: June 26, 2019
+Last Modified: July 27, 2019
+
+Dev: K4YT3X
+
+Description: This class is a high-level wrapper
+for waifu2x_ncnn_vulkan.
+"""
+
+# built-in imports
+import os
+import subprocess
+import threading
+
+# third-party imports
+from avalon_framework import Avalon
+
+
+class Waifu2xNcnnVulkan:
+    """This class communicates with waifu2x ncnn vulkan engine
+
+    An object will be created for this class, containing information
+    about the binary address and the processing method. When being called
+    by the main program, other detailed information will be passed to
+    the upscale function.
+    """
+
+    def __init__(self, waifu2x_settings):
+        self.waifu2x_settings = waifu2x_settings
+
+        # arguments passed through command line overwrites config file values
+
+        # waifu2x_ncnn_vulkan can't find its own model directory if its not in the current dir
+        #   so change to it
+        os.chdir(os.path.join(self.waifu2x_settings['waifu2x_ncnn_vulkan_path'], '..'))
+
+        self.print_lock = threading.Lock()
+
+    def upscale(self, input_directory, output_directory, scale_ratio, upscaler_exceptions):
+        """This is the core function for WAIFU2X class
+
+        Arguments:
+            input_directory {string} -- source directory path
+            output_directory {string} -- output directory path
+            ratio {int} -- output video ratio
+        """
+
+        try:
+            # overwrite config file settings
+            self.waifu2x_settings['input'] = input_directory
+            self.waifu2x_settings['output'] = output_directory
+
+            # print thread start message
+            self.print_lock.acquire()
+            Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} started')
+            self.print_lock.release()
+
+            # waifu2x_ncnn_vulkan does not have long-opts, we'll have a dictionary that maps "our" config long-opt
+            # names to their short opts
+            waifu2x_ncnn_vulkan_opt_flag = {
+                'input': '-i',
+                'output': '-o',
+                'noise-level': '-n',
+                'scale-ratio': '-s',
+                'tile-size': '-t',
+                'model-path': '-m',
+                'gpu': '-g',
+                'load-proc-save_threads': '-j',
+                'verbose': '-v'
+            }
+
+            execute = [self.waifu2x_settings['waifu2x_ncnn_vulkan_path']]
+            for key in self.waifu2x_settings.keys():
+                value = self.waifu2x_settings[key]
+                if key == 'waifu2x_ncnn_vulkan_path':
+                    continue
+                elif key == 'input':
+                    execute.append(waifu2x_ncnn_vulkan_opt_flag[key])
+                    execute.append(input_directory)
+                elif key == 'output':
+                    execute.append(waifu2x_ncnn_vulkan_opt_flag[key])
+                    execute.append(output_directory)
+                elif key == 'scale-ratio':
+                    execute.append(waifu2x_ncnn_vulkan_opt_flag[key])
+                    # waifu2x_ncnn_vulkan does not accept an arbitrary scale ratio, max is 2
+                    if scale_ratio == 1:
+                        execute.append('1')
+                    else:
+                        execute.append('2')
+                # allow upper if cases to take precedence
+                elif value is None or value is False:
+                    continue
+                else:
+                    execute.append(waifu2x_ncnn_vulkan_opt_flag[key])
+                    execute.append(str(value))
+
+            Avalon.debug_info(f'Executing: {execute}')
+            subprocess.run(execute, check=True, stderr=subprocess.DEVNULL)
+
+            # print thread exiting message
+            self.print_lock.acquire()
+            Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} exiting')
+            self.print_lock.release()
+
+            return 0
+        except Exception as e:
+            upscaler_exceptions.append(e)