diff --git a/src/upscaler.py b/src/upscaler.py index 684be25..d77bc0d 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 3, 2020 +Last Modified: May 4, 2020 Description: This file contains the Upscaler class. Each instance of the Upscaler class is an upscaler on an image or @@ -14,17 +14,13 @@ a folder. # local imports from exceptions import * from image_cleaner import ImageCleaner -from wrappers.anime4kcpp import Anime4kCpp from wrappers.ffmpeg import Ffmpeg -from wrappers.srmd_ncnn_vulkan import SrmdNcnnVulkan -from wrappers.waifu2x_caffe import Waifu2xCaffe -from wrappers.waifu2x_converter_cpp import Waifu2xConverterCpp -from wrappers.waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan # built-in imports from fractions import Fraction import contextlib import copy +import importlib import os import pathlib import re @@ -39,7 +35,15 @@ import traceback from avalon_framework import Avalon from tqdm import tqdm -AVAILABLE_DRIVERS = ['waifu2x_caffe', 'waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan', 'srmd_ncnn_vulkan', 'anime4kcpp'] +# these names are consistent for +# - driver selection in command line +# - driver wrapper file names +# - config file keys +AVAILABLE_DRIVERS = ['waifu2x_caffe', + 'waifu2x_converter_cpp', + 'waifu2x_ncnn_vulkan', + 'srmd_ncnn_vulkan', + 'anime4kcpp'] class Upscaler: @@ -63,7 +67,6 @@ class Upscaler: self.scale_width = None self.scale_height = None self.scale_ratio = None - self.model_dir = None self.processes = 1 self.video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x' self.image_format = 'png' @@ -214,27 +217,21 @@ class Upscaler: # create threads and start them for process_directory in process_directories: + DriverWrapperMain = getattr(importlib.import_module(f'wrappers.{self.driver}'), 'WrapperMain') + driver = DriverWrapperMain(copy.deepcopy(self.driver_settings)) + # if the driver being used is waifu2x-caffe if self.driver == 'waifu2x_caffe': - driver = Waifu2xCaffe(copy.deepcopy(self.driver_settings), self.model_dir, self.bit_depth) - if self.scale_ratio: - upscaler_processes.append(driver.upscale(process_directory, - self.upscaled_frames, - self.scale_ratio, - False, - False, - self.image_format)) - else: - upscaler_processes.append(driver.upscale(process_directory, - self.upscaled_frames, - False, - self.scale_width, - self.scale_height, - self.image_format)) + upscaler_processes.append(driver.upscale(process_directory, + self.upscaled_frames, + self.scale_ratio, + self.scale_width, + self.scale_height, + self.image_format, + self.bit_depth)) # if the driver being used is waifu2x-converter-cpp elif self.driver == 'waifu2x_converter_cpp': - driver = Waifu2xConverterCpp(self.driver_settings, self.model_dir) upscaler_processes.append(driver.upscale(process_directory, self.upscaled_frames, self.scale_ratio, @@ -243,14 +240,12 @@ class Upscaler: # if the driver being used is waifu2x-ncnn-vulkan elif self.driver == 'waifu2x_ncnn_vulkan': - driver = Waifu2xNcnnVulkan(copy.deepcopy(self.driver_settings)) upscaler_processes.append(driver.upscale(process_directory, self.upscaled_frames, self.scale_ratio)) # if the driver being used is srmd_ncnn_vulkan elif self.driver == 'srmd_ncnn_vulkan': - driver = SrmdNcnnVulkan(copy.deepcopy(self.driver_settings)) upscaler_processes.append(driver.upscale(process_directory, self.upscaled_frames, self.scale_ratio)) diff --git a/src/video2x.py b/src/video2x.py index a63d041..cbed695 100755 --- a/src/video2x.py +++ b/src/video2x.py @@ -13,7 +13,7 @@ __ __ _ _ ___ __ __ Name: Video2X Controller Creator: K4YT3X Date Created: Feb 24, 2018 -Last Modified: May 3, 2020 +Last Modified: May 4, 2020 Editor: BrianPetkovsek Last Modified: June 17, 2019 @@ -55,6 +55,7 @@ from upscaler import Upscaler # built-in imports import argparse import contextlib +import importlib import pathlib import re import shutil @@ -68,7 +69,7 @@ import yaml from avalon_framework import Avalon -VERSION = '3.3.0' +VERSION = '4.0.0' LEGAL_INFO = f'''Video2X Version: {VERSION} Author: K4YT3X @@ -87,26 +88,19 @@ LOGO = r''' def parse_arguments(): - """Processes CLI arguments - - This function parses all arguments - This allows users to customize options - for the output video. + """ parse CLI arguments """ - parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser = argparse.ArgumentParser(prog='video2x', formatter_class=argparse.ArgumentDefaultsHelpFormatter) # video options - file_options = parser.add_argument_group('File Options') - file_options.add_argument('-i', '--input', type=pathlib.Path, help='source video file/directory', action='store', required=True) - file_options.add_argument('-o', '--output', type=pathlib.Path, help='output video file/directory', action='store', required=True) - - # upscaler options - upscaler_options = parser.add_argument_group('Upscaler Options') - upscaler_options.add_argument('-d', '--driver', help='upscaling driver', action='store', default='waifu2x_caffe', choices=AVAILABLE_DRIVERS) - upscaler_options.add_argument('-y', '--model_dir', type=pathlib.Path, help='directory containing model JSON files', action='store') - upscaler_options.add_argument('-p', '--processes', help='number of processes to use for upscaling', action='store', type=int, default=1) - upscaler_options.add_argument('-c', '--config', type=pathlib.Path, help='video2x config file location', action='store', default=pathlib.Path(sys.argv[0]).parent.absolute() / 'video2x.yaml') - upscaler_options.add_argument('-b', '--batch', help='enable batch mode (select all default values to questions)', action='store_true') + general_options = parser.add_argument_group('General Options') + general_options.add_argument('-i', '--input', type=pathlib.Path, help='source video file/directory') + general_options.add_argument('-o', '--output', type=pathlib.Path, help='output video file/directory') + general_options.add_argument('-c', '--config', type=pathlib.Path, help='video2x config file location', action='store', + default=pathlib.Path(__file__).parent.absolute() / 'video2x.yaml') + general_options.add_argument('-d', '--driver', help='upscaling driver', choices=AVAILABLE_DRIVERS, default='waifu2x_caffe') + general_options.add_argument('-p', '--processes', help='number of processes to use for upscaling', action='store', type=int, default=1) + general_options.add_argument('-v', '--version', help='display version, lawful information and exit', action='store_true') # scaling options scaling_options = parser.add_argument_group('Scaling Options') @@ -114,12 +108,17 @@ def parse_arguments(): scaling_options.add_argument('--height', help='output video height', action='store', type=int) scaling_options.add_argument('-r', '--ratio', help='scaling ratio', action='store', type=float) - # extra options - extra_options = parser.add_argument_group('Extra Options') - extra_options.add_argument('-v', '--version', help='display version, lawful information and exit', action='store_true') + # if no driver arguments are specified + if '--' not in sys.argv: + video2x_args = parser.parse_args() + return video2x_args, None - # parse arguments - return parser.parse_args() + # if driver arguments are specified + else: + video2x_args = parser.parse_args(sys.argv[1:sys.argv.index('--')]) + wrapper = getattr(importlib.import_module(f'wrappers.{video2x_args.driver}'), 'WrapperMain') + driver_args = wrapper.parse_arguments(sys.argv[sys.argv.index('--') + 1:]) + return video2x_args, driver_args def print_logo(): @@ -154,18 +153,24 @@ if __name__ != '__main__': print_logo() # parse command line arguments -args = parse_arguments() +video2x_args, driver_args = parse_arguments() # display version and lawful informaition -if args.version: +if video2x_args.version: print(LEGAL_INFO) sys.exit(0) # read configurations from configuration file -config = read_config(args.config) +config = read_config(video2x_args.config) # load waifu2x configuration -driver_settings = config[args.driver] +driver_settings = config[video2x_args.driver] + +# overwrite driver_settings with driver_args +driver_args_dict = vars(driver_args) +for key in driver_args_dict: + if driver_args_dict[key] is not None: + driver_settings[key] = driver_args_dict[key] # check if driver path exists if not pathlib.Path(driver_settings['path']).exists(): @@ -211,33 +216,32 @@ try: begin_time = time.time() # if input specified is a single file - if args.input.is_file(): + if video2x_args.input.is_file(): # upscale single video file - Avalon.info(f'Upscaling single video file: {args.input}') + Avalon.info(f'Upscaling single video file: {video2x_args.input}') # check for input output format mismatch - if args.output.is_dir(): + if video2x_args.output.is_dir(): Avalon.error('Input and output path type mismatch') Avalon.error('Input is single file but output is directory') raise Exception('input output path type mismatch') - if not re.search(r'.*\..*$', str(args.output)): + if not re.search(r'.*\..*$', str(video2x_args.output)): Avalon.error('No suffix found in output file path') Avalon.error('Suffix must be specified for FFmpeg') raise Exception('No suffix specified') - upscaler = Upscaler(input_video=args.input, - output_video=args.output, + upscaler = Upscaler(input_video=video2x_args.input, + output_video=video2x_args.output, driver_settings=driver_settings, ffmpeg_settings=ffmpeg_settings) # set optional options - upscaler.driver = args.driver - upscaler.scale_width = args.width - upscaler.scale_height = args.height - upscaler.scale_ratio = args.ratio - upscaler.model_dir = args.model_dir - upscaler.processes = args.processes + upscaler.driver = video2x_args.driver + upscaler.scale_width = video2x_args.width + upscaler.scale_height = video2x_args.height + upscaler.scale_ratio = video2x_args.ratio + upscaler.processes = video2x_args.processes upscaler.video2x_cache_directory = video2x_cache_directory upscaler.image_format = image_format upscaler.preserve_frames = preserve_frames @@ -246,27 +250,26 @@ try: upscaler.run() # if input specified is a directory - elif args.input.is_dir(): + elif video2x_args.input.is_dir(): # upscale videos in a directory - Avalon.info(f'Upscaling videos in directory: {args.input}') + Avalon.info(f'Upscaling videos in directory: {video2x_args.input}') # make output directory if it doesn't exist - args.output.mkdir(parents=True, exist_ok=True) + video2x_args.output.mkdir(parents=True, exist_ok=True) - for input_video in [f for f in args.input.iterdir() if f.is_file()]: - output_video = args.output / input_video.name + for input_video in [f for f in video2x_args.input.iterdir() if f.is_file()]: + output_video = video2x_args.output / input_video.name upscaler = Upscaler(input_video=input_video, output_video=output_video, driver_settings=driver_settings, ffmpeg_settings=ffmpeg_settings) # set optional options - upscaler.driver = args.driver - upscaler.scale_width = args.width - upscaler.scale_height = args.height - upscaler.scale_ratio = args.ratio - upscaler.model_dir = args.model_dir - upscaler.processes = args.processes + upscaler.driver = video2x_args.driver + upscaler.scale_width = video2x_args.width + upscaler.scale_height = video2x_args.height + upscaler.scale_ratio = video2x_args.ratio + upscaler.processes = video2x_args.processes upscaler.video2x_cache_directory = video2x_cache_directory upscaler.image_format = image_format upscaler.preserve_frames = preserve_frames @@ -275,7 +278,7 @@ try: upscaler.run() else: Avalon.error('Input path is neither a file nor a directory') - raise FileNotFoundError(f'{args.input} is neither file nor directory') + raise FileNotFoundError(f'{video2x_args.input} is neither file nor directory') Avalon.info(f'Program completed, taking {round((time.time() - begin_time), 5)} seconds') diff --git a/src/wrappers/anime4kcpp.py b/src/wrappers/anime4kcpp.py index 4af01d3..904f3f2 100644 --- a/src/wrappers/anime4kcpp.py +++ b/src/wrappers/anime4kcpp.py @@ -4,13 +4,14 @@ Name: Waifu2x Caffe Driver Author: K4YT3X Date Created: May 3, 2020 -Last Modified: May 3, 2020 +Last Modified: May 4, 2020 Description: This class is a high-level wrapper for waifu2x-caffe. """ # built-in imports +import argparse import os import shlex import subprocess @@ -20,7 +21,7 @@ import threading from avalon_framework import Avalon -class Anime4kCpp: +class WrapperMain: """ Anime4K CPP wrapper """ @@ -28,6 +29,31 @@ class Anime4kCpp: 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.add_argument('--help', action='help', help='show this help message and exit') + # parser.add_argument('-i', '--input', type=pathlib.Path, help='File for loading') + # parser.add_argument('-o', '--output', type=pathlib.Path, help='File for outputting') + parser.add_argument('-p', '--passes', type=int, help='Passes for processing') + parser.add_argument('-n', '--pushColorCount', type=int, help='Limit the number of color pushes') + parser.add_argument('-c', '--strengthColor', type=float, help='Strength for pushing color,range 0 to 1,higher for thinner') + parser.add_argument('-g', '--strengthGradient', type=float, help='Strength for pushing gradient,range 0 to 1,higher for sharper') + parser.add_argument('-z', '--zoomFactor', type=float, help='zoom factor for resizing') + parser.add_argument('-t', '--threads', type=int, help='Threads count for video processing') + parser.add_argument('-f', '--fastMode', action='store_true', help='Faster but maybe low quality') + # parser.add_argument('-v', '--videoMode', action='store_true', help='Video process') + parser.add_argument('-s', '--preview', action='store_true', help='Preview image') + parser.add_argument('-b', '--preProcessing', action='store_true', help='Enable pre processing') + parser.add_argument('-a', '--postProcessing', action='store_true', help='Enable post processing') + parser.add_argument('-r', '--preFilters', type=int, help='Enhancement filter, only working when preProcessing is true,there are 5 options by binary:Median blur=0000001, Mean blur=0000010, CAS Sharpening=0000100, Gaussian blur weak=0001000, Gaussian blur=0010000, Bilateral filter=0100000, Bilateral filter faster=1000000, you can freely combine them, eg: Gaussian blur weak + Bilateral filter = 0001000 | 0100000 = 0101000 = 40(D)') + parser.add_argument('-e', '--postFilters', type=int, help='Enhancement filter, only working when postProcessing is true,there are 5 options by binary:Median blur=0000001, Mean blur=0000010, CAS Sharpening=0000100, Gaussian blur weak=0001000, Gaussian blur=0010000, Bilateral filter=0100000, Bilateral filter faster=1000000, you can freely combine them, eg: Gaussian blur weak + Bilateral filter = 0001000 | 0100000 = 0101000 = 40(D), so you can put 40 to enable Gaussian blur weak and Bilateral filter, which also is what I recommend for image that < 1080P, 48 for image that >= 1080P, and for performance I recommend to use 72 for video that < 1080P, 80 for video that >=1080P') + parser.add_argument('-q', '--GPUMode', action='store_true', help='Enable GPU acceleration') + parser.add_argument('-l', '--listGPUs', action='store_true', help='list GPUs') + parser.add_argument('-h', '--platformID', type=int, help='Specify the platform ID') + parser.add_argument('-d', '--deviceID', type=int, help='Specify the device ID') + return parser.parse_args(arguments) + def upscale(self, input_file, output_file, zoom_factor, threads): """This is the core function for WAIFU2X class @@ -52,7 +78,7 @@ class Anime4kCpp: value = self.driver_settings[key] - # is executable key or null or None means that leave this option out (keep default) + # null or None means that leave this option out (keep default) if value is None or value is False: continue else: diff --git a/src/wrappers/srmd_ncnn_vulkan.py b/src/wrappers/srmd_ncnn_vulkan.py index 04db35c..cf67b89 100644 --- a/src/wrappers/srmd_ncnn_vulkan.py +++ b/src/wrappers/srmd_ncnn_vulkan.py @@ -4,13 +4,14 @@ Name: SRMD NCNN Vulkan Driver Creator: K4YT3X Date Created: April 26, 2020 -Last Modified: April 26, 2020 +Last Modified: May 4, 2020 Description: This class is a high-level wrapper for srmd_ncnn_vulkan. """ # built-in imports +import argparse import os import shlex import subprocess @@ -20,7 +21,7 @@ import threading from avalon_framework import Avalon -class SrmdNcnnVulkan: +class WrapperMain: """This class communicates with SRMD NCNN Vulkan engine An object will be created for this class, containing information @@ -40,6 +41,22 @@ class SrmdNcnnVulkan: self.print_lock = threading.Lock() + @staticmethod + def parse_arguments(arguments): + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, add_help=False) + 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=pathlib.Path, help='input image path (jpg/png) or directory') + # parser.add_argument('-o', type=pathlib.Path, help='output image path (png) or directory') + parser.add_argument('-n', type=int, choices=range(-1, 11), help='denoise level') + parser.add_argument('-s', type=int, choices=range(2, 5), help='upscale ratio') + parser.add_argument('-t', type=int, help='tile size (>=32)') + parser.add_argument('-m', type=str, help='srmd 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 upscale(self, input_directory, output_directory, scale_ratio): """This is the core function for SRMD ncnn Vulkan class @@ -52,7 +69,7 @@ class SrmdNcnnVulkan: # overwrite config file settings self.driver_settings['i'] = input_directory self.driver_settings['o'] = output_directory - self.driver_settings['s'] = int(scale_ratio) + self.driver_settings['s'] = scale_ratio # list to be executed # initialize the list with the binary path as the first element @@ -62,15 +79,18 @@ class SrmdNcnnVulkan: value = self.driver_settings[key] - # is executable key or null or None means that leave this option out (keep default) - if key == 'path' or value is None or value is False: + # null or None means that leave this option out (keep default) + if value is None or value is False: continue else: if len(key) == 1: execute.append(f'-{key}') else: execute.append(f'--{key}') - execute.append(str(value)) + + # 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() diff --git a/src/wrappers/waifu2x_caffe.py b/src/wrappers/waifu2x_caffe.py index 147abfb..50df5ab 100644 --- a/src/wrappers/waifu2x_caffe.py +++ b/src/wrappers/waifu2x_caffe.py @@ -4,13 +4,14 @@ Name: Waifu2x Caffe Driver Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: May 3, 2020 +Last Modified: May 4, 2020 Description: This class is a high-level wrapper for waifu2x-caffe. """ # built-in imports +import argparse import os import shlex import subprocess @@ -20,7 +21,7 @@ import threading from avalon_framework import Avalon -class Waifu2xCaffe: +class WrapperMain: """This class communicates with waifu2x cui engine An object will be created for this class, containing information @@ -29,16 +30,36 @@ class Waifu2xCaffe: the upscale function. """ - def __init__(self, driver_settings, model_dir, bit_depth): + def __init__(self, driver_settings): self.driver_settings = driver_settings - self.driver_settings['model_dir'] = model_dir - self.driver_settings['output_depth'] = bit_depth - - # arguments passed through command line overwrites config file values - self.model_dir = model_dir self.print_lock = threading.Lock() - def upscale(self, input_directory, output_directory, scale_ratio, scale_width, scale_height, image_format): + @staticmethod + def parse_arguments(arguments): + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, add_help=False) + parser.add_argument('--help', action='help', help='show this help message and exit') + parser.add_argument('-t', '--tta', type=int, choices=range(2), help='8x slower and slightly high quality') + parser.add_argument('--gpu', type=int, help='gpu device no') + parser.add_argument('-b', '--batch_size', type=int, help='input batch size') + parser.add_argument('--crop_h', type=int, help='input image split size(height)') + parser.add_argument('--crop_w', type=int, help='input image split size(width)') + parser.add_argument('-c', '--crop_size', type=int, help='input image split size') + parser.add_argument('-d', '--output_depth', type=int, help='output image chaneel depth bit') + parser.add_argument('-q', '--output_quality', type=int, help='output image quality') + parser.add_argument('-p', '--process', choices=['cpu', 'gpu', 'cudnn'], help='process mode') + parser.add_argument('--model_dir', type=str, help='path to custom model directory (don\'t append last / )') + parser.add_argument('-h', '--scale_height', type=int, help='custom scale height') + parser.add_argument('-w', '--scale_width', type=int, help='custom scale width') + parser.add_argument('-s', '--scale_ratio', type=float, help='custom scale ratio') + parser.add_argument('-n', '--noise_level', type=int, choices=range(4), help='noise reduction level') + parser.add_argument('-m', '--mode', choices=['noise', 'scale', 'noise_scale'], help='image processing mode') + parser.add_argument('-e', '--output_extension', type=str, help='extention to output image file when output_path is (auto) or input_path is folder') + parser.add_argument('-l', '--input_extention_list', type=str, help='extention to input image file when input_path is folder') + # parser.add_argument('-o', '--output', type=pathlib.Path, help='path to output image file (when input_path is folder, output_path must be folder)') + # parser.add_argument('-i', '--input_file', type=pathlib.Path, help='(required) path to input image file') + return parser.parse_args(arguments) + + def upscale(self, input_directory, output_directory, scale_ratio, scale_width, scale_height, image_format, bit_depth): """This is the core function for WAIFU2X class Arguments: @@ -59,6 +80,7 @@ class Waifu2xCaffe: self.driver_settings['scale_height'] = scale_height self.driver_settings['output_extention'] = image_format + self.driver_settings['output_depth'] = bit_depth # list to be executed # initialize the list with waifu2x binary path as the first element @@ -68,7 +90,7 @@ class Waifu2xCaffe: value = self.driver_settings[key] - # is executable key or null or None means that leave this option out (keep default) + # null or None means that leave this option out (keep default) if value is None or value is False: continue else: diff --git a/src/wrappers/waifu2x_converter_cpp.py b/src/wrappers/waifu2x_converter_cpp.py index f4b97e6..1922b45 100644 --- a/src/wrappers/waifu2x_converter_cpp.py +++ b/src/wrappers/waifu2x_converter_cpp.py @@ -4,13 +4,14 @@ Name: Waifu2x Converter CPP Driver Author: K4YT3X Date Created: February 8, 2019 -Last Modified: May 3, 2020 +Last Modified: May 4, 2020 Description: This class is a high-level wrapper for waifu2x-converter-cpp. """ # built-in imports +import argparse import os import pathlib import shlex @@ -21,7 +22,7 @@ import threading from avalon_framework import Avalon -class Waifu2xConverterCpp: +class WrapperMain: """This class communicates with waifu2x cui engine An object will be created for this class, containing information @@ -30,11 +31,40 @@ class Waifu2xConverterCpp: the upscale function. """ - def __init__(self, driver_settings, model_dir): + def __init__(self, driver_settings): self.driver_settings = driver_settings - self.driver_settings['model_dir'] = model_dir self.print_lock = threading.Lock() + @staticmethod + def parse_arguments(arguments): + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, add_help=False) + parser.add_argument('--help', action='help', help='show this help message and exit') + parser.add_argument('--list-supported-formats', action='store_true', help='dump currently supported format list') + parser.add_argument('--list-opencv-formats', action='store_true', help='(deprecated. Use --list-supported-formats) dump opencv supported format list') + parser.add_argument('-l', '--list-processor', action='store_true', help='dump processor list') + parser.add_argument('-f', '--output-format', choices=['png', 'jpg'], help='The format used when running in recursive/folder mode\nSee --list-supported-formats for a list of supported formats/extensions.') + parser.add_argument('-c', '--png-compression', type=int, choices=range(10), help='Set PNG compression level (0-9), 9 = Max compression (slowest & smallest)') + parser.add_argument('-q', '--image-quality', type=int, choices=range(100), help='JPEG & WebP Compression quality (0-101, 0 being smallest size and lowest quality), use 101 for lossless WebP') + parser.add_argument('--block-size', type=int, help='block size') + parser.add_argument('--disable-gpu', action='store_true', help='disable GPU') + parser.add_argument('--force-OpenCL', action='store_true', help='force to use OpenCL on Intel Platform') + parser.add_argument('-p', '--processor', type=int, help='set target processor') + parser.add_argument('-j', '--jobs', type=int, help='number of threads launching at the same time') + parser.add_argument('--model-dir', type=str, help='path to custom model directory (don\'t append last / )') + parser.add_argument('--scale-ratio', type=float, help='custom scale ratio') + parser.add_argument('--noise-level', type=int, choices=range(4), help='noise reduction level') + parser.add_argument('-m', '--mode', choices=['noise', 'scale', 'noise-scale'], help='image processing mode') + parser.add_argument('-v', '--log-level', type=int, choices=range(5), help='Set log level') + parser.add_argument('-s', '--silent', action='store_true', help='Enable silent mode. (same as --log-level 1)') + parser.add_argument('-t', '--tta', type=int, choices=range(2), help='Enable Test-Time Augmentation mode.') + parser.add_argument('-g', '--generate-subdir', type=int, choices=range(2), help='Generate sub folder when recursive directory is enabled.') + parser.add_argument('-a', '--auto-naming', type=int, choices=range(2), help='Add postfix to output name when output path is not specified.\nSet 0 to disable this.') + parser.add_argument('-r', '--recursive-directory', type=int, choices=range(2), help='Search recursively through directories to find more images to process.') + # parser.add_argument('-o', '--output', type=pathlib.Pathh, help='path to output image file or directory (you should use the full path)') + # parser.add_argument('-i', '--input', type=pathlib.Path, help='(required) path to input image file or directory (you should use the full path)') + parser.add_argument('--version', action='store_true', help='Displays version information and exits.') + return parser.parse_args(arguments) + def upscale(self, input_directory, output_directory, scale_ratio, jobs, image_format): """ Waifu2x Converter Driver Upscaler This method executes the upscaling of extracted frames. diff --git a/src/wrappers/waifu2x_ncnn_vulkan.py b/src/wrappers/waifu2x_ncnn_vulkan.py index c3a7707..8b62a65 100644 --- a/src/wrappers/waifu2x_ncnn_vulkan.py +++ b/src/wrappers/waifu2x_ncnn_vulkan.py @@ -4,7 +4,7 @@ Name: Waifu2x NCNN Vulkan Driver Creator: SAT3LL Date Created: June 26, 2019 -Last Modified: November 15, 2019 +Last Modified: May 4, 2020 Editor: K4YT3X Last Modified: February 22, 2020 @@ -14,6 +14,7 @@ for waifu2x_ncnn_vulkan. """ # built-in imports +import argparse import os import shlex import subprocess @@ -23,7 +24,7 @@ import threading from avalon_framework import Avalon -class Waifu2xNcnnVulkan: +class WrapperMain: """This class communicates with waifu2x ncnn vulkan engine An object will be created for this class, containing information @@ -43,6 +44,22 @@ class Waifu2xNcnnVulkan: self.print_lock = threading.Lock() + @staticmethod + def parse_arguments(arguments): + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, add_help=False) + 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=pathlib.Path, help='input image path (jpg/png) or directory') + # parser.add_argument('-o', type=pathlib.Path, help='output image path (png) or directory') + parser.add_argument('-n', type=int, choices=range(-1, 4), help='denoise level') + parser.add_argument('-s', type=int, choices=range(1, 3), help='upscale ratio') + parser.add_argument('-t', type=int, help='tile size (>=32)') + parser.add_argument('-m', type=str, help='waifu2x 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 upscale(self, input_directory, output_directory, scale_ratio): """This is the core function for WAIFU2X class @@ -59,21 +76,24 @@ class Waifu2xNcnnVulkan: # list to be executed # initialize the list with waifu2x binary path as the first element - execute = [str(self.driver_settings['path'])] + execute = [self.driver_settings.pop('path')] for key in self.driver_settings.keys(): value = self.driver_settings[key] - # is executable key or null or None means that leave this option out (keep default) - if key == 'path' or value is None or value is False: + # null or None means that leave this option out (keep default) + if value is None or value is False: continue else: if len(key) == 1: execute.append(f'-{key}') else: execute.append(f'--{key}') - execute.append(str(value)) + + # 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()