added RealSR NCNN Vulkan support

This commit is contained in:
k4yt3x 2020-05-26 06:35:53 -04:00
parent 387a233daf
commit 4c5ca73e38
6 changed files with 284 additions and 8 deletions

View File

@ -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:

View File

@ -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])

View File

@ -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':

View File

@ -1,6 +1,6 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE QtCreatorProject>
<!-- Written by QtCreator 4.12.0, 2020-05-23T08:51:47. -->
<!-- Written by QtCreator 4.12.0, 2020-05-26T06:35:20. -->
<qtcreator>
<data>
<variable>EnvironmentId</variable>

View File

@ -7,7 +7,7 @@
<x>0</x>
<y>0</y>
<width>673</width>
<height>802</height>
<height>844</height>
</rect>
</property>
<property name="acceptDrops">
@ -342,6 +342,11 @@
<string>SRMD NCNN Vulkan</string>
</property>
</item>
<item>
<property name="text">
<string>RealSR NCNN Vulkan</string>
</property>
</item>
<item>
<property name="text">
<string>Anime4KCPP</string>
@ -1384,6 +1389,131 @@
</item>
</layout>
</widget>
<widget class="QWidget" name="realsrNcnnVulkanTab">
<attribute name="title">
<string>RealSR NCNN Vulkan</string>
</attribute>
<layout class="QVBoxLayout" name="verticalLayout_34">
<item>
<layout class="QHBoxLayout" name="realsrNcnnVulkanPathHorizontalLayout">
<item>
<widget class="QLineEdit" name="realsrNcnnVulkanPathLineEdit"/>
</item>
<item>
<widget class="QPushButton" name="realsrNcnnVulkanPathSelectButton">
<property name="text">
<string>Select Binary Path</string>
</property>
</widget>
</item>
</layout>
</item>
<item>
<layout class="QHBoxLayout" name="realsrNcnnVulkanTileSizeHorizontalLayout">
<item>
<widget class="QLabel" name="realsrNcnnVulkanTileSizeLabel">
<property name="text">
<string>Tile Size</string>
</property>
</widget>
</item>
<item>
<widget class="QSpinBox" name="realsrNcnnVulkanTileSizeSpinBox">
<property name="minimum">
<number>0</number>
</property>
<property name="maximum">
<number>99999</number>
</property>
<property name="value">
<number>0</number>
</property>
</widget>
</item>
</layout>
</item>
<item>
<layout class="QHBoxLayout" name="realsrNcnnVulkanModelHorizontalLayout">
<item>
<widget class="QLabel" name="realsrNcnnVulkanModelLabel">
<property name="text">
<string>Model</string>
</property>
</widget>
</item>
<item>
<widget class="QComboBox" name="realsrNcnnVulkanModelComboBox">
<property name="toolTip">
<string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;If the input file contains very little noise, DF2K will produce sharper outputs.&lt;/p&gt;&lt;p&gt;However, if the input is noisy, DF2K might produce artifacts. DF2K_JPEG will then be preferred.&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
</property>
<item>
<property name="text">
<string>models-DF2K_JPEG</string>
</property>
</item>
<item>
<property name="text">
<string>models-DF2K</string>
</property>
</item>
</widget>
</item>
</layout>
</item>
<item>
<layout class="QHBoxLayout" name="realsrNcnnVulkanGpuIdHorizontalLayout">
<item>
<widget class="QLabel" name="realsrNcnnVulkanGpuIdLabel">
<property name="text">
<string>GPU ID</string>
</property>
</widget>
</item>
<item>
<widget class="QSpinBox" name="realsrNcnnVulkanGpuIdSpinBox"/>
</item>
</layout>
</item>
<item>
<layout class="QHBoxLayout" name="realsrNcnnVulkanJobsHorizontalLayout">
<item>
<widget class="QLabel" name="realsrNcnnVulkanJobsLabel">
<property name="text">
<string>Jobs</string>
</property>
</widget>
</item>
<item>
<widget class="QLineEdit" name="realsrNcnnVulkanJobsLineEdit">
<property name="text">
<string>1:2:2</string>
</property>
</widget>
</item>
</layout>
</item>
<item>
<widget class="QCheckBox" name="realsrNcnnVulkanTtaCheckBox">
<property name="text">
<string>TTA</string>
</property>
</widget>
</item>
<item>
<spacer name="realsrNcnnVulkanVerticalSpacer">
<property name="orientation">
<enum>Qt::Vertical</enum>
</property>
<property name="sizeHint" stdset="0">
<size>
<width>20</width>
<height>331</height>
</size>
</property>
</spacer>
</item>
</layout>
</widget>
<widget class="QWidget" name="anime4kCppTab">
<attribute name="title">
<string>Anime4K CPP</string>

View File

@ -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)