diff --git a/bin/config_generator.py b/bin/config_generator.py new file mode 100755 index 0000000..74f3750 --- /dev/null +++ b/bin/config_generator.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Name: Video2x Config Generator +Author: K4YT3X +Date Created: October 23, 2018 +Last Modified: November 26, 2018 + +Licensed under the GNU General Public License Version 3 (GNU GPL v3), + available at: https://www.gnu.org/licenses/gpl-3.0.txt + +(C) 2018 K4YT3X +""" +from avalon_framework import Avalon +import json +import os + +VERSION = '1.0.0' + + +def get_path(text): + """ Get path and validate + """ + while True: + path = Avalon.gets(text) + if os.path.isdir(path): + return path + Avalon.error('{} id not a directory / folder'.format(path)) + + +def enroll_settings(): + settings = {} + + settings['waifu2x_path'] = get_path('waifu2x-caffe-cui.exe path: ') + settings['ffmpeg_path'] = get_path('ffmpeg binaries directory: ') + settings['ffmpeg_arguments'] = Avalon.gets('Extra arguments passed to ffmpeg: ') + + settings['ffmpeg_hwaccel'] = Avalon.gets('ffmpeg hardware acceleration method (cuda): ') + if settings['ffmpeg_hwaccel'] == '': + settings['ffmpeg_hwaccel'] = 'cuda' + + return settings + + +def write_config(settings): + with open('video2x.json', 'w') as config: + json.dump(settings, config, indent=2) + config.close() + + +try: + print('Video2X Config Generator {}'.format(VERSION)) + write_config(enroll_settings()) +except KeyboardInterrupt: + Avalon.warning('Exiting...') diff --git a/bin/ffmpeg.py b/bin/ffmpeg.py new file mode 100755 index 0000000..338133d --- /dev/null +++ b/bin/ffmpeg.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Name: FFMPEG Class +Author: K4YT3X +Date Created: Feb 24, 2018 +Last Modified: November 2, 2018 + +Description: This class handles all FFMPEG related +operations. + +Version 2.0.5 +""" +import subprocess + + +class Ffmpeg: + """This class communicates with ffmpeg + + This class deals with ffmpeg. It handles extracitng + frames, stripping audio, converting images into videos + and inserting audio tracks to videos. + """ + + def __init__(self, ffmpeg_path, outfile, ffmpeg_arguments, hardware_acc=False): + self.ffmpeg_path = '{}ffmpeg.exe'.format(ffmpeg_path) + self.outfile = outfile + self.hardware_acc = hardware_acc + self.ffmpeg_arguments = ffmpeg_arguments + + def extract_frames(self, videoin, outpath): + """Extract every frame from original videos + + This method extracts every frame from videoin + using ffmpeg + + Arguments: + videoin {string} -- input video path + outpath {string} -- video output folder + """ + execute = '\"{}\" -i \"{}\" {}\\extracted_%0d.png -y {}'.format( + self.ffmpeg_path, videoin, outpath, ' '.join(self.ffmpeg_arguments)) + print(execute) + subprocess.call(execute) + + def extract_audio(self, videoin, outpath): + """Strips audio tracks from videos + + This method strips audio tracks from videos + into the output folder in aac format. + + Arguments: + videoin {string} -- input video path + outpath {string} -- video output folder + """ + execute = '\"{}\" -i \"{}\" -vn -acodec copy {}\\output-audio.aac -y {}'.format( + self.ffmpeg_path, videoin, outpath, ' '.join(self.ffmpeg_arguments)) + print(execute) + subprocess.call(execute) + + def convert_video(self, framerate, resolution, upscaled, ): + """Converts images into videos + + This method converts a set of images into a + video. + + Arguments: + framerate {float} -- target video framerate + resolution {string} -- target video resolution + upscaled {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_path, framerate, resolution, upscaled, upscaled, ' '.join(self.ffmpeg_arguments)) + print(execute) + subprocess.call(execute) + + def insert_audio_track(self, upscaled): + """Insert audio into video + + Inserts the WAV audio track stripped from + the original video into final video. + + Arguments: + upscaled {string} -- upscaled image folder + """ + execute = '\"{}\" -i {}\\no_audio.mp4 -i {}\\output-audio.aac -shortest -codec copy {} -y {}'.format( + self.ffmpeg_path, upscaled, upscaled, self.outfile, ' '.join(self.ffmpeg_arguments)) + print(execute) + subprocess.call(execute) diff --git a/bin/upscaler.py b/bin/upscaler.py new file mode 100644 index 0000000..b7254ba --- /dev/null +++ b/bin/upscaler.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Name: Video2X Upscaler +Author: K4YT3X +Date Created: December 10, 2018 +Last Modified: December 10, 2018 + +Licensed under the GNU General Public License Version 3 (GNU GPL v3), + available at: https://www.gnu.org/licenses/gpl-3.0.txt + +(C) 2018 K4YT3X +""" + +from avalon_framework import Avalon +from ffmpeg import Ffmpeg +from fractions import Fraction +from waifu2x import Waifu2x +import json +import os +import psutil +import shutil +import subprocess +import tempfile +import threading + +# Each thread might take up to 2.5 GB during initialization. +MEM_PER_THREAD = 2.5 + + +class ArgumentError(Exception): + def __init__(self, message): + super().__init__(message) + + +class Upscaler: + + def __init__(self, input_video, output_video, method, waifu2x_path, ffmpeg_path, ffmpeg_arguments=[], ffmpeg_hwaccel='gpu', output_width=False, output_height=False, factor=False, model_type='anime_style_art_rgb', threads=3): + # 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 + + # Optional arguments + self.ffmpeg_arguments = ffmpeg_arguments + self.ffmpeg_hwaccel = ffmpeg_hwaccel + self.output_width = output_width + self.output_height = output_height + self.factor = factor + self.model_type = model_type + self.threads = threads + + # Make temporary directories + self.extracted_frames = tempfile.mkdtemp() + self.upscaled_frames = tempfile.mkdtemp() + + def _get_video_info(self): + """Gets input video information + returns input video information using ffprobe in dictionary. + """ + json_str = subprocess.check_output('\"{}ffprobe.exe\" -v quiet -print_format json -show_format -show_streams \"{}\"'.format(self.ffmpeg_path, self.input_video)) + return json.loads(json_str.decode('utf-8')) + + def _check_model_type(self, args): + """ Validate upscaling model + """ + models_available = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo', + 'anime_style_art_rgb', 'photo', 'anime_style_art_y'] + if self.model_type not in models_available: + Avalon.error('Specified model type not found!') + Avalon.info('Available models:') + for model in models_available: + print(model) + exit(1) + + 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.factor: + raise ArgumentError('You must specify output video width and height or upscale factor') + elif not self.video_output: + raise ArgumentError('You need to specify the output video name') + elif not self.method: + raise ArgumentError('You need to specify the enlarging processing unit') + + def _check_memory(self): + """ Check usable memory + Warn the user if insufficient memory is available for + the number of threads that the user have chosen. + """ + memory_available = psutil.virtual_memory().available / (1024 ** 3) + + # If user doesn't even have enough memory to run even one thread + if memory_available < MEM_PER_THREAD: + Avalon.warning('You might have an insufficient amount of memory available to run this program ({} GB)'.format(memory_available)) + Avalon.warning('Proceed with caution') + # If memory available is less than needed, warn the user + elif memory_available < (MEM_PER_THREAD * self.threads): + Avalon.warning('Each waifu2x-caffe thread will require up to 2.5 GB during initialization') + Avalon.warning('You demanded {} threads to be created, but you only have {} GB memory available'.format(self.threads, round(memory_available, 4))) + Avalon.warning('{} GB of memory is recommended for {} threads'.format(MEM_PER_THREAD * self.threads, self.threads)) + Avalon.warning('With your current amount of memory available, {} threads is recommended'.format(int(memory_available // MEM_PER_THREAD))) + + # Ask the user if he / she wants to change to the recommended + # number of threads + if Avalon.ask('Change to the recommended value?', True): + self.threads = int(memory_available // MEM_PER_THREAD) + else: + Avalon.warning('Proceed with caution') + + def _upscale_frames(self, w2): + """ Upscale video frames with waifu2x-caffe + + This function upscales all the frames extracted + by ffmpeg using the waifu2x-caffe binary. + + Arguments: + w2 {Waifu2x Object} -- initialized waifu2x object + """ + + # Create a container for all upscaler threads + upscaler_threads = [] + + # List all images in the extracted frames + frames = [os.path.join(self.extracted_frames, f) for f in os.listdir(self.extracted_frames) if os.path.isfile(os.path.join(self.extracted_frames, f))] + + # If we have less images than threads, + # create only the threads necessary + if len(frames) < self.threads: + self.threads = len(frames) + + # Create a folder for each thread and append folder + # name into a list + + thread_pool = [] + for thread_id in range(self.threads): + thread_folder = '{}\\{}'.format(self.extracted_frames, str(thread_id)) + + # Delete old folders and create new folders + if os.path.isdir(thread_folder): + shutil.rmtree(thread_folder) + os.mkdir(thread_folder) + + # Append folder path into list + thread_pool.append((thread_folder, thread_id)) + + # Evenly distribute images into each folder + # until there is none left in the folder + for image in frames: + # Move image + shutil.move(image, thread_pool[0][0]) + # Rotate list + thread_pool = thread_pool[-1:] + thread_pool[:-1] + + # 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.width, self.height)) + thread.name = thread_info[1] + + # Add threads into the pool + upscaler_threads.append(thread) + + # Start all threads + for thread in upscaler_threads: + thread.start() + + # Wait for threads to finish + for thread in upscaler_threads: + thread.join() + + def run(self): + """Main controller for Video2X + + This function controls the flow of video conversion + and handles all necessary functions. + """ + + # Parse arguments for waifu2x + # Check argument sanity + self._check_model_type(self.model_type) + self._check_arguments() + self._check_memory() + + # Convert paths to absolute paths + 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): + raise FileNotFoundError(self.waifu2x_path) + + # Initialize objects for ffmpeg and waifu2x-caffe + fm = Ffmpeg(self.ffmpeg_path, self.video_output, self.ffmpeg_arguments) + w2 = Waifu2x(self.waifu2x_path, self.method, self.model_type) + + # Extract frames from video + fm.extract_frames(self.input_video, self.extracted_frames) + + Avalon.info('Reading video information') + info = self._get_video_info() + # Analyze original video with ffprobe and retrieve framerate + # width, height = info['streams'][0]['width'], info['streams'][0]['height'] + + # Find index of video stream + video_stream_index = None + for stream in info['streams']: + if stream['codec_type'] == 'video': + video_stream_index = stream['index'] + break + + # Exit if no video stream found + if video_stream_index is None: + Avalon.error('Aborting: No video stream found') + + # Get average frame rate of video stream + framerate = float(Fraction(info['streams'][video_stream_index]['avg_frame_rate'])) + Avalon.info('Framerate: {}'.format(framerate)) + + # Upscale images one by one using waifu2x + Avalon.info('Starting to upscale extracted images') + self._upscale_frames(w2) + Avalon.info('Upscaling completed') + + # Frames to Video + Avalon.info('Converting extracted frames into video') + + # Width/height will be coded width/height x upscale factor + if self.factor: + coded_width = info['streams'][video_stream_index]['coded_width'] + coded_height = info['streams'][video_stream_index]['coded_height'] + fm.convert_video(framerate, '{}x{}'.format(self.factor * coded_width, self.factor * coded_height), self.upscaled_frames) + + # Use user defined output size + else: + fm.convert_video(framerate, '{}x{}'.format(self.width, self.height), self.upscaled_frames) + Avalon.info('Conversion completed') + + # Extract and press audio in + Avalon.info('Stripping audio track from original video') + fm.extract_audio(self.input_video, self.upscaled_frames) + Avalon.info('Inserting audio track into new video') + fm.insert_audio_track(self.upscaled_frames) diff --git a/bin/video2x.json b/bin/video2x.json new file mode 100644 index 0000000..b78106b --- /dev/null +++ b/bin/video2x.json @@ -0,0 +1,6 @@ +{ + "waifu2x_path": "C:/Program Files (x86)/waifu2x-caffe/waifu2x-caffe-cui.exe", + "ffmpeg_path": "C:/Program Files (x86)/ffmpeg/bin/", + "ffmpeg_arguments": [], + "ffmpeg_hwaccel": "cuda" +} \ No newline at end of file diff --git a/bin/video2x.py b/bin/video2x.py new file mode 100755 index 0000000..6961783 --- /dev/null +++ b/bin/video2x.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" + +__ __ _ _ ___ __ __ +\ \ / / (_) | | |__ \ \ \ / / + \ \ / / _ __| | ___ ___ ) | \ V / + \ \/ / | | / _` | / _ \ / _ \ / / > < + \ / | | | (_| | | __/ | (_) | / /_ / . \ + \/ |_| \__,_| \___| \___/ |____| /_/ \_\ + + +Name: Video2x Controller +Author: K4YT3X +Date Created: Feb 24, 2018 +Last Modified: December 10, 2018 + +Licensed under the GNU General Public License Version 3 (GNU GPL v3), + available at: https://www.gnu.org/licenses/gpl-3.0.txt + +(C) 2018 K4YT3X + +Description: Video2X is an automation software based on +waifu2x image enlarging engine. It extracts frames from a +video, enlarge it by a number of times without losing any +details or quality, keeping lines smooth and edges sharp. +""" +from avalon_framework import Avalon +from upscaler import Upscaler +import argparse +import inspect +import json +import os +import time +import traceback + +VERSION = '2.1.7' + +EXEC_PATH = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) +FRAMES = '{}\\frames'.format(EXEC_PATH) # Folder containing extracted frames +UPSCALED = '{}\\upscaled'.format(EXEC_PATH) # Folder containing enlarges frames + + +def process_arguments(): + """Processes CLI arguments + + This function parses all arguments + This allows users to customize options + for the output video. + """ + parser = argparse.ArgumentParser() + + # Video options + basic_options = parser.add_argument_group('Basic Options') + basic_options.add_argument('-i', '--input', help='Specify source video file', action='store', default=False, required=True) + basic_options.add_argument('-o', '--output', help='Specify output video file', action='store', default=False, required=True) + basic_options.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb') + 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') + + # Scaling options + scaling_options = parser.add_argument_group('Scaling Options', required=True) + scaling_options.add_argument('--width', help='Output video width', action='store', type=int, default=False) + scaling_options.add_argument('--height', help='Output video height', action='store', type=int, default=False) + scaling_options.add_argument('-f', '--factor', help='Factor to upscale the videos by', action='store', type=int, default=False) + + # Render drivers, at least one option must be specified + driver_group = parser.add_argument_group('Render Drivers', required=True) + driver_group.add_argument('--cpu', help='Use CPU for enlarging', action='store_true', default=False) + driver_group.add_argument('--gpu', help='Use GPU for enlarging', action='store_true', default=False) + driver_group.add_argument('--cudnn', help='Use CUDNN for enlarging', action='store_true', default=False) + + # Parse arguments + return parser.parse_args() + + +def print_logo(): + print('__ __ _ _ ___ __ __') + print('\\ \\ / / (_) | | |__ \\ \\ \\ / /') + print(' \\ \\ / / _ __| | ___ ___ ) | \\ V /') + print(' \\ \\/ / | | / _` | / _ \\ / _ \\ / / > <') + print(' \\ / | | | (_| | | __/ | (_) | / /_ / . \\') + print(' \\/ |_| \\__,_| \\___| \\___/ |____| /_/ \\_\\') + print('\n Video2X Video Enlarger') + spaces = ((44 - len("Version " + VERSION)) // 2) * " " + print('{}\n{} Version {}\n{}'.format(Avalon.FM.BD, spaces, VERSION, Avalon.FM.RST)) + + +def read_config(config_file): + """ Reads configuration file + + Returns a dictionary read by json. + """ + with open(config_file, 'r') as raw_config: + config = json.load(raw_config) + return config + + +# /////////////////// Execution /////////////////// # + +# This is not a library +if __name__ != '__main__': + Avalon.error('This file cannot be imported') + exit(1) + +print_logo() + +# Process cli arguments +args = process_arguments() + +# Read configurations from JSON +config = read_config() +waifu2x_path = config['waifu2x_path'] +ffmpeg_path = config['ffmpeg_path'] +ffmpeg_arguments = config['ffmpeg_arguments'] +ffmpeg_hwaccel = config['ffmpeg_hwaccel'] + +# Parse arguments for waifu2x +if args.cpu: + method = 'cpu' +elif args.gpu: + method = 'gpu' + ffmpeg_arguments.append('-hwaccel {}'.format(ffmpeg_hwaccel)) +elif args.cudnn: + method = 'cudnn' + ffmpeg_arguments.append('-hwaccel {}'.format(ffmpeg_hwaccel)) + +# Start execution +try: + # Start timer + begin_time = time.time() + + # Initialize and run upscaling + upscaler = Upscaler(input_video=args.video, output_video=args.output, method=method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, factor=args.factor, model_type=args.model_type, threads=args.threads) + upscaler.run() + + Avalon.info('Program completed, taking {} seconds'.format(round((time.time() - begin_time), 5))) +except Exception: + Avalon.error('An exception occurred') + traceback.print_exc() diff --git a/bin/video2x_setup.py b/bin/video2x_setup.py new file mode 100644 index 0000000..fcd745f --- /dev/null +++ b/bin/video2x_setup.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Name: Video2X Setup Script +Author: K4YT3X +Date Created: November 28, 2018 +Last Modified: November 29, 2018 + +Licensed under the GNU General Public License Version 3 (GNU GPL v3), + available at: https://www.gnu.org/licenses/gpl-3.0.txt + +(C) 2018 K4YT3X + +Description: This script helps installing all dependencies of video2x +and generates a configuration for it. + +Installation Details: +- ffmpeg: %LOCALAPPDATA%\\video2x\\ffmpeg +- waifu2x-caffe: %LOCALAPPDATA%\\video2x\\waifu2x-caffe + +""" +import json +import os +import subprocess +import tempfile +import traceback +import zipfile + +# Requests doesn't come with windows, therefore +# it will be installed as a dependency and imported +# later in the script. +# import requests + +VERSION = '1.0.0' + + +class Video2xSetup: + """ Install dependencies for video2x video enlarger + + This library is meant to be executed as a stand-alone + script. All files will be installed under %LOCALAPPDATA%\\video2x. + """ + + def __init__(self): + self.trash = [] + + def run(self): + + print('\nInstalling Python libraries') + self._install_python_requirements() + + print('\nInstalling FFMPEG') + self._install_ffmpeg() + + print('\nInstalling waifu2x-caffe') + self._install_waifu2x_caffe() + + print('\nGenerating Video2X configuration file') + self._generate_config() + + print('\nCleaning up temporary files') + self._cleanup() + + def _install_python_requirements(self): + """ Read requirements.txt and return its content + """ + with open('requirements.txt', 'r') as req: + for line in req: + package = line.split(' ')[0] + pip_install(package) + + def _cleanup(self): + """ Cleanup all the temp files downloaded + """ + for file in self.trash: + try: + print('Deleting: {}'.format(file)) + os.remove(file) + except FileNotFoundError: + pass + + def _install_ffmpeg(self): + """ Install FFMPEG + """ + latest_release = 'https://ffmpeg.zeranoe.com/builds/win64/static/ffmpeg-4.1-win64-static.zip' + + ffmpeg_zip = download(latest_release, tempfile.gettempdir()) + self.trash.append(ffmpeg_zip) + + with zipfile.ZipFile(ffmpeg_zip) as zipf: + zipf.extractall('{}\\video2x'.format(os.getenv('localappdata'))) + + def _install_waifu2x_caffe(self): + """ Install waifu2x_caffe + """ + import requests + + # Get latest release of waifu2x-caffe via GitHub API + latest_release = json.loads(requests.get('https://api.github.com/repos/lltcggie/waifu2x-caffe/releases/latest').content) + + for a in latest_release['assets']: + if 'waifu2x-caffe.zip' in a['browser_download_url']: + waifu2x_caffe_zip = download(a['browser_download_url'], tempfile.gettempdir()) + self.trash.append(waifu2x_caffe_zip) + + with zipfile.ZipFile(waifu2x_caffe_zip) as zipf: + zipf.extractall('{}\\video2x'.format(os.getenv('localappdata'))) + + def _generate_config(self): + """ Generate video2x config + """ + settings = {} + + settings['waifu2x_path'] = '{}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe'.format(os.getenv('localappdata')) + settings['ffmpeg_path'] = '{}\\video2x\\ffmpeg-4.1-win64-static\\bin'.format(os.getenv('localappdata')) + settings['ffmpeg_arguments'] = '' + settings['ffmpeg_hwaccel'] = 'cuda' + + with open('video2x.json', 'w') as config: + json.dump(settings, config, indent=2) + config.close() + + +def download(url, save_path, chunk_size=4096): + """ Download file to local with requests library + """ + import requests + + output_file = '{}\\{}'.format(save_path, url.split('/')[-1]) + print('Downloading: {}'.format(url)) + print('Chunk size: {}'.format(chunk_size)) + print('Saving to: {}'.format(output_file)) + + stream = requests.get(url, stream=True) + + # Write content into file + with open(output_file, 'wb') as output: + for chunk in stream.iter_content(chunk_size=chunk_size): + if chunk: + print('!', end='') + output.write(chunk) + print() + + return output_file + + +def pip_install(package): + """ Install python package via python pip module + + pip.main() is not available after pip 9.0.1, thus + pip module is not used in this case. + """ + return subprocess.run(['python', '-m', 'pip', 'install', '-U', package]).returncode + + +if __name__ == "__main__": + try: + print('Video2x Setup Script') + print('Version: {}'.format(VERSION)) + setup = Video2xSetup() + setup.run() + print('\n Script finished successfully') + except Exception: + traceback.print_exc() + print('An error has occurred') + print('Video2X Automatic Setup has failed') + exit(1) diff --git a/bin/waifu2x.py b/bin/waifu2x.py new file mode 100755 index 0000000..aea14ba --- /dev/null +++ b/bin/waifu2x.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Name: FFMPEG Class +Author: K4YT3X +Date Created: Feb 24, 2018 +Last Modified: October 23, 2018 + +Description: This class controls waifu2x +engine + +Version 2.0.4 +""" +from avalon_framework import Avalon +import subprocess +import threading + + +class Waifu2x: + """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, waifu2x_path, method, model_type): + self.waifu2x_path = waifu2x_path + self.method = method + self.model_type = model_type + self.print_lock = threading.Lock() + + def upscale(self, folderin, folderout, width, height): + """This is the core function for WAIFU2X class + + Arguments: + folderin {string} -- source folder path + folderout {string} -- output folder path + width {int} -- output video width + height {int} -- output video 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 = '\"{}\" -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) + subprocess.call(execute) + + # Print thread exiting message + self.print_lock.acquire() + Avalon.debug_info('[upscaler] Thread {} exiting'.format(threading.current_thread().name)) + self.print_lock.release()