mirror of
https://github.com/k4yt3x/video2x.git
synced 2024-12-29 16:09:10 +00:00
4.0.0 new argument parsing mechanism and project structure
This commit is contained in:
k4yt3x
parent
88299d404a
commit
89860f22cb
@ -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))
|
||||
|
||||
107
src/video2x.py
107
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')
|
||||
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user