From 7224c5399714e244c18f8a028769341cc0261484 Mon Sep 17 00:00:00 2001 From: k4yt3x Date: Sun, 26 Apr 2020 20:01:37 -0400 Subject: [PATCH] 3.2.0 added SRMD NCNN Vulkan support --- src/srmd_ncnn_vulkan.py | 79 +++++++++++++++++++++++++++++++++++++++++ src/upscaler.py | 12 +++++-- src/video2x.py | 7 ++-- 3 files changed, 94 insertions(+), 4 deletions(-) create mode 100644 src/srmd_ncnn_vulkan.py diff --git a/src/srmd_ncnn_vulkan.py b/src/srmd_ncnn_vulkan.py new file mode 100644 index 0000000..04db35c --- /dev/null +++ b/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) diff --git a/src/upscaler.py b/src/upscaler.py index 176092e..70a4e64 100755 --- a/src/upscaler.py +++ b/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() diff --git a/src/video2x.py b/src/video2x.py index ec703a2..b8dd0f4 100755 --- a/src/video2x.py +++ b/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