diff --git a/bin/upscaler.py b/bin/upscaler.py index 6d80d1c..25bd0d5 100644 --- a/bin/upscaler.py +++ b/bin/upscaler.py @@ -4,7 +4,7 @@ Name: Video2X Upscaler Author: K4YT3X Date Created: December 10, 2018 -Last Modified: March 24, 2019 +Last Modified: March 30, 2019 Licensed under the GNU General Public License Version 3 (GNU GPL v3), available at: https://www.gnu.org/licenses/gpl-3.0.txt @@ -37,36 +37,31 @@ class Upscaler: ArgumentError -- if argument is not valid """ - 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_dir=None, threads=5, video2x_cache_folder='{}\\video2x'.format(tempfile.gettempdir()), preserve_frames=False): + def __init__(self, input_video, output_video, method, waifu2x_settings, ffmpeg_settings): # mandatory arguments self.input_video = input_video self.output_video = output_video self.method = method self.waifu2x_settings = waifu2x_settings self.ffmpeg_settings = ffmpeg_settings - self.waifu2x_driver = waifu2x_driver - - # check sanity of waifu2x_driver option - if waifu2x_driver != 'waifu2x_caffe' and waifu2x_driver != 'waifu2x_converter': - raise Exception('Unrecognized waifu2x driver: {}'.format(waifu2x_driver)) # optional arguments - self.scale_width = scale_width - self.scale_height = scale_height - self.scale_ratio = scale_ratio - self.model_dir = model_dir - self.threads = threads + self.waifu2x_driver = 'waifu2x_caffe' + self.scale_width = None + self.scale_height = None + self.scale_ratio = None + self.model_dir = None + self.threads = 5 + self.video2x_cache_folder = '{}\\video2x'.format(tempfile.gettempdir()) + self.image_format = 'png' + self.preserve_frames = False # create temporary folder/directories - self.video2x_cache_folder = video2x_cache_folder self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder) Avalon.debug_info('Extracted frames are being saved to: {}'.format(self.extracted_frames)) - self.upscaled_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder) Avalon.debug_info('Upscaled frames are being saved to: {}'.format(self.upscaled_frames)) - self.preserve_frames = preserve_frames - def cleanup(self): # delete temp directories when done # avalon framework cannot be used if python is shutting down @@ -153,9 +148,9 @@ 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.scale_ratio, self.threads, self.upscaler_exceptions) + w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.image_format, self.upscaler_exceptions) 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) + renamed = re.sub('_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{}'.format(self.image_format), '.{}'.format(self.image_format), image) shutil.move('{}\\{}'.format(self.upscaled_frames, image), '{}\\{}'.format(self.upscaled_frames, renamed)) self.progress_bar_exit_signal = True @@ -202,9 +197,9 @@ class Upscaler: for thread_info in thread_pool: # create thread if self.scale_ratio: - thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.upscaler_exceptions)) + thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.image_format, self.upscaler_exceptions)) else: - thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, self.upscaler_exceptions)) + thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, self.image_format, self.upscaler_exceptions)) thread.name = thread_info[1] # add threads into the pool @@ -252,7 +247,7 @@ class Upscaler: self.output_video = os.path.abspath(self.output_video) # initialize objects for ffmpeg and waifu2x-caffe - fm = Ffmpeg(self.ffmpeg_settings) + fm = Ffmpeg(self.ffmpeg_settings, self.image_format) # initialize waifu2x driver if self.waifu2x_driver == 'waifu2x_caffe': diff --git a/bin/waifu2x_caffe.py b/bin/waifu2x_caffe.py index 7e2237e..f716422 100644 --- a/bin/waifu2x_caffe.py +++ b/bin/waifu2x_caffe.py @@ -4,7 +4,7 @@ Name: Waifu2x Caffe Driver Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: March 24, 2019 +Last Modified: March 30, 2019 Description: This class is a high-level wrapper for waifu2x-caffe. @@ -33,7 +33,7 @@ class Waifu2xCaffe: self.model_dir = model_dir self.print_lock = threading.Lock() - def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height, upscaler_exceptions): + def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height, image_format, upscaler_exceptions): """This is the core function for WAIFU2X class Arguments: @@ -54,6 +54,8 @@ class Waifu2xCaffe: self.waifu2x_settings['scale_width'] = scale_width self.waifu2x_settings['scale_height'] = scale_height + self.waifu2x_settings['output_extention'] = image_format + # print thread start message self.print_lock.acquire() Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) diff --git a/bin/waifu2x_converter.py b/bin/waifu2x_converter.py index 8e10e3a..97f8245 100644 --- a/bin/waifu2x_converter.py +++ b/bin/waifu2x_converter.py @@ -4,7 +4,7 @@ Name: Waifu2x Converter CPP Driver Author: K4YT3X Date Created: February 8, 2019 -Last Modified: March 24, 2019 +Last Modified: March 30, 2019 Description: This class is a high-level wrapper for waifu2x-converter-cpp. @@ -28,7 +28,7 @@ class Waifu2xConverter: self.waifu2x_settings['model_dir'] = model_dir self.print_lock = threading.Lock() - def upscale(self, input_folder, output_folder, scale_ratio, jobs, upscaler_exceptions): + def upscale(self, input_folder, output_folder, scale_ratio, jobs, image_format, upscaler_exceptions): """ Waifu2x Converter Driver Upscaler This method executes the upscaling of extracted frames. @@ -54,6 +54,7 @@ class Waifu2xConverter: self.waifu2x_settings['scale-ratio'] = scale_ratio self.waifu2x_settings['jobs'] = jobs + self.waifu2x_settings['output-format'] = image_format # models_rgb must be specified manually for waifu2x-converter-cpp # if it's not specified in the arguments, create automatically