mirror of
https://github.com/k4yt3x/video2x.git
synced 2024-12-28 06:59:11 +00:00
2.6.0 complete redesign of configuration file, with lots of enhancements
This commit is contained in:
k4yt3x
parent
fe06c3ec16
commit
9639b0b990
@ -4,11 +4,12 @@
|
||||
Name: FFMPEG Class
|
||||
Author: K4YT3X
|
||||
Date Created: Feb 24, 2018
|
||||
Last Modified: March 4, 2019
|
||||
Last Modified: March 9, 2019
|
||||
|
||||
Description: This class handles all FFMPEG related
|
||||
operations.
|
||||
"""
|
||||
from avalon_framework import Avalon
|
||||
import json
|
||||
import subprocess
|
||||
|
||||
@ -21,11 +22,22 @@ class Ffmpeg:
|
||||
and inserting audio tracks to videos.
|
||||
"""
|
||||
|
||||
def __init__(self, ffmpeg_path, ffmpeg_arguments, hardware_acc=False):
|
||||
self.ffmpeg_path = ffmpeg_path
|
||||
self.ffmpeg_binary = '\"{}ffmpeg.exe\"'.format(ffmpeg_path)
|
||||
self.hardware_acc = hardware_acc
|
||||
self.ffmpeg_arguments = ffmpeg_arguments
|
||||
def __init__(self, ffmpeg_settings):
|
||||
self.ffmpeg_settings = ffmpeg_settings
|
||||
self._parse_settings()
|
||||
|
||||
def _parse_settings(self):
|
||||
""" Parse ffmpeg settings
|
||||
"""
|
||||
self.ffmpeg_path = self.ffmpeg_settings['ffmpeg_path']
|
||||
# Add a forward slash to directory if not present
|
||||
# otherwise there will be a format error
|
||||
if self.ffmpeg_path[-1] != '/' and self.ffmpeg_path[-1] != '\\':
|
||||
self.ffmpeg_path = '{}/'.format(self.ffmpeg_path)
|
||||
|
||||
self.ffmpeg_binary = '\"{}ffmpeg.exe\"'.format(self.ffmpeg_path)
|
||||
self.ffmpeg_hwaccel = self.ffmpeg_settings['ffmpeg_hwaccel']
|
||||
self.extra_arguments = self.ffmpeg_settings['extra_arguments']
|
||||
|
||||
def get_video_info(self, input_video):
|
||||
""" Gets input video information
|
||||
@ -39,7 +51,9 @@ class Ffmpeg:
|
||||
Returns:
|
||||
dictionary -- JSON text of input video information
|
||||
"""
|
||||
json_str = subprocess.check_output('\"{}ffprobe.exe\" -v quiet -print_format json -show_format -show_streams \"{}\"'.format(self.ffmpeg_path, input_video))
|
||||
execute = '\"{}ffprobe.exe\" -v quiet -print_format json -show_format -show_streams \"{}\"'.format(self.ffmpeg_path, input_video)
|
||||
Avalon.debug_info('Executing: {}'.format(execute))
|
||||
json_str = subprocess.check_output(execute)
|
||||
return json.loads(json_str.decode('utf-8'))
|
||||
|
||||
def extract_frames(self, input_video, extracted_frames):
|
||||
@ -53,8 +67,8 @@ class Ffmpeg:
|
||||
extracted_frames {string} -- video output folder
|
||||
"""
|
||||
execute = '{} -i \"{}\" \"{}\"\\extracted_%0d.png -y {}'.format(
|
||||
self.ffmpeg_binary, input_video, extracted_frames, ' '.join(self.ffmpeg_arguments))
|
||||
print('Executing: {}'.format(execute))
|
||||
self.ffmpeg_binary, input_video, extracted_frames, ' '.join(self.extra_arguments))
|
||||
Avalon.debug_info('Executing: {}'.format(execute))
|
||||
subprocess.run(execute, shell=True, check=True)
|
||||
|
||||
def convert_video(self, framerate, resolution, upscaled_frames):
|
||||
@ -69,8 +83,8 @@ class Ffmpeg:
|
||||
upscaled_frames {string} -- source images folder
|
||||
"""
|
||||
execute = '{} -r {} -f image2 -s {} -i \"{}\"\\extracted_%d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p \"{}\"\\no_audio.mp4 -y {}'.format(
|
||||
self.ffmpeg_binary, framerate, resolution, upscaled_frames, upscaled_frames, ' '.join(self.ffmpeg_arguments))
|
||||
print('Executing: {}'.format(execute))
|
||||
self.ffmpeg_binary, framerate, resolution, upscaled_frames, upscaled_frames, ' '.join(self.extra_arguments))
|
||||
Avalon.debug_info('Executing: {}'.format(execute))
|
||||
subprocess.run(execute, shell=True, check=True)
|
||||
|
||||
def migrate_audio_tracks_subtitles(self, input_video, output_video, upscaled_frames):
|
||||
@ -82,8 +96,8 @@ class Ffmpeg:
|
||||
upscaled_frames {string} -- directory containing upscaled frames
|
||||
"""
|
||||
execute = '{} -i \"{}\"\\no_audio.mp4 -i \"{}\" -map 0:v:0? -map 1? -c copy -map -1:v? \"{}\" -y {}'.format(
|
||||
self.ffmpeg_binary, upscaled_frames, input_video, output_video, ' '.join(self.ffmpeg_arguments))
|
||||
self.ffmpeg_binary, upscaled_frames, input_video, output_video, ' '.join(self.extra_arguments))
|
||||
# execute = '{} -i \"{}\"\\no_audio.mp4 -i \"{}\" -c:a copy -c:v copy -c:s copy -map 0:v? -map 1:a? -map 1:s? \"{}\" -y {}'.format(
|
||||
# self.ffmpeg_binary, upscaled_frames, input_video, output_video, ' '.join(self.ffmpeg_arguments))
|
||||
print('Executing: {}'.format(execute))
|
||||
Avalon.debug_info('Executing: {}'.format(execute))
|
||||
subprocess.run(execute, shell=True, check=True)
|
||||
|
||||
@ -4,7 +4,7 @@
|
||||
Name: Video2X Upscaler
|
||||
Author: K4YT3X
|
||||
Date Created: December 10, 2018
|
||||
Last Modified: March 4, 2019
|
||||
Last Modified: March 9, 2019
|
||||
|
||||
Licensed under the GNU General Public License Version 3 (GNU GPL v3),
|
||||
available at: https://www.gnu.org/licenses/gpl-3.0.txt
|
||||
@ -32,13 +32,13 @@ MODELS_AVAILABLE = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo', 'anime_sty
|
||||
|
||||
class Upscaler:
|
||||
|
||||
def __init__(self, input_video, output_video, method, waifu2x_path, ffmpeg_path, waifu2x_driver='waifu2x_caffe', ffmpeg_arguments=[], ffmpeg_hwaccel='auto', output_width=False, output_height=False, ratio=False, model_type='anime_style_art_rgb', threads=3, video2x_cache_folder='{}\\video2x'.format(tempfile.gettempdir()), preserve_frames=False):
|
||||
def __init__(self, input_video, output_video, method, waifu2x_settings, ffmpeg_settings, waifu2x_driver='waifu2x_caffe', scale_width=False, scale_height=False, scale_ratio=False, model_type='models/cunet', threads=5, video2x_cache_folder='{}\\video2x'.format(tempfile.gettempdir()), preserve_frames=False):
|
||||
# Mandatory arguments
|
||||
self.input_video = input_video
|
||||
self.output_video = output_video
|
||||
self.method = method
|
||||
self.waifu2x_path = waifu2x_path
|
||||
self.ffmpeg_path = ffmpeg_path
|
||||
self.waifu2x_settings = waifu2x_settings
|
||||
self.ffmpeg_settings = ffmpeg_settings
|
||||
self.waifu2x_driver = waifu2x_driver
|
||||
|
||||
# Check sanity of waifu2x_driver option
|
||||
@ -46,11 +46,9 @@ class Upscaler:
|
||||
raise Exception('Unrecognized waifu2x driver: {}'.format(waifu2x_driver))
|
||||
|
||||
# Optional arguments
|
||||
self.ffmpeg_arguments = ffmpeg_arguments
|
||||
self.ffmpeg_hwaccel = ffmpeg_hwaccel
|
||||
self.output_width = output_width
|
||||
self.output_height = output_height
|
||||
self.ratio = ratio
|
||||
self.scale_width = scale_width
|
||||
self.scale_height = scale_height
|
||||
self.scale_ratio = scale_ratio
|
||||
self.model_type = model_type
|
||||
self.threads = threads
|
||||
|
||||
@ -66,10 +64,6 @@ class Upscaler:
|
||||
|
||||
self.preserve_frames = preserve_frames
|
||||
|
||||
# If hardware acceleration enabled, append arguments
|
||||
if self.ffmpeg_hwaccel:
|
||||
self.ffmpeg_arguments.append('-hwaccel {}'.format(self.ffmpeg_hwaccel))
|
||||
|
||||
def cleanup(self):
|
||||
# Delete temp directories when done
|
||||
# Avalon framework cannot be used if python is shutting down
|
||||
@ -84,14 +78,16 @@ class Upscaler:
|
||||
""" Validate upscaling model
|
||||
"""
|
||||
if self.model_type not in MODELS_AVAILABLE:
|
||||
raise InvalidModelType('Specified model type not available')
|
||||
if self.model_type.split('/')[-1] not in MODELS_AVAILABLE:
|
||||
if self.model_type.split('\\')[-1] not in MODELS_AVAILABLE:
|
||||
raise InvalidModelType('Specified model type not available')
|
||||
|
||||
def _check_arguments(self):
|
||||
# Check if arguments are valid / all necessary argument
|
||||
# values are specified
|
||||
if not self.input_video:
|
||||
raise ArgumentError('You need to specify the video to process')
|
||||
elif (not self.output_width or not self.output_height) and not self.ratio:
|
||||
elif (not self.scale_width or not self.scale_height) and not self.scale_ratio:
|
||||
raise ArgumentError('You must specify output video width and height or upscale factor')
|
||||
elif not self.output_video:
|
||||
raise ArgumentError('You need to specify the output video name')
|
||||
@ -155,7 +151,7 @@ class Upscaler:
|
||||
progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],))
|
||||
progress_bar.start()
|
||||
|
||||
w2.upscale(self.extracted_frames, self.upscaled_frames, self.ratio, self.threads)
|
||||
w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads)
|
||||
for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]:
|
||||
renamed = re.sub('_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.png', '.png', image)
|
||||
shutil.move('{}\\{}'.format(self.upscaled_frames, image), '{}\\{}'.format(self.upscaled_frames, renamed))
|
||||
@ -203,7 +199,10 @@ class Upscaler:
|
||||
# Create threads and start them
|
||||
for thread_info in thread_pool:
|
||||
# Create thread
|
||||
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.output_width, self.output_height))
|
||||
if self.scale_ratio:
|
||||
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height))
|
||||
else:
|
||||
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False))
|
||||
thread.name = thread_info[1]
|
||||
|
||||
# Add threads into the pool
|
||||
@ -239,25 +238,14 @@ class Upscaler:
|
||||
self.input_video = os.path.abspath(self.input_video)
|
||||
self.output_video = os.path.abspath(self.output_video)
|
||||
|
||||
# Add a forward slash to directory if not present
|
||||
# otherwise there will be a format error
|
||||
if self.ffmpeg_path[-1] != '/' and self.ffmpeg_path[-1] != '\\':
|
||||
self.ffmpeg_path = '{}/'.format(self.ffmpeg_path)
|
||||
|
||||
# Check if FFMPEG and waifu2x are present
|
||||
if not os.path.isdir(self.ffmpeg_path):
|
||||
raise FileNotFoundError(self.ffmpeg_path)
|
||||
if not os.path.isfile(self.waifu2x_path) and not os.path.isdir(self.waifu2x_path):
|
||||
raise FileNotFoundError(self.waifu2x_path)
|
||||
|
||||
# Initialize objects for ffmpeg and waifu2x-caffe
|
||||
fm = Ffmpeg(self.ffmpeg_path, self.ffmpeg_arguments)
|
||||
fm = Ffmpeg(self.ffmpeg_settings)
|
||||
|
||||
# Initialize waifu2x driver
|
||||
if self.waifu2x_driver == 'waifu2x_caffe':
|
||||
w2 = Waifu2xCaffe(self.waifu2x_path, self.method, self.model_type)
|
||||
w2 = Waifu2xCaffe(self.waifu2x_settings, self.method, self.model_type)
|
||||
elif self.waifu2x_driver == 'waifu2x_converter':
|
||||
w2 = Waifu2xConverter(self.waifu2x_path)
|
||||
w2 = Waifu2xConverter(self.waifu2x_settings)
|
||||
else:
|
||||
raise Exception('Unrecognized waifu2x driver: {}'.format(self.waifu2x_driver))
|
||||
|
||||
@ -286,11 +274,11 @@ class Upscaler:
|
||||
Avalon.info('Framerate: {}'.format(framerate))
|
||||
|
||||
# Width/height will be coded width/height x upscale factor
|
||||
if self.ratio:
|
||||
if self.scale_ratio:
|
||||
coded_width = video_info['streams'][video_stream_index]['coded_width']
|
||||
coded_height = video_info['streams'][video_stream_index]['coded_height']
|
||||
self.output_width = self.ratio * coded_width
|
||||
self.output_height = self.ratio * coded_height
|
||||
self.scale_width = self.scale_ratio * coded_width
|
||||
self.scale_height = self.scale_ratio * coded_height
|
||||
|
||||
# Upscale images one by one using waifu2x
|
||||
Avalon.info('Starting to upscale extracted images')
|
||||
@ -301,7 +289,7 @@ class Upscaler:
|
||||
Avalon.info('Converting extracted frames into video')
|
||||
|
||||
# Use user defined output size
|
||||
fm.convert_video(framerate, '{}x{}'.format(self.output_width, self.output_height), self.upscaled_frames)
|
||||
fm.convert_video(framerate, '{}x{}'.format(self.scale_width, self.scale_height), self.upscaled_frames)
|
||||
Avalon.info('Conversion completed')
|
||||
|
||||
# Migrate audio tracks and subtitles
|
||||
|
||||
@ -1,8 +1,49 @@
|
||||
{
|
||||
"waifu2x_path": "C:/Program Files (x86)/waifu2x-caffe/waifu2x-caffe-cui.exe",
|
||||
"ffmpeg_path": "C:/Program Files (x86)/ffmpeg/bin/",
|
||||
"ffmpeg_arguments": [],
|
||||
"ffmpeg_hwaccel": "auto",
|
||||
"video2x_cache_folder": false,
|
||||
"preserve_frames": false
|
||||
}
|
||||
"waifu2x_caffe": {
|
||||
"waifu2x_caffe_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe",
|
||||
"mode": "noise_scale",
|
||||
"scale_ratio": null,
|
||||
"scale_width": null,
|
||||
"scale_height": null,
|
||||
"noise_level": 3,
|
||||
"process": "gpu",
|
||||
"crop_size": 128,
|
||||
"output_quality": -1,
|
||||
"output_depth": 8,
|
||||
"batch_size": 1,
|
||||
"gpu": 0,
|
||||
"tta": 0,
|
||||
"input_path": null,
|
||||
"output_path": null,
|
||||
"model_dir": "models/cunet",
|
||||
"crop_w": null,
|
||||
"crop_h": null
|
||||
},
|
||||
|
||||
"waifu2x_converter": {
|
||||
"waifu2x_converter_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\waifu2x-converter-cpp",
|
||||
"block_size": null,
|
||||
"disable-gpu": null,
|
||||
"force-OpenCL": null,
|
||||
"processor": null,
|
||||
"jobs": null,
|
||||
"model_dir": null,
|
||||
"scale_ratio": null,
|
||||
"noise_level": 3,
|
||||
"mode": "noise_scale",
|
||||
"quiet": true,
|
||||
"output": null,
|
||||
"input": null
|
||||
},
|
||||
|
||||
"ffmpeg":{
|
||||
"ffmpeg_path": "C:\\Users\\K4YT3X\\AppData\\Local\\video2x\\ffmpeg-4.1-win64-static\\bin",
|
||||
"ffmpeg_hwaccel": "auto",
|
||||
"extra_arguments": []
|
||||
},
|
||||
|
||||
"video2x": {
|
||||
"video2x_cache_folder": false,
|
||||
"preserve_frames": false
|
||||
}
|
||||
}
|
||||
@ -13,7 +13,7 @@ __ __ _ _ ___ __ __
|
||||
Name: Video2X Controller
|
||||
Author: K4YT3X
|
||||
Date Created: Feb 24, 2018
|
||||
Last Modified: March 4, 2019
|
||||
Last Modified: March 9, 2019
|
||||
|
||||
Licensed under the GNU General Public License Version 3 (GNU GPL v3),
|
||||
available at: https://www.gnu.org/licenses/gpl-3.0.txt
|
||||
@ -38,7 +38,7 @@ import tempfile
|
||||
import time
|
||||
import traceback
|
||||
|
||||
VERSION = '2.5.0'
|
||||
VERSION = '2.6.0'
|
||||
|
||||
# Each thread might take up to 2.5 GB during initialization.
|
||||
# (system memory, not to be confused with GPU memory)
|
||||
@ -53,7 +53,7 @@ def process_arguments():
|
||||
This allows users to customize options
|
||||
for the output video.
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
|
||||
# Video options
|
||||
basic_options = parser.add_argument_group('Basic Options')
|
||||
@ -61,7 +61,7 @@ def process_arguments():
|
||||
basic_options.add_argument('-o', '--output', help='Specify output video file/directory', action='store', default=False, required=True)
|
||||
basic_options.add_argument('-m', '--method', help='Specify upscaling method', action='store', default='gpu', choices=['cpu', 'gpu', 'cudnn'], required=True)
|
||||
basic_options.add_argument('-d', '--driver', help='Waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter'])
|
||||
basic_options.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb', choices=MODELS_AVAILABLE)
|
||||
basic_options.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='models/cunet', choices=MODELS_AVAILABLE)
|
||||
basic_options.add_argument('-t', '--threads', help='Specify number of threads to use for upscaling', action='store', type=int, default=5)
|
||||
basic_options.add_argument('-c', '--config', help='Manually specify config file', action='store', default='{}\\video2x.json'.format(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
@ -179,12 +179,19 @@ check_memory()
|
||||
|
||||
# Read configurations from JSON
|
||||
config = read_config(args.config)
|
||||
waifu2x_path = config['waifu2x_path']
|
||||
ffmpeg_path = config['ffmpeg_path']
|
||||
ffmpeg_arguments = config['ffmpeg_arguments']
|
||||
ffmpeg_hwaccel = config['ffmpeg_hwaccel']
|
||||
video2x_cache_folder = config['video2x_cache_folder']
|
||||
preserve_frames = config['preserve_frames']
|
||||
|
||||
# load waifu2x configuration
|
||||
if args.driver == 'waifu2x_caffe':
|
||||
waifu2x_settings = config['waifu2x_caffe']
|
||||
elif args.driver == 'waifu2x_converter':
|
||||
waifu2x_settings = config['waifu2x_converter']
|
||||
|
||||
# read FFMPEG configuration
|
||||
ffmpeg_settings = config['ffmpeg']
|
||||
|
||||
# load video2x settings
|
||||
video2x_cache_folder = config['video2x']['video2x_cache_folder']
|
||||
preserve_frames = config['video2x']['preserve_frames']
|
||||
|
||||
# Create temp directories if they don't exist
|
||||
if not video2x_cache_folder:
|
||||
@ -214,7 +221,7 @@ try:
|
||||
if os.path.isfile(args.input):
|
||||
""" Upscale single video file """
|
||||
Avalon.info('Upscaling single video file: {}'.format(args.input))
|
||||
upscaler = Upscaler(input_video=args.input, output_video=args.output, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, waifu2x_driver=args.driver, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, ratio=args.ratio, model_type=args.model_type, threads=args.threads, video2x_cache_folder=video2x_cache_folder)
|
||||
upscaler = Upscaler(input_video=args.input, output_video=args.output, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings, waifu2x_driver=args.driver, scale_width=args.width, scale_height=args.height, scale_ratio=args.ratio, model_type=args.model_type, threads=args.threads, video2x_cache_folder=video2x_cache_folder)
|
||||
upscaler.run()
|
||||
upscaler.cleanup()
|
||||
elif os.path.isdir(args.input):
|
||||
@ -222,7 +229,7 @@ try:
|
||||
Avalon.info('Upscaling videos in folder/directory: {}'.format(args.input))
|
||||
for input_video in [f for f in os.listdir(args.input) if os.path.isfile(os.path.join(args.input, f))]:
|
||||
output_video = '{}\\{}'.format(args.output, input_video)
|
||||
upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, waifu2x_driver=args.driver, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, ratio=args.ratio, model_type=args.model_type, threads=args.threads, video2x_cache_folder=video2x_cache_folder)
|
||||
upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings, waifu2x_driver=args.driver, scale_width=args.width, scale_height=args.height, scale_ratio=args.ratio, model_type=args.model_type, threads=args.threads, video2x_cache_folder=video2x_cache_folder)
|
||||
upscaler.run()
|
||||
upscaler.cleanup()
|
||||
else:
|
||||
|
||||
@ -4,7 +4,7 @@
|
||||
Name: Waifu2x Caffe Driver
|
||||
Author: K4YT3X
|
||||
Date Created: Feb 24, 2018
|
||||
Last Modified: March 4, 2019
|
||||
Last Modified: March 9, 2019
|
||||
|
||||
Description: This class controls waifu2x
|
||||
engine
|
||||
@ -23,33 +23,70 @@ class Waifu2xCaffe:
|
||||
the upscale function.
|
||||
"""
|
||||
|
||||
def __init__(self, waifu2x_path, method, model_type):
|
||||
self.waifu2x_path = waifu2x_path
|
||||
self.method = method
|
||||
self.model_type = model_type
|
||||
def __init__(self, waifu2x_settings, process, model_dir):
|
||||
self.waifu2x_settings = waifu2x_settings
|
||||
self.waifu2x_settings['process'] = process
|
||||
self.waifu2x_settings['model_dir'] = model_dir
|
||||
|
||||
# arguments passed through command line overwrites config file values
|
||||
self.process = process
|
||||
self.model_dir = model_dir
|
||||
self.print_lock = threading.Lock()
|
||||
|
||||
def upscale(self, folderin, folderout, width, height):
|
||||
def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height):
|
||||
"""This is the core function for WAIFU2X class
|
||||
|
||||
Arguments:
|
||||
folderin {string} -- source folder path
|
||||
folderout {string} -- output folder path
|
||||
input_folder {string} -- source folder path
|
||||
output_folder {string} -- output folder path
|
||||
width {int} -- output video width
|
||||
height {int} -- output video height
|
||||
"""
|
||||
|
||||
# overwrite config file settings
|
||||
self.waifu2x_settings['input_path'] = input_folder
|
||||
self.waifu2x_settings['output_path'] = output_folder
|
||||
|
||||
if scale_ratio:
|
||||
self.waifu2x_settings['scale_ratio'] = scale_ratio
|
||||
elif scale_width and scale_height:
|
||||
self.waifu2x_settings['scale_width'] = scale_width
|
||||
self.waifu2x_settings['scale_height'] = scale_height
|
||||
|
||||
# Print thread start message
|
||||
self.print_lock.acquire()
|
||||
Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name))
|
||||
self.print_lock.release()
|
||||
|
||||
"""
|
||||
# Create string for execution
|
||||
# execute = ['{}'.format(self.waifu2x_path), '-p', self.method, '-I', 'png', '-i', '\"{}\"'.format(folderin), '-e', 'png', '-o', folderout, '-w', str(width), '-h', str(height), '-n', '3', '-m', 'noise_scale', '-y', self.model_type]
|
||||
execute = '\"{}\" -p {} -I png -i \"{}\" -e png -o {} -w {} -h {} -n 3 -m noise_scale -y {}'.format(
|
||||
self.waifu2x_path, self.method, folderin, folderout, width, height, self.model_type)
|
||||
print('Executing: {}'.format(execute))
|
||||
subprocess.run(execute, shell=True, check=True)
|
||||
self.waifu2x_path, self.process, input_folder, output_folder, width, height, self.model_dir)
|
||||
"""
|
||||
|
||||
execute = []
|
||||
|
||||
for key in self.waifu2x_settings.keys():
|
||||
|
||||
value = self.waifu2x_settings[key]
|
||||
|
||||
# The key doesn't need to be passed in this case
|
||||
if key == 'waifu2x_caffe_path':
|
||||
execute.append(str(value))
|
||||
|
||||
# Null or None means that leave this option out (keep default)
|
||||
elif value is None or value is False:
|
||||
continue
|
||||
|
||||
else:
|
||||
if len(key) == 1:
|
||||
execute.append('-{}'.format(key))
|
||||
else:
|
||||
execute.append('--{}'.format(key))
|
||||
execute.append(str(value))
|
||||
|
||||
Avalon.debug_info('Executing: {}'.format(execute))
|
||||
subprocess.run(execute, check=True)
|
||||
|
||||
# Print thread exiting message
|
||||
self.print_lock.acquire()
|
||||
|
||||
@ -4,7 +4,7 @@
|
||||
Name: Waifu2x Converter CPP Driver
|
||||
Author: K4YT3X
|
||||
Date Created: February 8, 2019
|
||||
Last Modified: March 4, 2019
|
||||
Last Modified: March 9, 2019
|
||||
|
||||
Description: This class controls waifu2x
|
||||
engine
|
||||
@ -23,30 +23,62 @@ class Waifu2xConverter:
|
||||
the upscale function.
|
||||
"""
|
||||
|
||||
def __init__(self, waifu2x_path):
|
||||
self.waifu2x_path = waifu2x_path
|
||||
def __init__(self, waifu2x_settings):
|
||||
self.waifu2x_settings = waifu2x_settings
|
||||
self.print_lock = threading.Lock()
|
||||
|
||||
def upscale(self, folderin, folderout, scale_ratio, threads):
|
||||
def upscale(self, input_folder, output_folder, scale_ratio, jobs):
|
||||
""" Waifu2x Converter Driver Upscaler
|
||||
This method executes the upscaling of extracted frames.
|
||||
|
||||
Arguments:
|
||||
folderin {string} -- source folder path
|
||||
folderout {string} -- output folder path
|
||||
input_folder {string} -- source folder path
|
||||
output_folder {string} -- output folder path
|
||||
scale_ratio {int} -- frames' scale ratio
|
||||
threads {int} -- number of threads
|
||||
"""
|
||||
|
||||
# overwrite config file settings
|
||||
self.waifu2x_settings['input'] = input_folder
|
||||
self.waifu2x_settings['output'] = output_folder
|
||||
self.waifu2x_settings['scale_ratio'] = scale_ratio
|
||||
self.waifu2x_settings['jobs'] = jobs
|
||||
|
||||
# models_rgb must be specified manually
|
||||
self.waifu2x_settings['model_dir'] = '{}\\models_rgb'.format(self.waifu2x_settings['waifu2x_converter_path'])
|
||||
|
||||
# Print thread start message
|
||||
self.print_lock.acquire()
|
||||
Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name))
|
||||
self.print_lock.release()
|
||||
|
||||
# Create string for execution
|
||||
execute = '\"{}\\waifu2x-converter-cpp.exe\" -q -i \"{}\" -o {} --scale_ratio {} --noise_level 3 -m noise_scale -j {} --model_dir {}\\models_rgb'.format(
|
||||
self.waifu2x_path, folderin, folderout, scale_ratio, threads, self.waifu2x_path)
|
||||
print('Executing: {}'.format(execute))
|
||||
# Create list for execution
|
||||
execute = []
|
||||
|
||||
for key in self.waifu2x_settings.keys():
|
||||
|
||||
value = self.waifu2x_settings[key]
|
||||
|
||||
# The key doesn't need to be passed in this case
|
||||
if key == 'waifu2x_converter_path':
|
||||
execute.append('{}\\waifu2x-converter-cpp.exe'.format(str(value)))
|
||||
|
||||
# Null or None means that leave this option out (keep default)
|
||||
elif value is None or value is False:
|
||||
continue
|
||||
else:
|
||||
if len(key) == 1:
|
||||
execute.append('-{}'.format(key))
|
||||
else:
|
||||
execute.append('--{}'.format(key))
|
||||
|
||||
# True means key is an option
|
||||
if value is True:
|
||||
continue
|
||||
|
||||
execute.append(str(value))
|
||||
|
||||
Avalon.debug_info('Executing: {}'.format(execute))
|
||||
subprocess.run(execute, check=True)
|
||||
|
||||
# Print thread exiting message
|
||||
|
||||
Loading…
Reference in New Issue
Block a user