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
Author: K4YT3X
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),
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 waifu2x_caffe import Waifu2xCaffe
from waifu2x_converter import Waifu2xConverter
from waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan
import copy
import os
import re
@ -147,7 +150,7 @@ class Upscaler:
self.upscaler_exceptions = []
# 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}')
# it's easier to do multi-threading with waifu2x_converter
@ -166,84 +169,159 @@ class Upscaler:
self.progress_bar_exit_signal = True
progress_bar.join()
return
elif self.waifu2x_driver == 'waifu2x_caffe':
# create a container for all upscaler threads
upscaler_threads = []
# create a container for all upscaler threads
upscaler_threads = []
# list all images in the extracted frames
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
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)
# 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
# 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)
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)
# 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))
# 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]
# 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 threads and start them
for thread_info in thread_pool:
# create a separate w2 instance for each thread
w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir)
# create a separate w2 instance for each thread
w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir)
# create thread
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
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]
# add threads into the pool
upscaler_threads.append(thread)
# 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()
# 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()
# 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()
# start all threads
for thread in upscaler_threads:
thread.start()
# wait for threads to finish
for thread in upscaler_threads:
thread.join()
# 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()
# upscaling done, kill the clearer
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True
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:
raise(self.upscaler_exceptions[0])
# list all images in the extracted frames
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):
"""Main controller for Video2X

8
bin/video2x.json
View File

@ -39,6 +39,14 @@
"output": 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_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\ffmpeg-latest-win64-static\\bin",
"video_to_frames": {

23
bin/video2x.py
View File

@ -15,7 +15,7 @@ __ __ _ _ ___ __ __
Name: Video2X Controller
Author: K4YT3X
Date Created: Feb 24, 2018
Last Modified: June 15, 2019
Last Modified: June 26, 2019
Dev: BrianPetkovsek
Dev: SAT3LL
@ -82,7 +82,7 @@ def process_arguments():
# 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('-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('-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'))
@ -203,6 +203,10 @@ def absolutify_paths(config):
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'])
# 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
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'])
@ -244,8 +248,11 @@ if not args.input:
if not args.output:
Avalon.error('You must specify output video file/directory path')
exit(1)
if args.driver == 'waifu2x_converter' and args.width and args.height:
Avalon.error('Waifu2x Converter CPP accepts only scaling ratio')
if (args.driver == 'waifu2x_converter' or args.driver == 'waifu2x-ncnn-vulkan') and args.width and args.height:
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)
if (args.width or args.height) and args.ratio:
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':
waifu2x_settings = config['waifu2x_converter']
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')
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
ffmpeg_settings = config['ffmpeg']

48
bin/video2x_setup.py
View File

@ -7,7 +7,9 @@ Name: Video2X Setup Script
Author: K4YT3X
Author: BrianPetkovsek
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),
available at: https://www.gnu.org/licenses/gpl-3.0.txt
@ -20,6 +22,8 @@ and generates a configuration for it.
Installation Details:
- ffmpeg: %LOCALAPPDATA%\\video2x\\ffmpeg
- waifu2x-caffe: %LOCALAPPDATA%\\video2x\\waifu2x-caffe
- waifu2x-cpp-converter: %LOCALAPPDATA%\\video2x\\waifu2x-converter-cpp
- waifu2x-ncnn-vulkan: %LOCALAPPDATA%\\video2x\\waifu2x-ncnn-vulkan
"""
import argparse
@ -46,7 +50,7 @@ def process_arguments():
# video 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
return parser.parse_args()
@ -75,10 +79,13 @@ class Video2xSetup:
if self.driver == 'all':
self._install_waifu2x_caffe()
self._install_waifu2x_converter_cpp()
self._install_waifu2x_ncnn_vulkan()
elif self.driver == 'waifu2x_caffe':
self._install_waifu2x_caffe()
elif self.driver == 'waifu2x_converter':
self._install_waifu2x_converter_cpp()
elif self.driver == 'waifu2x-ncnn-vulkan':
self._install_waifu2x_ncnn_vulkan()
print('\nGenerating Video2X configuration file')
self._generate_config()
@ -96,8 +103,12 @@ class Video2xSetup:
"""
for file in self.trash:
try:
print('Deleting: {}'.format(file))
os.remove(file)
if os.path.isfile(file):
print('Deleting: {}'.format(file))
os.remove(file)
else:
print('Deleting: {}'.format(file))
os.rmdir(file)
except FileNotFoundError:
pass
@ -147,6 +158,30 @@ class Video2xSetup:
with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf:
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):
""" Generate video2x config
"""
@ -161,10 +196,15 @@ class Video2xSetup:
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_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':
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':
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['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)