Merge pull request #112 from sat3ll/waifu2x-ncnn-vk

Add support for waifu2x-ncnn-vulkan engine

bin/upscaler.py

@@ -6,7 +6,9 @@
 Name: Video2X Upscaler
 Author: K4YT3X
 Date Created: December 10, 2018
-Last Modified: June 15, 2019
+Last Modified: June 26, 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
@@ -22,6 +24,7 @@ from image_cleaner import ImageCleaner
 from tqdm import tqdm
 from waifu2x_caffe import Waifu2xCaffe
 from waifu2x_converter import Waifu2xConverter
+from waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan
 import copy
 import os
 import re
@@ -147,7 +150,8 @@ class Upscaler:
         self.upscaler_exceptions = []
 
         # initialize waifu2x driver
-        if self.waifu2x_driver != 'waifu2x_caffe' and self.waifu2x_driver != 'waifu2x_converter':
+        drivers = ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan']
+        if self.waifu2x_driver not in drivers:
             raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
 
         # it's easier to do multi-threading with waifu2x_converter
@@ -166,84 +170,109 @@ class Upscaler:
             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
+            # list all images in the extracted frames
-        upscaler_threads = []
+            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))]
 
-        # list all images in the extracted frames
+            # if we have less images than threads,
-        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))]
+            # create only the threads necessary
+            if len(frames) < self.threads:
+                self.threads = len(frames)
 
-        # if we have less images than threads,
+            # create a directory for each thread and append directory
-        # create only the threads necessary
+            # name into a list
-        if len(frames) < self.threads:
-            self.threads = len(frames)
 
-        # create a directory for each thread and append directory
+            thread_pool = []
-        # name into a list
+            thread_directories = []
+            for thread_id in range(self.threads):
+                thread_directory = os.path.join(self.extracted_frames, str(thread_id))
+                thread_directories.append(thread_directory)
 
-        thread_pool = []
+                # delete old directories and create new directories
-        thread_directories = []
+                if os.path.isdir(thread_directory):
-        for thread_id in range(self.threads):
+                    shutil.rmtree(thread_directory)
-            thread_directory = os.path.join(self.extracted_frames, str(thread_id))
+                os.mkdir(thread_directory)
-            thread_directories.append(thread_directory)
 
-            # delete old directories and create new directories
+                # append directory path into list
-            if os.path.isdir(thread_directory):
+                thread_pool.append((thread_directory, thread_id))
-                shutil.rmtree(thread_directory)
-            os.mkdir(thread_directory)
 
-            # append directory path into list
+            # evenly distribute images into each directory
-            thread_pool.append((thread_directory, thread_id))
+            # 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]
 
-        # evenly distribute images into each directory
+            # create threads and start them
-        # until there is none left in the directory
+            for thread_info in thread_pool:
-        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
+                # create a separate w2 instance for each thread
-        for thread_info in thread_pool:
+                if self.waifu2x_driver == 'waifu2x_caffe':
+                    w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir)
+                    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))
 
-            # create a separate w2 instance for each thread
+                # if the driver being used is waifu2x_ncnn_vulkan 
-            w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir)
+                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))
 
-            # create thread
+                # create thread
-            if self.scale_ratio:
+                thread.name = thread_info[1]
-                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]
 
-            # add threads into the pool
+                # add threads into the pool
-            upscaler_threads.append(thread)
+                upscaler_threads.append(thread)
 
-        # start progress bar in a different thread
+            # start progress bar in a different thread
-        progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,))
+            progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,))
-        progress_bar.start()
+            progress_bar.start()
 
-        # create the clearer and start it
+            # create the clearer and start it
-        Avalon.debug_info('Starting upscaled image cleaner')
+            Avalon.debug_info('Starting upscaled image cleaner')
-        image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
+            image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
-        image_cleaner.start()
+            image_cleaner.start()
 
-        # start all threads
+            # start all threads
-        for thread in upscaler_threads:
+            for thread in upscaler_threads:
-            thread.start()
+                thread.start()
 
-        # wait for threads to finish
+            # wait for threads to finish
-        for thread in upscaler_threads:
+            for thread in upscaler_threads:
-            thread.join()
+                thread.join()
 
-        # upscaling done, kill the clearer
+            # upscaling done, kill the clearer
-        Avalon.debug_info('Killing upscaled image cleaner')
+            Avalon.debug_info('Killing upscaled image cleaner')
-        image_cleaner.stop()
+            image_cleaner.stop()
 
-        self.progress_bar_exit_signal = True
+            self.progress_bar_exit_signal = True
+
+            if len(self.upscaler_exceptions) != 0:
+                raise(self.upscaler_exceptions[0])
 
-        if len(self.upscaler_exceptions) != 0:
-            raise(self.upscaler_exceptions[0])
 
     def run(self):
         """Main controller for Video2X

bin/video2x.json

@@ -39,6 +39,14 @@
     "output": null,
     "input": null
   },
+  "waifu2x_ncnn_vulkan": {
+    "waifu2x_ncnn_vulkan_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\waifu2x-ncnn-vulkan\\waifu2x.exe",
+    "input": null,
+    "output": null,
+    "noise-level": 2,
+    "scale-ratio": null,
+    "block-size": 400
+  },
   "ffmpeg": {
     "ffmpeg_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\ffmpeg-latest-win64-static\\bin",
     "video_to_frames": {

bin/video2x.py

@@ -15,7 +15,7 @@ __      __  _       _                  ___   __   __
 Name: Video2X Controller
 Author: K4YT3X
 Date Created: Feb 24, 2018
-Last Modified: June 15, 2019
+Last Modified: June 26, 2019
 
 Dev: BrianPetkovsek
 Dev: SAT3LL
@@ -82,7 +82,7 @@ def process_arguments():
     # 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'])
-    upscaler_options.add_argument('-d', '--driver', help='waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter'])
+    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', 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=os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), 'video2x.json'))
@@ -203,6 +203,10 @@ def absolutify_paths(config):
     if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE):
         config['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(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'] = os.path.join(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'] = os.path.join(current_directory, config['ffmpeg']['ffmpeg_path'])
@@ -244,8 +248,11 @@ if not args.input:
 if not args.output:
     Avalon.error('You must specify output video file/directory path')
     exit(1)
-if args.driver == 'waifu2x_converter' and args.width and args.height:
+if (args.driver == 'waifu2x_converter' or args.driver == 'waifu2x_ncnn_vulkan') and args.width and args.height:
-    Avalon.error('Waifu2x Converter CPP accepts only scaling ratio')
+    Avalon.error('Waifu2x Converter CPP/NCNN accepts only scaling ratio')
+    exit(1)
+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')
     exit(1)
 if (args.width or args.height) and args.ratio:
     Avalon.error('You can only specify either scaling ratio or output width and height')
@@ -271,9 +278,15 @@ if args.driver == 'waifu2x_caffe':
 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')
+        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 os.path.isfile(waifu2x_settings['waifu2x_ncnn_vulkan_path']):
+        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'])
 
 # read FFmpeg configuration
 ffmpeg_settings = config['ffmpeg']

bin/video2x_setup.py

@@ -7,7 +7,9 @@ Name: Video2X Setup Script
 Author: K4YT3X
 Author: BrianPetkovsek
 Date Created: November 28, 2018
-Last Modified: June 15, 2019
+Last Modified: June 26, 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
@@ -20,11 +22,14 @@ 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
 
 """
 import argparse
 import json
 import os
+import shutil
 import subprocess
 import sys
 import tempfile
@@ -36,7 +41,7 @@ import zipfile
 # later in the script.
 # import requests
 
-VERSION = '1.2.1'
+VERSION = '1.3.0'
 
 
 def process_arguments():
@@ -46,7 +51,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()
@@ -75,10 +80,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()
@@ -96,8 +104,12 @@ class Video2xSetup:
         """
         for file in self.trash:
             try:
-                print('Deleting: {}'.format(file))
+                if os.path.isfile(file):
-                os.remove(file)
+                    print('Deleting: {}'.format(file))
+                    os.remove(file)
+                else:
+                    print('Deleting: {}'.format(file))
+                    shutil.rmtree(file)
             except FileNotFoundError:
                 pass
 
@@ -147,6 +159,29 @@ class Video2xSetup:
         with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf:
             zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x', 'waifu2x-converter-cpp'))
 
+    def _install_waifu2x_ncnn_vulkan(self):
+        """ Install waifu2x-ncnn-vulkan
+        """
+        print('\nInstalling waifu2x-ncnn-vulkan')
+        import re
+        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 then move (to remove the top level directory)
+        with zipfile.ZipFile(waifu2x_ncnn_vulkan_zip) as zipf:
+            extraction_path = os.path.join(tempfile.gettempdir(), 'waifu2x-ncnn-vulkan-ext')
+            zipf.extractall(extraction_path)
+            shutil.move(os.path.join(extraction_path, os.listdir(extraction_path)[0]), os.path.join(os.getenv('localappdata'), 'video2x', 'waifu2x-ncnn-vulkan'))
+            self.trash.append(extraction_path)
+
+
     def _generate_config(self):
         """ Generate video2x config
         """
@@ -161,10 +196,15 @@ class Video2xSetup:
         if self.driver == 'all':
             template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-caffe', 'waifu2x-caffe-cui.exe')
             template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-converter-cpp')
+            # TODO: after version 20190611 executable changes to waifu2x-ncnn-vulkan so rename this when it breaks
+            template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-ncnn-vulkan', 'waifu2x.exe')
         elif self.driver == 'waifu2x_caffe':
             template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-caffe', 'waifu2x-caffe-cui.exe')
         elif self.driver == 'waifu2x_converter':
             template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-converter-cpp')
+        elif self.driver == 'waifu2x_ncnn_vulkan':
+            # TODO: after version 20190611 executable changes to waifu2x-ncnn-vulkan so rename this when it breaks
+            template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-ncnn-vulkan', 'waifu2x.exe')
 
         template_dict['ffmpeg']['ffmpeg_path'] = os.path.join(local_app_data, 'video2x', 'ffmpeg-latest-win64-static', 'bin')
         template_dict['video2x']['video2x_cache_directory'] = None

bin/waifu2x_ncnn_vulkan.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+# -*- coding: future_fstrings -*-
+
+
+"""
+Name: Waifu2x NCNN Vulkan Driver
+Author: SAT3LL
+Date Created: June 26, 2019
+Last Modified: June 26, 2019
+
+Dev: K4YT3X
+
+Description: This class is a high-level wrapper
+for waifu2x_ncnn_vulkan.
+"""
+from avalon_framework import Avalon
+import os
+import subprocess
+import threading
+
+
+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_path'] = input_directory
+            self.waifu2x_settings['output_path'] = 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 accepts arguments in a positional manner
+            # See: https://github.com/nihui/waifu2x_ncnn_vulkan#usage
+            # waifu2x_ncnn_vulkan.exe [input image] [output png] [noise=-1/0/1/2/3] [scale=1/2] [blocksize=400]
+            #     noise = noise level, large value means strong denoise effect, -1=no effect
+            #     scale = scale level, 1=no scale, 2=upscale 2x
+            #     blocksize = tile size, use smaller value to reduce GPU memory usage, default is 400
+
+            # waifu2x_ncnn_vulkan does not accept an arbitrary scale ratio, max is 2
+            if scale_ratio == 1:
+                for raw_frame in os.listdir(input_directory):
+                    command = [
+                        os.path.join(input_directory, raw_frame),
+                        os.path.join(output_directory, raw_frame),
+                        str(self.waifu2x_settings['noise-level']),
+                        '1',
+                        str(self.waifu2x_settings['block-size'])
+                    ]
+                    execute = [self.waifu2x_settings['waifu2x_ncnn_vulkan_path']]
+                    execute.extend(command)
+
+                    Avalon.debug_info(f'Executing: {execute}')
+                    subprocess.run(execute, check=True, stderr=subprocess.DEVNULL)
+            else:
+                for raw_frame in os.listdir(input_directory):
+                    command = [
+                        os.path.join(input_directory, raw_frame),
+                        os.path.join(output_directory, raw_frame),
+                        str(self.waifu2x_settings['noise-level']),
+                        '2',
+                        str(self.waifu2x_settings['block-size'])
+                    ]
+                    execute = [self.waifu2x_settings['waifu2x_ncnn_vulkan_path']]
+                    execute.extend(command)
+
+                    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)