diff --git a/bin/upscaler.py b/bin/upscaler.py index befbd0c..29b0e51 100644 --- a/bin/upscaler.py +++ b/bin/upscaler.py @@ -142,6 +142,10 @@ class Upscaler: # progress bar thread exit signal self.progress_bar_exit_signal = False + # create a container for exceptions in threads + # if this thread is not empty, then an exception has occured + self.upscaler_exceptions = [] + # it's easier to do multi-threading with waifu2x_converter # the number of threads can be passed directly to waifu2x_converter if self.waifu2x_driver == 'waifu2x_converter': @@ -149,7 +153,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.scale_ratio, self.threads) + w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, 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) shutil.move('{}\\{}'.format(self.upscaled_frames, image), '{}\\{}'.format(self.upscaled_frames, renamed)) @@ -161,10 +165,6 @@ 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))] @@ -202,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, self.threads_exceptions)) + thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, 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.threads_exceptions)) + thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, self.upscaler_exceptions)) thread.name = thread_info[1] # add threads into the pool @@ -233,8 +233,8 @@ class Upscaler: self.progress_bar_exit_signal = True - if len(self.threads_exceptions) != 0: - raise(self.threads_exceptions[0]) + if len(self.upscaler_exceptions) != 0: + raise(self.upscaler_exceptions[0]) def run(self): """Main controller for Video2X diff --git a/bin/waifu2x_caffe.py b/bin/waifu2x_caffe.py index 43487ac..7e2237e 100644 --- a/bin/waifu2x_caffe.py +++ b/bin/waifu2x_caffe.py @@ -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, threads_exceptions): + def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height, upscaler_exceptions): """This is the core function for WAIFU2X class Arguments: @@ -88,4 +88,4 @@ class Waifu2xCaffe: # return command execution return code return completed_command.returncode except Exception as e: - threads_exceptions.append(e) + upscaler_exceptions.append(e) diff --git a/bin/waifu2x_converter.py b/bin/waifu2x_converter.py index ec78325..36d14a8 100644 --- a/bin/waifu2x_converter.py +++ b/bin/waifu2x_converter.py @@ -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, threads_exceptions): + def upscale(self, input_folder, output_folder, scale_ratio, jobs, upscaler_exceptions): """ Waifu2x Converter Driver Upscaler This method executes the upscaling of extracted frames. @@ -93,4 +93,4 @@ class Waifu2xConverter: return subprocess.run(execute, check=True).returncode except Exception as e: - threads_exceptions.append(e) + upscaler_exceptions.append(e)