mirror of
https://github.com/k4yt3x/video2x.git
synced 2026-02-15 01:34:47 +08:00
111 lines
5.0 KiB
Python
Executable File
111 lines
5.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Name: Waifu2x Caffe Driver
|
|
Author: K4YT3X
|
|
Date Created: Feb 24, 2018
|
|
Last Modified: June 7, 2020
|
|
|
|
Description: This class is a high-level wrapper
|
|
for waifu2x-caffe.
|
|
"""
|
|
|
|
# built-in imports
|
|
import argparse
|
|
import os
|
|
import pathlib
|
|
import shlex
|
|
import subprocess
|
|
import threading
|
|
|
|
# third-party imports
|
|
from avalon_framework import Avalon
|
|
|
|
|
|
class WrapperMain:
|
|
"""This class communicates with waifu2x cui 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, driver_settings):
|
|
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.error = lambda message: (_ for _ in ()).throw(AttributeError(message))
|
|
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', 'auto_scale'], help='image processing mode')
|
|
parser.add_argument('-e', '--output_extention', 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_path', type=str, help=argparse.SUPPRESS) # help='path to output image file (when input_path is folder, output_path must be folder)')
|
|
parser.add_argument('-i', '--input_path', type=str, help=argparse.SUPPRESS) # help='(required) path to input image file')
|
|
return parser.parse_args(arguments)
|
|
|
|
def load_configurations(self, upscaler):
|
|
# use scale width and scale height if specified
|
|
self.driver_settings['scale_ratio'] = upscaler.scale_ratio
|
|
self.driver_settings['output_extention'] = upscaler.extracted_frame_format
|
|
|
|
# bit_depth will be 12 at this point
|
|
# it will up updated later
|
|
self.driver_settings['output_depth'] = 12
|
|
|
|
def upscale(self, input_directory, output_directory):
|
|
""" start upscaling process
|
|
"""
|
|
|
|
# change the working directory to the binary's parent directory
|
|
# so the binary can find shared object files and other files
|
|
os.chdir(pathlib.Path(self.driver_settings['path']).parent)
|
|
|
|
# overwrite config file settings
|
|
self.driver_settings['input_path'] = input_directory
|
|
self.driver_settings['output_path'] = output_directory
|
|
|
|
# list to be executed
|
|
# initialize the list with waifu2x binary path as the first element
|
|
execute = [self.driver_settings['path']]
|
|
|
|
for key in self.driver_settings.keys():
|
|
|
|
value = self.driver_settings[key]
|
|
|
|
# null or None means that leave this option out (keep default)
|
|
if key == 'path' or value is None or value is False:
|
|
continue
|
|
else:
|
|
if len(key) == 1:
|
|
execute.append(f'-{key}')
|
|
else:
|
|
execute.append(f'--{key}')
|
|
|
|
# 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()
|
|
Avalon.debug_info(f'[upscaler] Subprocess {os.getpid()} executing: {shlex.join(execute)}')
|
|
self.print_lock.release()
|
|
return subprocess.Popen(execute)
|