patch for potential problem with upscaler initialization

This commit is contained in:
k4yt3x 2019-03-24 21:57:22 -04:00
parent 7fd8228a1b
commit f9432b0fe8

View File

@ -30,13 +30,12 @@ import time
class Upscaler: class Upscaler:
""" An instance of this class is a upscaler that will """ An instance of this class is a upscaler that will
upscale all images in the given folder. upscale all images in the given folder.
Raises: Raises:
Exception -- all exceptions Exception -- all exceptions
ArgumentError -- if argument is not valid 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, 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):
# mandatory arguments # mandatory arguments
self.input_video = input_video self.input_video = input_video
@ -152,7 +151,7 @@ class Upscaler:
for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]: 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+)?)\]\.png', '.png', image)
shutil.move('{}\\{}'.format(self.upscaled_frames, image), '{}\\{}'.format(self.upscaled_frames, renamed)) shutil.move('{}\\{}'.format(self.upscaled_frames, image), '{}\\{}'.format(self.upscaled_frames, renamed))
self.progress_bar_exit_signal = True self.progress_bar_exit_signal = True
progress_bar.join() progress_bar.join()
return return
@ -197,9 +196,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, False, self.scale_width, self.scale_height))
else:
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))
else:
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height))
thread.name = thread_info[1] thread.name = thread_info[1]
# add threads into the pool # add threads into the pool
@ -216,7 +215,7 @@ class Upscaler:
# wait for threads to finish # wait for threads to finish
for thread in upscaler_threads: for thread in upscaler_threads:
thread.join() thread.join()
self.progress_bar_exit_signal = True self.progress_bar_exit_signal = True
def run(self): def run(self):