make script stop executing on waifu2x exception

This commit is contained in:
k4yt3x 2019-03-24 23:04:50 -04:00
parent c59d113a11
commit 1d57c6652e
3 changed files with 102 additions and 91 deletions

View File

@ -59,12 +59,10 @@ class Upscaler:
# create temporary folder/directories # create temporary folder/directories
self.video2x_cache_folder = video2x_cache_folder self.video2x_cache_folder = video2x_cache_folder
self.extracted_frames_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder) self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder)
self.extracted_frames = self.extracted_frames_object.name
Avalon.debug_info('Extracted frames are being saved to: {}'.format(self.extracted_frames)) 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 = tempfile.mkdtemp(dir=self.video2x_cache_folder)
self.upscaled_frames = self.upscaled_frames_object.name
Avalon.debug_info('Upscaled frames are being saved to: {}'.format(self.upscaled_frames)) Avalon.debug_info('Upscaled frames are being saved to: {}'.format(self.upscaled_frames))
self.preserve_frames = preserve_frames self.preserve_frames = preserve_frames
@ -74,11 +72,11 @@ class Upscaler:
# avalon framework cannot be used if python is shutting down # avalon framework cannot be used if python is shutting down
# therefore, plain print is used # therefore, plain print is used
if not self.preserve_frames: if not self.preserve_frames:
for directory in [self.extracted_frames, self.upscaled_frames]:
try: try:
print('Cleaning up cache directory: {}'.format(self.extracted_frames)) print('Cleaning up cache directory: {}'.format())
self.extracted_frames_object.cleanup() shutil.rmtree(directory)
print('Cleaning up cache directory: {}'.format(self.upscaled_frames))
self.upscaled_frames_object.cleanup()
except (OSError, FileNotFoundError): except (OSError, FileNotFoundError):
pass pass
@ -163,6 +161,10 @@ class Upscaler:
# create a container for all upscaler threads # create a container for all upscaler threads
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 # 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))] 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: for thread_info in thread_pool:
# create thread # create thread
if self.scale_ratio: 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: 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] thread.name = thread_info[1]
# add threads into the pool # add threads into the pool
@ -231,6 +233,9 @@ class Upscaler:
self.progress_bar_exit_signal = True self.progress_bar_exit_signal = True
if len(self.threads_exceptions) != 0:
raise(self.threads_exceptions[0])
def run(self): def run(self):
"""Main controller for Video2X """Main controller for Video2X

View File

@ -4,7 +4,7 @@
Name: Waifu2x Caffe Driver Name: Waifu2x Caffe Driver
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 Date Created: Feb 24, 2018
Last Modified: March 19, 2019 Last Modified: March 24, 2019
Description: This class is a high-level wrapper Description: This class is a high-level wrapper
for waifu2x-caffe. for waifu2x-caffe.
@ -33,7 +33,7 @@ class Waifu2xCaffe:
self.model_dir = model_dir self.model_dir = model_dir
self.print_lock = threading.Lock() 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 """This is the core function for WAIFU2X class
Arguments: Arguments:
@ -43,6 +43,7 @@ class Waifu2xCaffe:
height {int} -- output video height height {int} -- output video height
""" """
try:
# overwrite config file settings # overwrite config file settings
self.waifu2x_settings['input_path'] = input_folder self.waifu2x_settings['input_path'] = input_folder
self.waifu2x_settings['output_path'] = output_folder self.waifu2x_settings['output_path'] = output_folder
@ -86,3 +87,5 @@ class Waifu2xCaffe:
# return command execution return code # return command execution return code
return completed_command.returncode return completed_command.returncode
except Exception as e:
threads_exceptions.append(e)

View File

@ -4,7 +4,7 @@
Name: Waifu2x Converter CPP Driver Name: Waifu2x Converter CPP Driver
Author: K4YT3X Author: K4YT3X
Date Created: February 8, 2019 Date Created: February 8, 2019
Last Modified: March 19, 2019 Last Modified: March 24, 2019
Description: This class is a high-level wrapper Description: This class is a high-level wrapper
for waifu2x-converter-cpp. for waifu2x-converter-cpp.
@ -28,7 +28,7 @@ class Waifu2xConverter:
self.waifu2x_settings['model_dir'] = model_dir self.waifu2x_settings['model_dir'] = model_dir
self.print_lock = threading.Lock() 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 """ Waifu2x Converter Driver Upscaler
This method executes the upscaling of extracted frames. This method executes the upscaling of extracted frames.
@ -39,6 +39,7 @@ class Waifu2xConverter:
threads {int} -- number of threads threads {int} -- number of threads
""" """
try:
# overwrite config file settings # overwrite config file settings
self.waifu2x_settings['input'] = input_folder self.waifu2x_settings['input'] = input_folder
self.waifu2x_settings['output'] = output_folder self.waifu2x_settings['output'] = output_folder
@ -91,3 +92,5 @@ class Waifu2xConverter:
Avalon.debug_info('Executing: {}'.format(execute)) Avalon.debug_info('Executing: {}'.format(execute))
return subprocess.run(execute, check=True).returncode return subprocess.run(execute, check=True).returncode
except Exception as e:
threads_exceptions.append(e)