make script stop executing on waifu2x exception

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,11 +72,11 @@ class Upscaler:
         # avalon framework cannot be used if python is shutting down
         # therefore, plain print is used
         if not self.preserve_frames:
+
+            for directory in [self.extracted_frames, self.upscaled_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()
+                    print('Cleaning up cache directory: {}'.format())
+                    shutil.rmtree(directory)
                 except (OSError, FileNotFoundError):
                     pass
 
@@ -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
 

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,6 +43,7 @@ class Waifu2xCaffe:
             height {int} -- output video height
         """
 
+        try:
             # overwrite config file settings
             self.waifu2x_settings['input_path'] = input_folder
             self.waifu2x_settings['output_path'] = output_folder
@@ -86,3 +87,5 @@ class Waifu2xCaffe:
 
             # return command execution return code
             return completed_command.returncode
+        except Exception as e:
+            threads_exceptions.append(e)

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,6 +39,7 @@ class Waifu2xConverter:
             threads {int} -- number of threads
         """
 
+        try:
             # overwrite config file settings
             self.waifu2x_settings['input'] = input_folder
             self.waifu2x_settings['output'] = output_folder
@@ -91,3 +92,5 @@ class Waifu2xConverter:
             Avalon.debug_info('Executing: {}'.format(execute))
             return subprocess.run(execute, check=True).returncode
 
+        except Exception as e:
+            threads_exceptions.append(e)