3.2.0 added SRMD NCNN Vulkan support

src/srmd_ncnn_vulkan.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Name: SRMD NCNN Vulkan Driver
+Creator: K4YT3X
+Date Created: April 26, 2020
+Last Modified: April 26, 2020
+
+Description: This class is a high-level wrapper
+for srmd_ncnn_vulkan.
+"""
+
+# built-in imports
+import os
+import shlex
+import subprocess
+import threading
+
+# third-party imports
+from avalon_framework import Avalon
+
+
+class SrmdNcnnVulkan:
+    """This class communicates with SRMD 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, driver_settings):
+        self.driver_settings = driver_settings
+
+        # arguments passed through command line overwrites config file values
+
+        # srmd_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.driver_settings['path'], '..'))
+
+        self.print_lock = threading.Lock()
+
+    def upscale(self, input_directory, output_directory, scale_ratio):
+        """This is the core function for SRMD ncnn Vulkan class
+
+        Arguments:
+            input_directory {string} -- source directory path
+            output_directory {string} -- output directory path
+            ratio {int} -- output video ratio
+        """
+
+        # overwrite config file settings
+        self.driver_settings['i'] = input_directory
+        self.driver_settings['o'] = output_directory
+        self.driver_settings['s'] = int(scale_ratio)
+
+        # list to be executed
+        # initialize the list with the binary path as the first element
+        execute = [str(self.driver_settings['path'])]
+
+        for key in self.driver_settings.keys():
+
+            value = self.driver_settings[key]
+
+            # is executable key or null or None means that leave this option out (keep default)
+            if key == 'path' or value is None or value is False:
+                continue
+            else:
+                if len(key) == 1:
+                    execute.append(f'-{key}')
+                else:
+                    execute.append(f'--{key}')
+                execute.append(str(value))
+
+        # return the Popen object of the new process created
+        self.print_lock.acquire()
+        Avalon.debug_info(f'[upscaler] Subprocess {os.getpid()} executing: {shlex.join(execute)}')
+        self.print_lock.release()
+        return subprocess.Popen(execute)

src/upscaler.py

@@ -4,7 +4,7 @@
 Name: Video2X Upscaler
 Author: K4YT3X
 Date Created: December 10, 2018
-Last Modified: April 4, 2020
+Last Modified: April 26, 2020
 
 Description: This file contains the Upscaler class. Each
 instance of the Upscaler class is an upscaler on an image or
@@ -16,6 +16,7 @@ from anime4k import Anime4k
 from exceptions import *
 from ffmpeg import Ffmpeg
 from image_cleaner import ImageCleaner
+from srmd_ncnn_vulkan import SrmdNcnnVulkan
 from waifu2x_caffe import Waifu2xCaffe
 from waifu2x_converter import Waifu2xConverter
 from waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan
@@ -37,7 +38,7 @@ import traceback
 from avalon_framework import Avalon
 from tqdm import tqdm
 
-AVAILABLE_DRIVERS = ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan', 'anime4k']
+AVAILABLE_DRIVERS = ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan', 'anime4k', 'srmd_ncnn_vulkan']
 
 
 class Upscaler:
@@ -242,6 +243,13 @@ class Upscaler:
                                                      self.scale_ratio,
                                                      self.processes)
 
+            # if the driver being used is srmd_ncnn_vulkan
+            elif self.waifu2x_driver == 'srmd_ncnn_vulkan':
+                driver = SrmdNcnnVulkan(copy.deepcopy(self.driver_settings))
+                upscaler_processes.append(driver.upscale(process_directory,
+                                                         self.upscaled_frames,
+                                                         self.scale_ratio))
+
         # start progress bar in a different thread
         progress_bar = threading.Thread(target=self._progress_bar, args=(process_directories,))
         progress_bar.start()

src/video2x.py

@@ -72,7 +72,7 @@ import GPUtil
 import psutil
 
 
-VERSION = '3.1.0'
+VERSION = '3.2.0'
 
 LEGAL_INFO = f'''Video2X Version: {VERSION}
 Author: K4YT3X
@@ -269,6 +269,9 @@ if (args.driver in ['waifu2x_converter', 'waifu2x_ncnn_vulkan', 'anime4k']) and
 if args.driver == 'waifu2x_ncnn_vulkan' and args.ratio is not None 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.driver == 'srmd_ncnn_vulkan' and args.ratio is not None and (args.ratio not in [2, 3, 4]):
+    Avalon.error('Scaling ratio must be one of 2, 3 or 4 for srmd_ncnn_vulkan')
+    raise ArgumentError('scaling ratio must be one of 2, 3 or 4 for srmd_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')
     raise ArgumentError('both scaling ration and width/height specified')
@@ -277,7 +280,7 @@ if (args.width and not args.height) or (not args.width and args.height):
     raise ArgumentError('only one of width or height is specified')
 
 # check available memory if driver is waifu2x-based
-if args.driver in ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan']:
+if args.driver in ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan', 'srmd_ncnn_vulkan']:
     check_memory()
 
 # anime4k runs significantly faster with more processes