diff --git a/bin/upscaler.py b/bin/upscaler.py index b74e1dc..ce6a6aa 100644 --- a/bin/upscaler.py +++ b/bin/upscaler.py @@ -59,12 +59,10 @@ class Upscaler: # create temporary folder/directories self.video2x_cache_folder = video2x_cache_folder - self.extracted_frames_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder) - self.extracted_frames = self.extracted_frames_object.name + 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_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder) - self.upscaled_frames = self.upscaled_frames_object.name + 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 @@ -74,13 +72,13 @@ class Upscaler: # avalon framework cannot be used if python is shutting down # therefore, plain print is used if not self.preserve_frames: - try: - print('Cleaning up cache directory: {}'.format(self.extracted_frames)) - self.extracted_frames_object.cleanup() - print('Cleaning up cache directory: {}'.format(self.upscaled_frames)) - self.upscaled_frames_object.cleanup() - except (OSError, FileNotFoundError): - pass + + for directory in [self.extracted_frames, self.upscaled_frames]: + try: + print('Cleaning up cache directory: {}'.format()) + shutil.rmtree(directory) + except (OSError, FileNotFoundError): + pass def _check_arguments(self): # check if arguments are valid / all necessary argument @@ -163,6 +161,10 @@ class Upscaler: # create a container for all upscaler threads upscaler_threads = [] + # create a container for exceptions in threads + # if this thread is not empty, then an exception has occured + self.threads_exceptions = [] + # 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))] @@ -200,9 +202,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)) + thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.threads_exceptions)) else: - thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height)) + thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, self.threads_exceptions)) thread.name = thread_info[1] # add threads into the pool @@ -231,6 +233,9 @@ class Upscaler: self.progress_bar_exit_signal = True + if len(self.threads_exceptions) != 0: + raise(self.threads_exceptions[0]) + def run(self): """Main controller for Video2X diff --git a/bin/waifu2x_caffe.py b/bin/waifu2x_caffe.py index 8403c02..43487ac 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 19, 2019 +Last Modified: March 24, 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): + def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height, threads_exceptions): """This is the core function for WAIFU2X class Arguments: @@ -43,46 +43,49 @@ class Waifu2xCaffe: height {int} -- output video height """ - # overwrite config file settings - self.waifu2x_settings['input_path'] = input_folder - self.waifu2x_settings['output_path'] = output_folder + try: + # 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 + 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() + # print thread start message + self.print_lock.acquire() + Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) + self.print_lock.release() - # list to be executed - execute = [] + # list to be executed + execute = [] - execute.append(self.waifu2x_settings['waifu2x_caffe_path']) - for key in self.waifu2x_settings.keys(): + execute.append(self.waifu2x_settings['waifu2x_caffe_path']) + for key in self.waifu2x_settings.keys(): - value = self.waifu2x_settings[key] + value = self.waifu2x_settings[key] - #is executable key or null or None means that leave this option out (keep default) - if key == 'waifu2x_caffe_path' or value is None or value is False: - continue - else: - if len(key) == 1: - execute.append('-{}'.format(key)) + # is executable key or null or None means that leave this option out (keep default) + if key == 'waifu2x_caffe_path' or value is None or value is False: + continue else: - execute.append('--{}'.format(key)) - execute.append(str(value)) - - Avalon.debug_info('Executing: {}'.format(execute)) - completed_command = subprocess.run(execute, check=True) + if len(key) == 1: + execute.append('-{}'.format(key)) + else: + execute.append('--{}'.format(key)) + execute.append(str(value)) - # print thread exiting message - self.print_lock.acquire() - Avalon.debug_info('[upscaler] Thread {} exiting'.format(threading.current_thread().name)) - self.print_lock.release() + Avalon.debug_info('Executing: {}'.format(execute)) + completed_command = subprocess.run(execute, check=True) - # return command execution return code - return completed_command.returncode + # print thread exiting message + self.print_lock.acquire() + Avalon.debug_info('[upscaler] Thread {} exiting'.format(threading.current_thread().name)) + self.print_lock.release() + + # return command execution return code + return completed_command.returncode + except Exception as e: + threads_exceptions.append(e) diff --git a/bin/waifu2x_converter.py b/bin/waifu2x_converter.py index 56f5f90..ec78325 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 19, 2019 +Last Modified: March 24, 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): + def upscale(self, input_folder, output_folder, scale_ratio, jobs, threads_exceptions): """ Waifu2x Converter Driver Upscaler This method executes the upscaling of extracted frames. @@ -39,55 +39,58 @@ class Waifu2xConverter: threads {int} -- number of threads """ - # overwrite config file settings - self.waifu2x_settings['input'] = input_folder - self.waifu2x_settings['output'] = output_folder + try: + # overwrite config file settings + self.waifu2x_settings['input'] = input_folder + self.waifu2x_settings['output'] = output_folder - # temporary fix for https://github.com/DeadSix27/waifu2x-converter-cpp/issues/109 - self.waifu2x_settings['i'] = input_folder - self.waifu2x_settings['o'] = output_folder - self.waifu2x_settings['input'] = None - self.waifu2x_settings['output'] = None + # temporary fix for https://github.com/DeadSix27/waifu2x-converter-cpp/issues/109 + self.waifu2x_settings['i'] = input_folder + self.waifu2x_settings['o'] = output_folder + self.waifu2x_settings['input'] = None + self.waifu2x_settings['output'] = None - self.waifu2x_settings['scale_ratio'] = scale_ratio - self.waifu2x_settings['jobs'] = jobs + self.waifu2x_settings['scale_ratio'] = scale_ratio + self.waifu2x_settings['jobs'] = jobs - # models_rgb must be specified manually for waifu2x-converter-cpp - # if it's not specified in the arguments, create automatically - if self.waifu2x_settings['model_dir'] is None: - self.waifu2x_settings['model_dir'] = '{}\\models_rgb'.format(self.waifu2x_settings['waifu2x_converter_path']) + # models_rgb must be specified manually for waifu2x-converter-cpp + # if it's not specified in the arguments, create automatically + if self.waifu2x_settings['model_dir'] is None: + 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() + # print thread start message + self.print_lock.acquire() + Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) + self.print_lock.release() - # list to be executed - execute = [] + # list to be executed + execute = [] - for key in self.waifu2x_settings.keys(): + for key in self.waifu2x_settings.keys(): - value = self.waifu2x_settings[key] + 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))) + # 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: + # 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)) - return subprocess.run(execute, check=True).returncode + # true means key is an option + if value is True: + continue + execute.append(str(value)) + + Avalon.debug_info('Executing: {}'.format(execute)) + return subprocess.run(execute, check=True).returncode + + except Exception as e: + threads_exceptions.append(e)