mirror of
https://github.com/k4yt3x/video2x.git
synced 2026-05-17 16:37:37 +08:00
4.0.0 new argument parsing mechanism and project structure
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user