diff --git a/waifu2x.py b/waifu2x.py index 53507ab..97aecee 100644 --- a/waifu2x.py +++ b/waifu2x.py @@ -4,15 +4,14 @@ Name: FFMPEG Class Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: Feb 25, 2018 +Last Modified: May 19, 2018 Description: This class controls waifu2x engine -Version 1.1 +Version 2.0 """ - -import os +import subprocess class WAIFU2X: @@ -22,24 +21,27 @@ class WAIFU2X: 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. - - TODO: Make enhancement model customizable """ - def __init__(self, waifu2x_path, method): + def __init__(self, waifu2x_path, method, model_type): self.waifu2x_path = waifu2x_path self.method = method + self.model_type = model_type - def upscale(self, folderin, folderout, width, height, model_type): + def upscale(self, file, upscaled, width, height): """This is the core function for WAIFU2X class [description] Arguments: - folderin {string} -- source folder path - folderout {string} -- output folder path + file {string} -- input image + upscaled {string} -- output folder path width {int} -- output video width height {int} -- output video height + model_type {string} -- model to use for upscaling """ - os.system("{} -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, model_type)) + file_id = file.split('extracted_')[-1].split('.png')[0] + output_file = '{}\\{}{}{}'.format(upscaled, 'extracted_', file_id, '.png') + execute = "{} -p {} -I png -i {} -e png -o {} -w {} -h {} -n 3 -m noise_scale -y {}".format( + self.waifu2x_path, self.method, file, output_file, width, height, self.model_type) + subprocess.call(execute, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)