diff --git a/bin/upscaler.py b/bin/upscaler.py index e4e7e2f..c3b603d 100644 --- a/bin/upscaler.py +++ b/bin/upscaler.py @@ -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,8 @@ class Upscaler: self.upscaler_exceptions = [] # initialize waifu2x driver - if self.waifu2x_driver != 'waifu2x_caffe' and self.waifu2x_driver != 'waifu2x_converter': + drivers = ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan'] + if self.waifu2x_driver not in drivers: raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}') # it's easier to do multi-threading with waifu2x_converter @@ -166,84 +170,109 @@ class Upscaler: self.progress_bar_exit_signal = True progress_bar.join() return + else: + # 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 + if self.waifu2x_driver == 'waifu2x_caffe': + 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)) - # create a separate w2 instance for each thread - w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir) + # if the driver being used is waifu2x_ncnn_vulkan + elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan': + w2 = Waifu2xNcnnVulkan(copy.deepcopy(self.waifu2x_settings)) + thread = threading.Thread(target=w2.upscale, + args=(thread_info[0], + self.upscaled_frames, + self.scale_ratio, + self.upscaler_exceptions)) - # 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 + 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]) - if len(self.upscaler_exceptions) != 0: - raise(self.upscaler_exceptions[0]) def run(self): """Main controller for Video2X diff --git a/bin/video2x.json b/bin/video2x.json index 1e78a42..19b4d73 100644 --- a/bin/video2x.json +++ b/bin/video2x.json @@ -39,6 +39,14 @@ "output": null, "input": null }, + "waifu2x_ncnn_vulkan": { + "waifu2x_ncnn_vulkan_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\waifu2x-ncnn-vulkan\\waifu2x.exe", + "input": null, + "output": null, + "noise-level": 2, + "scale-ratio": null, + "block-size": 400 + }, "ffmpeg": { "ffmpeg_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\ffmpeg-latest-win64-static\\bin", "video_to_frames": { diff --git a/bin/video2x.py b/bin/video2x.py index 6100996..6178bf9 100644 --- a/bin/video2x.py +++ b/bin/video2x.py @@ -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/NCNN accepts only scaling ratio') + exit(1) +if args.driver == 'waifu2x_ncnn_vulkan' and (args.ratio > 2 or not args.ratio.is_integer()): + Avalon.error('Scaling ratio must be 1 or 2 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'] diff --git a/bin/video2x_setup.py b/bin/video2x_setup.py index a25fcd6..4dbd0b0 100644 --- a/bin/video2x_setup.py +++ b/bin/video2x_setup.py @@ -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,11 +22,14 @@ 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 import json import os +import shutil import subprocess import sys import tempfile @@ -36,7 +41,7 @@ import zipfile # later in the script. # import requests -VERSION = '1.2.1' +VERSION = '1.3.0' def process_arguments(): @@ -46,7 +51,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 +80,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 +104,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)) + shutil.rmtree(file) except FileNotFoundError: pass @@ -147,6 +159,29 @@ 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 + + # 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 diff --git a/bin/waifu2x_ncnn_vulkan.py b/bin/waifu2x_ncnn_vulkan.py new file mode 100644 index 0000000..38b51ad --- /dev/null +++ b/bin/waifu2x_ncnn_vulkan.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +# -*- coding: future_fstrings -*- + + +""" +Name: Waifu2x NCNN Vulkan Driver +Author: SAT3LL +Date Created: June 26, 2019 +Last Modified: June 26, 2019 + +Dev: K4YT3X + +Description: This class is a high-level wrapper +for waifu2x_ncnn_vulkan. +""" +from avalon_framework import Avalon +import os +import subprocess +import threading + + +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)