From 4c5ca73e38046985c4a97c9fb34569edcf2a56e0 Mon Sep 17 00:00:00 2001 From: k4yt3x Date: Tue, 26 May 2020 06:35:53 -0400 Subject: [PATCH] added RealSR NCNN Vulkan support --- src/upscaler.py | 7 +- src/video2x.yaml | 13 ++- src/video2x_gui.py | 36 +++++++- src/video2x_gui.pyproject.user | 2 +- src/video2x_gui.ui | 132 ++++++++++++++++++++++++++++- src/wrappers/realsr_ncnn_vulkan.py | 102 ++++++++++++++++++++++ 6 files changed, 284 insertions(+), 8 deletions(-) create mode 100644 src/wrappers/realsr_ncnn_vulkan.py diff --git a/src/upscaler.py b/src/upscaler.py index bff223c..f5b25de 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: May 22, 2020 +Last Modified: May 26, 2020 Description: This file contains the Upscaler class. Each instance of the Upscaler class is an upscaler on an image or @@ -50,7 +50,7 @@ language.install() _ = language.gettext # version information -UPSCALER_VERSION = '4.1.1' +UPSCALER_VERSION = '4.2.0' # these names are consistent for # - driver selection in command line @@ -60,6 +60,7 @@ AVAILABLE_DRIVERS = ['waifu2x_caffe', 'waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan', 'srmd_ncnn_vulkan', + 'realsr_ncnn_vulkan', 'anime4kcpp'] @@ -281,7 +282,7 @@ class Upscaler: process_directory.mkdir(parents=True, exist_ok=True) # waifu2x-converter-cpp will perform multi-threading within its own process - if self.driver in ['waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan', 'srmd_ncnn_vulkan']: + if self.driver in ['waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan', 'srmd_ncnn_vulkan', 'realsr_ncnn_vulkan']: process_directories = [self.extracted_frames] else: diff --git a/src/video2x.yaml b/src/video2x.yaml index 43358d0..a0db3c7 100644 --- a/src/video2x.yaml +++ b/src/video2x.yaml @@ -1,7 +1,7 @@ # Name: Video2X Configuration File # Creator: K4YT3X # Date Created: October 23, 2018 -# Last Modified: May 23, 2020 +# Last Modified: May 26, 2020 # Values here are the default values. Change the value here to # save the default value permanently. # Items commented out are parameters irrelevant to this context @@ -76,6 +76,17 @@ srmd_ncnn_vulkan: g: 0 # gpu device to use (default=0) j: '1:2:2' # thread count for load/proc/save (default=1:2:2) x: false # enable tta mode +realsr_ncnn_vulkan: + path: '%LOCALAPPDATA%\video2x\realsr-ncnn-vulkan\realsr-ncnn-vulkan' + v: null # verbose output + #i: null # input-path: input image path (jpg/png) or directory + #o: null # output-path: output image path (png) or directory + s: 4 # upscale ratio (4, default=4) + t: 0 # tile size (>=32/0=auto, default=0) + m: null # realsr model path (default=models-DF2K) + g: 0 # gpu device to use (default=0) + j: '1:2:2' # thread count for load/proc/save (default=1:2:2) + x: false # enable tta mode anime4kcpp: path: '%LOCALAPPDATA%\video2x\anime4kcpp\CLI\Anime4KCPP_CLI\Anime4KCPP_CLI' #input: null # File for loading (string [=./pic/p1.png]) diff --git a/src/video2x_gui.py b/src/video2x_gui.py index 9991fea..94ba316 100755 --- a/src/video2x_gui.py +++ b/src/video2x_gui.py @@ -4,7 +4,7 @@ Creator: Video2X GUI Author: K4YT3X Date Created: May 5, 2020 -Last Modified: May 23, 2020 +Last Modified: May 26, 2020 """ # local imports @@ -32,7 +32,7 @@ from PyQt5.QtGui import * from PyQt5.QtWidgets import * import magic -GUI_VERSION = '2.4.0' +GUI_VERSION = '2.5.0' LEGAL_INFO = f'''Video2X GUI Version: {GUI_VERSION}\\ Upscaler Version: {UPSCALER_VERSION}\\ @@ -46,6 +46,7 @@ AVAILABLE_DRIVERS = { 'Waifu2X Converter CPP': 'waifu2x_converter_cpp', 'Waifu2X NCNN Vulkan': 'waifu2x_ncnn_vulkan', 'SRMD NCNN Vulkan': 'srmd_ncnn_vulkan', + 'RealSR NCNN Vulkan': 'realsr_ncnn_vulkan', 'Anime4KCPP': 'anime4kcpp' } @@ -346,6 +347,17 @@ class Video2XMainWindow(QMainWindow): self.srmd_ncnn_vulkan_jobs_line_edit = self.findChild(QLineEdit, 'srmdNcnnVulkanJobsLineEdit') self.srmd_ncnn_vulkan_tta_check_box = self.findChild(QCheckBox, 'srmdNcnnVulkanTtaCheckBox') + # realsr-ncnn-vulkan + self.realsr_ncnn_vulkan_path_line_edit = self.findChild(QLineEdit, 'realsrNcnnVulkanPathLineEdit') + self.enable_line_edit_file_drop(self.realsr_ncnn_vulkan_path_line_edit) + self.realsr_ncnn_vulkan_path_select_button = self.findChild(QPushButton, 'realsrNcnnVulkanPathSelectButton') + self.realsr_ncnn_vulkan_path_select_button.clicked.connect(lambda: self.select_driver_binary_path(self.realsr_ncnn_vulkan_path_line_edit)) + self.realsr_ncnn_vulkan_tile_size_spin_box = self.findChild(QSpinBox, 'realsrNcnnVulkanTileSizeSpinBox') + self.realsr_ncnn_vulkan_model_combo_box = self.findChild(QComboBox, 'realsrNcnnVulkanModelComboBox') + self.realsr_ncnn_vulkan_gpu_id_spin_box = self.findChild(QSpinBox, 'realsrNcnnVulkanGpuIdSpinBox') + self.realsr_ncnn_vulkan_jobs_line_edit = self.findChild(QLineEdit, 'realsrNcnnVulkanJobsLineEdit') + self.realsr_ncnn_vulkan_tta_check_box = self.findChild(QCheckBox, 'realsrNcnnVulkanTtaCheckBox') + # anime4k self.anime4kcpp_path_line_edit = self.findChild(QLineEdit, 'anime4kCppPathLineEdit') self.enable_line_edit_file_drop(self.anime4kcpp_path_line_edit) @@ -497,6 +509,14 @@ class Video2XMainWindow(QMainWindow): self.srmd_ncnn_vulkan_jobs_line_edit.setText(settings['j']) self.srmd_ncnn_vulkan_tta_check_box.setChecked(settings['x']) + # realsr-ncnn-vulkan + settings = self.config['realsr_ncnn_vulkan'] + self.realsr_ncnn_vulkan_path_line_edit.setText(str(pathlib.Path(os.path.expandvars(settings['path'])).absolute())) + self.realsr_ncnn_vulkan_tile_size_spin_box.setValue(settings['t']) + self.realsr_ncnn_vulkan_gpu_id_spin_box.setValue(settings['g']) + self.realsr_ncnn_vulkan_jobs_line_edit.setText(settings['j']) + self.realsr_ncnn_vulkan_tta_check_box.setChecked(settings['x']) + # anime4k settings = self.config['anime4kcpp'] self.anime4kcpp_path_line_edit.setText(str(pathlib.Path(os.path.expandvars(settings['path'])).absolute())) @@ -598,6 +618,14 @@ class Video2XMainWindow(QMainWindow): self.config['srmd_ncnn_vulkan']['j'] = self.srmd_ncnn_vulkan_jobs_line_edit.text() self.config['srmd_ncnn_vulkan']['x'] = self.srmd_ncnn_vulkan_tta_check_box.isChecked() + # realsr-ncnn-vulkan + self.config['realsr_ncnn_vulkan']['path'] = os.path.expandvars(self.realsr_ncnn_vulkan_path_line_edit.text()) + self.config['realsr_ncnn_vulkan']['t'] = self.realsr_ncnn_vulkan_tile_size_spin_box.value() + self.config['realsr_ncnn_vulkan']['m'] = str((pathlib.Path(self.config['realsr_ncnn_vulkan']['path']).parent / self.realsr_ncnn_vulkan_model_combo_box.currentText()).absolute()) + self.config['realsr_ncnn_vulkan']['g'] = self.realsr_ncnn_vulkan_gpu_id_spin_box.value() + self.config['realsr_ncnn_vulkan']['j'] = self.realsr_ncnn_vulkan_jobs_line_edit.text() + self.config['realsr_ncnn_vulkan']['x'] = self.realsr_ncnn_vulkan_tta_check_box.isChecked() + # anime4k self.config['anime4kcpp']['path'] = os.path.expandvars(self.anime4kcpp_path_line_edit.text()) self.config['anime4kcpp']['passes'] = self.anime4kcpp_passes_spin_box.value() @@ -788,6 +816,10 @@ class Video2XMainWindow(QMainWindow): self.scale_ratio_double_spin_box.setMinimum(2.0) self.scale_ratio_double_spin_box.setMaximum(4.0) self.scale_ratio_double_spin_box.setValue(2.0) + elif current_driver == 'realsr_ncnn_vulkan': + self.scale_ratio_double_spin_box.setMinimum(4.0) + self.scale_ratio_double_spin_box.setMaximum(4.0) + self.scale_ratio_double_spin_box.setValue(4.0) # update preferred processes/threads count if current_driver == 'anime4kcpp': diff --git a/src/video2x_gui.pyproject.user b/src/video2x_gui.pyproject.user index 1ce5915..0c97afa 100644 --- a/src/video2x_gui.pyproject.user +++ b/src/video2x_gui.pyproject.user @@ -1,6 +1,6 @@ - + EnvironmentId diff --git a/src/video2x_gui.ui b/src/video2x_gui.ui index c5fdd3f..272450c 100644 --- a/src/video2x_gui.ui +++ b/src/video2x_gui.ui @@ -7,7 +7,7 @@ 0 0 673 - 802 + 844 @@ -342,6 +342,11 @@ SRMD NCNN Vulkan + + + RealSR NCNN Vulkan + + Anime4KCPP @@ -1384,6 +1389,131 @@ + + + RealSR NCNN Vulkan + + + + + + + + + + + Select Binary Path + + + + + + + + + + + Tile Size + + + + + + + 0 + + + 99999 + + + 0 + + + + + + + + + + + Model + + + + + + + <html><head/><body><p>If the input file contains very little noise, DF2K will produce sharper outputs.</p><p>However, if the input is noisy, DF2K might produce artifacts. DF2K_JPEG will then be preferred.</p></body></html> + + + + models-DF2K_JPEG + + + + + models-DF2K + + + + + + + + + + + + GPU ID + + + + + + + + + + + + + + Jobs + + + + + + + 1:2:2 + + + + + + + + + TTA + + + + + + + Qt::Vertical + + + + 20 + 331 + + + + + + Anime4K CPP diff --git a/src/wrappers/realsr_ncnn_vulkan.py b/src/wrappers/realsr_ncnn_vulkan.py new file mode 100644 index 0000000..f49ffa6 --- /dev/null +++ b/src/wrappers/realsr_ncnn_vulkan.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Name: RealSR NCNN Vulkan Driver +Creator: K4YT3X +Date Created: May 26, 2020 +Last Modified: May 26, 2020 + +Description: This class is a high-level wrapper +for realsr_ncnn_vulkan. +""" + +# built-in imports +import argparse +import os +import pathlib +import platform +import shlex +import subprocess +import threading + +# third-party imports +from avalon_framework import Avalon + + +class WrapperMain: + """This class communicates with RealSR 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 + self.print_lock = threading.Lock() + + @staticmethod + def parse_arguments(arguments): + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, add_help=False) + parser.error = lambda message: (_ for _ in ()).throw(AttributeError(message)) + parser.add_argument('--help', action='help', help='show this help message and exit') + parser.add_argument('-v', action='store_true', help='verbose output') + parser.add_argument('-i', type=str, help=argparse.SUPPRESS) # help='input image path (jpg/png) or directory') + parser.add_argument('-o', type=str, help=argparse.SUPPRESS) # help='output image path (png) or directory') + parser.add_argument('-s', type=int, choices=[4], help='upscale ratio') + parser.add_argument('-t', type=int, help='tile size (>=32/0=auto)') + parser.add_argument('-m', type=str, help='realsr model path') + parser.add_argument('-g', type=int, help='gpu device to use') + parser.add_argument('-j', type=str, help='thread count for load/proc/save') + parser.add_argument('-x', action='store_true', help='enable tta mode') + return parser.parse_args(arguments) + + def load_configurations(self, upscaler): + self.driver_settings['s'] = int(upscaler.scale_ratio) + self.driver_settings['j'] = '{}:{}:{}'.format(upscaler.processes, upscaler.processes, upscaler.processes) + + def upscale(self, input_directory, output_directory): + """This is the core function for RealSR 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 + + # by default, realsr-ncnn-vulkan will look for the models under the current working directory + # change the working directory to its containing folder if model directory not specified + if self.driver_settings['m'] is None and platform.system() == 'Windows': + os.chdir(pathlib.Path(self.driver_settings['path']).parent) + + # list to be executed + # initialize the list with the binary path as the first element + execute = [self.driver_settings['path']] + + for key in self.driver_settings.keys(): + + value = self.driver_settings[key] + + # 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}') + + # true means key is an option + if value is not True: + 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)