OpenSourceJoker/video2x
3
0
Fork 0
mirror of https://github.com/k4yt3x/video2x.git synced 2024-12-29 16:09:10 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add support for waifu2x-ncnn-vulkan engine

Browse Source
This commit is contained in:
sat3ll 2019-06-26 00:25:27 +01:00
parent b40aee0cd7
commit 4e3565b3ac
5 changed files with 309 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

200
bin/upscaler.py
View File

@ -6,7 +6,9 @@
Name: Video2X Upscaler Name: Video2X Upscaler
Author: K4YT3X Author: K4YT3X
Date Created: December 10, 2018 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), Licensed under the GNU General Public License Version 3 (GNU GPL v3),
available at: https://www.gnu.org/licenses/gpl-3.0.txt 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 tqdm import tqdm
from waifu2x_caffe import Waifu2xCaffe from waifu2x_caffe import Waifu2xCaffe
from waifu2x_converter import Waifu2xConverter from waifu2x_converter import Waifu2xConverter
from waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan
import copy import copy
import os import os
import re import re
@ -147,7 +150,7 @@ class Upscaler:
self.upscaler_exceptions = [] self.upscaler_exceptions = []
# initialize waifu2x driver # initialize waifu2x driver
if self.waifu2x_driver != 'waifu2x_caffe' and self.waifu2x_driver != 'waifu2x_converter': if self.waifu2x_driver != 'waifu2x_caffe' and self.waifu2x_driver != 'waifu2x_converter' and self.waifu2x_driver != 'waifu2x-ncnn-vulkan':
raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}') raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
# it's easier to do multi-threading with waifu2x_converter # it's easier to do multi-threading with waifu2x_converter
@ -166,84 +169,159 @@ class Upscaler:
self.progress_bar_exit_signal = True self.progress_bar_exit_signal = True
progress_bar.join() progress_bar.join()
return return
elif self.waifu2x_driver == 'waifu2x_caffe':
# 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: w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir)
# create a separate w2 instance for each thread # create thread
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))
thread.name = thread_info[1]
# create thread # add threads into the pool
if self.scale_ratio: upscaler_threads.append(thread)
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 # start progress bar in a different thread
upscaler_threads.append(thread) progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,))
progress_bar.start()
# start progress bar in a different thread # create the clearer and start it
progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,)) Avalon.debug_info('Starting upscaled image cleaner')
progress_bar.start() image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
image_cleaner.start()
# create the clearer and start it # start all threads
Avalon.debug_info('Starting upscaled image cleaner') for thread in upscaler_threads:
image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads)) thread.start()
image_cleaner.start()
# start all threads # wait for threads to finish
for thread in upscaler_threads: for thread in upscaler_threads:
thread.start() thread.join()
# wait for threads to finish # upscaling done, kill the clearer
for thread in upscaler_threads: Avalon.debug_info('Killing upscaled image cleaner')
thread.join() image_cleaner.stop()
# upscaling done, kill the clearer self.progress_bar_exit_signal = True
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True if len(self.upscaler_exceptions) != 0:
raise(self.upscaler_exceptions[0])
elif self.waifu2x_driver == 'waifu2x-ncnn-vulkan':
# create a container for all upscaler threads
upscaler_threads = []
if len(self.upscaler_exceptions) != 0: # list all images in the extracted frames
raise(self.upscaler_exceptions[0]) 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))]
# if we have less images than threads,
# create only the threads necessary
if len(frames) < self.threads:
self.threads = len(frames)
# create a directory for each thread and append directory
# name into a list
thread_pool = []
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)
# delete old directories and create new directories
if os.path.isdir(thread_directory):
shutil.rmtree(thread_directory)
os.mkdir(thread_directory)
# append directory path into list
thread_pool.append((thread_directory, thread_id))
# evenly distribute images into each directory
# 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]
# create threads and start them
for thread_info in thread_pool:
# create a separate w2 instance for each thread
w2 = Waifu2xNcnnVulkan(copy.deepcopy(self.waifu2x_settings))
# create thread
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, self.upscaler_exceptions))
thread.name = thread_info[1]
# add threads into the pool
upscaler_threads.append(thread)
# start progress bar in a different thread
progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,))
progress_bar.start()
# create the clearer and start it
Avalon.debug_info('Starting upscaled image cleaner')
image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
image_cleaner.start()
# start all threads
for thread in upscaler_threads:
thread.start()
# wait for threads to finish
for thread in upscaler_threads:
thread.join()
# upscaling done, kill the clearer
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True
if len(self.upscaler_exceptions) != 0:
raise(self.upscaler_exceptions[0])
def run(self): def run(self):
"""Main controller for Video2X """Main controller for Video2X

8
bin/video2x.json
View File

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

23
bin/video2x.py
View File

@ -15,7 +15,7 @@ __ __ _ _ ___ __ __
Name: Video2X Controller Name: Video2X Controller
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 Date Created: Feb 24, 2018
Last Modified: June 15, 2019 Last Modified: June 26, 2019
Dev: BrianPetkovsek Dev: BrianPetkovsek
Dev: SAT3LL Dev: SAT3LL
@ -82,7 +82,7 @@ def process_arguments():
# upscaler options # upscaler options
upscaler_options = parser.add_argument_group('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('-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('-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('-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')) 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): 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']) 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 # check ffmpeg path
if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE): 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']) 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: if not args.output:
Avalon.error('You must specify output video file/directory path') Avalon.error('You must specify output video file/directory path')
exit(1) 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 / waifu2x-ncnn-vulkan accepts only scaling ratio')
exit(1)
if args.driver == 'waifu2x-ncnn-vulkan' and (args.ratio > 2 or not args.ratio.is_integer()):
Avalon.error('scale ratio must be 2 or 1 for waifu2x-ncnn-vulkan')
exit(1) exit(1)
if (args.width or args.height) and args.ratio: if (args.width or args.height) and args.ratio:
Avalon.error('You can only specify either scaling ratio or output width and height') 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': elif args.driver == 'waifu2x_converter':
waifu2x_settings = config['waifu2x_converter'] waifu2x_settings = config['waifu2x_converter']
if not os.path.isdir(waifu2x_settings['waifu2x_converter_path']): 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') Avalon.error('Please check the configuration file settings')
raise FileNotFoundError(waifu2x_settings['waifu2x_converter_path']) 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 # read FFmpeg configuration
ffmpeg_settings = config['ffmpeg'] ffmpeg_settings = config['ffmpeg']

48
bin/video2x_setup.py
View File

@ -7,7 +7,9 @@ Name: Video2X Setup Script
Author: K4YT3X Author: K4YT3X
Author: BrianPetkovsek Author: BrianPetkovsek
Date Created: November 28, 2018 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), Licensed under the GNU General Public License Version 3 (GNU GPL v3),
available at: https://www.gnu.org/licenses/gpl-3.0.txt available at: https://www.gnu.org/licenses/gpl-3.0.txt
@ -20,6 +22,8 @@ and generates a configuration for it.
Installation Details: Installation Details:
- ffmpeg: %LOCALAPPDATA%\\video2x\\ffmpeg - ffmpeg: %LOCALAPPDATA%\\video2x\\ffmpeg
- waifu2x-caffe: %LOCALAPPDATA%\\video2x\\waifu2x-caffe - 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 argparse
@ -46,7 +50,7 @@ def process_arguments():
# video options # video options
general_options = parser.add_argument_group('General 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 # parse arguments
return parser.parse_args() return parser.parse_args()
@ -75,10 +79,13 @@ class Video2xSetup:
if self.driver == 'all': if self.driver == 'all':
self._install_waifu2x_caffe() self._install_waifu2x_caffe()
self._install_waifu2x_converter_cpp() self._install_waifu2x_converter_cpp()
self._install_waifu2x_ncnn_vulkan()
elif self.driver == 'waifu2x_caffe': elif self.driver == 'waifu2x_caffe':
self._install_waifu2x_caffe() self._install_waifu2x_caffe()
elif self.driver == 'waifu2x_converter': elif self.driver == 'waifu2x_converter':
self._install_waifu2x_converter_cpp() self._install_waifu2x_converter_cpp()
elif self.driver == 'waifu2x-ncnn-vulkan':
self._install_waifu2x_ncnn_vulkan()
print('\nGenerating Video2X configuration file') print('\nGenerating Video2X configuration file')
self._generate_config() self._generate_config()
@ -96,8 +103,12 @@ class Video2xSetup:
""" """
for file in self.trash: for file in self.trash:
try: 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))
os.rmdir(file)
except FileNotFoundError: except FileNotFoundError:
pass pass
@ -147,6 +158,30 @@ class Video2xSetup:
with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf: with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf:
zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x', 'waifu2x-converter-cpp')) 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
import shutil
# 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): def _generate_config(self):
""" Generate video2x config """ Generate video2x config
""" """
@ -161,10 +196,15 @@ class Video2xSetup:
if self.driver == 'all': 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_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') 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': 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') 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': elif self.driver == 'waifu2x_converter':
template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-converter-cpp') 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['ffmpeg']['ffmpeg_path'] = os.path.join(local_app_data, 'video2x', 'ffmpeg-latest-win64-static', 'bin')
template_dict['video2x']['video2x_cache_directory'] = None template_dict['video2x']['video2x_cache_directory'] = None

100
bin/waifu2x_ncnn_vulkan.py Normal file
View File

@ -0,0 +1,100 @@
#!/usr/bin/env python3
# -*- coding: future_fstrings -*-
"""
Name: Waifu2x NCNN Vulkan Driver
Author: K4YT3X, SAT3LL
Description: This class is a high-level wrapper
for waifu2x-ncnn-vulkan.
"""
from avalon_framework import Avalon
import subprocess
import threading
import os
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)