diff --git a/src/upscaler.py b/src/upscaler.py index d211b1d..594d79c 100755 --- a/src/upscaler.py +++ b/src/upscaler.py @@ -4,7 +4,7 @@ Name: Video2X Upscaler Author: K4YT3X Date Created: December 10, 2018 -Last Modified: May 8, 2020 +Last Modified: May 9, 2020 Description: This file contains the Upscaler class. Each instance of the Upscaler class is an upscaler on an image or @@ -257,7 +257,7 @@ class Upscaler: process_directory.mkdir(parents=True, exist_ok=True) # waifu2x-converter-cpp will perform multi-threading within its own process - if self.driver == 'waifu2x_converter_cpp': + if self.driver in ['waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan', 'srmd_ncnn_vulkan']: process_directories = [self.extracted_frames] else: @@ -297,13 +297,15 @@ class Upscaler: elif self.driver == 'waifu2x_ncnn_vulkan': self.process_pool.append(driver.upscale(process_directory, self.upscaled_frames, - self.scale_ratio)) + self.scale_ratio, + self.processes)) # if the driver being used is srmd_ncnn_vulkan elif self.driver == 'srmd_ncnn_vulkan': self.process_pool.append(driver.upscale(process_directory, self.upscaled_frames, - self.scale_ratio)) + self.scale_ratio, + self.processes)) # start progress bar in a different thread Avalon.debug_info(_('Starting progress monitor')) @@ -489,7 +491,7 @@ class Upscaler: self._wait() # get average frame rate of video stream - framerate = float(Fraction(video_info['streams'][video_stream_index]['avg_frame_rate'])) + framerate = float(Fraction(video_info['streams'][video_stream_index]['r_frame_rate'])) fm.pixel_format = video_info['streams'][video_stream_index]['pix_fmt'] # get a dict of all pixel formats and corresponding bit depth diff --git a/src/wrappers/srmd_ncnn_vulkan.py b/src/wrappers/srmd_ncnn_vulkan.py index 2be517d..9f64a47 100644 --- a/src/wrappers/srmd_ncnn_vulkan.py +++ b/src/wrappers/srmd_ncnn_vulkan.py @@ -4,7 +4,7 @@ Name: SRMD NCNN Vulkan Driver Creator: K4YT3X Date Created: April 26, 2020 -Last Modified: May 7, 2020 +Last Modified: May 9, 2020 Description: This class is a high-level wrapper for srmd_ncnn_vulkan. @@ -53,7 +53,7 @@ class WrapperMain: parser.add_argument('-x', action='store_true', help='enable tta mode') return parser.parse_args(arguments) - def upscale(self, input_directory, output_directory, scale_ratio): + def upscale(self, input_directory, output_directory, scale_ratio, threads): """This is the core function for SRMD ncnn Vulkan class Arguments: @@ -66,6 +66,7 @@ class WrapperMain: self.driver_settings['i'] = input_directory self.driver_settings['o'] = output_directory self.driver_settings['s'] = scale_ratio + self.driver_settings['j'] = '{}:{}:{}'.format(threads, threads, threads) # by default, srmd-ncnn-vulkan will look for the models under the current working directory # change the working directory to its containing folder if model directory not specified diff --git a/src/wrappers/waifu2x_ncnn_vulkan.py b/src/wrappers/waifu2x_ncnn_vulkan.py index a62bf21..36e0fe2 100644 --- a/src/wrappers/waifu2x_ncnn_vulkan.py +++ b/src/wrappers/waifu2x_ncnn_vulkan.py @@ -4,7 +4,7 @@ Name: Waifu2x NCNN Vulkan Driver Creator: SAT3LL Date Created: June 26, 2019 -Last Modified: May 7, 2020 +Last Modified: May 9, 2020 Editor: K4YT3X Last Modified: February 22, 2020 @@ -56,7 +56,7 @@ class WrapperMain: parser.add_argument('-x', action='store_true', help='enable tta mode') return parser.parse_args(arguments) - def upscale(self, input_directory, output_directory, scale_ratio): + def upscale(self, input_directory, output_directory, scale_ratio, threads): """This is the core function for WAIFU2X class Arguments: @@ -69,6 +69,7 @@ class WrapperMain: self.driver_settings['i'] = input_directory self.driver_settings['o'] = output_directory self.driver_settings['s'] = int(scale_ratio) + self.driver_settings['j'] = '{}:{}:{}'.format(threads, threads, threads) # by default, waifu2x-ncnn-vulkan will look for the models under the current working directory # change the working directory to its containing folder if model directory not specified