diff --git a/video2x/interpolator.py b/video2x/interpolator.py index 5800f34..193b83f 100755 --- a/video2x/interpolator.py +++ b/video2x/interpolator.py @@ -38,7 +38,7 @@ from PIL import ImageChops, ImageStat from loguru import logger -DRIVER_CLASSES = {"rife": Rife} +ALGORITHM_CLASSES = {"rife": Rife} class Interpolator(multiprocessing.Process): @@ -57,7 +57,7 @@ class Interpolator(multiprocessing.Process): def run(self): self.running = True logger.info(f"Interpolator process {self.name} initiating") - driver_objects = {} + processor_objects = {} while self.running: try: try: @@ -65,7 +65,7 @@ class Interpolator(multiprocessing.Process): ( frame_index, (image0, image1), - (difference_threshold, driver), + (difference_threshold, algorithm), ) = self.processing_queue.get(False) except queue.Empty: time.sleep(0.1) @@ -86,13 +86,13 @@ class Interpolator(multiprocessing.Process): # process the interpolation if difference_ratio < difference_threshold: - # select a driver object with the required settings + # select a processor object with the required settings # create a new object if none are available - driver_object = driver_objects.get(driver) - if driver_object is None: - driver_object = DRIVER_CLASSES[driver](0) - driver_objects[driver] = driver_object - interpolated_image = driver_object.process(image0, image1) + processor_object = processor_objects.get(algorithm) + if processor_object is None: + processor_object = ALGORITHM_CLASSES[algorithm](0) + processor_objects[algorithm] = processor_object + interpolated_image = processor_object.process(image0, image1) # if the difference is greater than threshold # there's a change in camera angle, ignore diff --git a/video2x/upscaler.py b/video2x/upscaler.py index f6c91d6..fad6670 100755 --- a/video2x/upscaler.py +++ b/video2x/upscaler.py @@ -40,15 +40,15 @@ import time from PIL import Image, ImageChops, ImageStat from loguru import logger -# fixed scaling ratios supported by the drivers +# fixed scaling ratios supported by the algorithms # that only support certain fixed scale ratios -DRIVER_FIXED_SCALING_RATIOS = { +ALGORITHM_FIXED_SCALING_RATIOS = { "waifu2x": [1, 2], "srmd": [2, 3, 4], "realsr": [4], } -DRIVER_CLASSES = {"waifu2x": Waifu2x, "srmd": Srmd, "realsr": Realsr} +ALGORITHM_CLASSES = {"waifu2x": Waifu2x, "srmd": Srmd, "realsr": Realsr} class Upscaler(multiprocessing.Process): @@ -69,7 +69,7 @@ class Upscaler(multiprocessing.Process): logger.opt(colors=True).info( f"Upscaler process {self.name} initiating" ) - driver_objects = {} + processor_objects = {} while self.running: try: try: @@ -82,7 +82,7 @@ class Upscaler(multiprocessing.Process): output_height, noise, difference_threshold, - driver, + algorithm, ), ) = self.processing_queue.get(False) @@ -123,9 +123,9 @@ class Upscaler(multiprocessing.Process): # calculate required minimum scale ratio output_scale = max(output_width / width, output_height / height) - # select the optimal driver scaling ratio to use + # select the optimal algorithm scaling ratio to use supported_scaling_ratios = sorted( - DRIVER_FIXED_SCALING_RATIOS[driver] + ALGORITHM_FIXED_SCALING_RATIOS[algorithm] ) remaining_scaling_ratio = math.ceil(output_scale) @@ -163,17 +163,17 @@ class Upscaler(multiprocessing.Process): for job in scaling_jobs: - # select a driver object with the required settings + # select a processor object with the required settings # create a new object if none are available - driver_object = driver_objects.get((driver, job)) - if driver_object is None: - driver_object = DRIVER_CLASSES[driver]( + processor_object = processor_objects.get((algorithm, job)) + if processor_object is None: + processor_object = ALGORITHM_CLASSES[algorithm]( scale=job, noise=noise ) - driver_objects[(driver, job)] = driver_object + processor_objects[(algorithm, job)] = processor_object - # process the image with the selected driver - image1 = driver_object.process(image1) + # process the image with the selected algorithm + image1 = processor_object.process(image1) # downscale the image to the desired output size and save the image to disk image1 = image1.resize((output_width, output_height), Image.LANCZOS) diff --git a/video2x/video2x.py b/video2x/video2x.py index 941a48f..c7a6929 100755 --- a/video2x/video2x.py +++ b/video2x/video2x.py @@ -81,21 +81,15 @@ Contact: k4yt3x@k4yt3x.com""".format( __version__ ) -UPSCALING_DRIVERS = [ +# algorithms available for upscaling tasks +UPSCALING_ALGORITHMS = [ "waifu2x", "srmd", "realsr", ] -INTERPOLATION_DRIVERS = ["rife"] - -# fixed scaling ratios supported by the drivers -# that only support certain fixed scale ratios -DRIVER_FIXED_SCALING_RATIOS = { - "waifu2x": [1, 2], - "srmd": [2, 3, 4], - "realsr": [4], -} +# algorithms available for frame interpolation tasks +INTERPOLATION_ALGORITHMS = ["rife"] # progress bar labels for different modes MODE_LABELS = {"upscale": "Upscaling", "interpolate": "Interpolating"} @@ -279,6 +273,10 @@ class Video2X: logger.exception(e) exception.append(e) + # if no exceptions were produced + else: + logger.success("Processing completed successfully") + finally: # mark processing queue as closed self.processing_queue.close() @@ -319,7 +317,7 @@ class Video2X: noise: int, processes: int, threshold: float, - driver: str, + algorithm: str, ) -> None: # get basic video information @@ -354,7 +352,7 @@ class Video2X: output_height, noise, threshold, - driver, + algorithm, ), ) @@ -364,7 +362,7 @@ class Video2X: output_path: pathlib.Path, processes: int, threshold: float, - driver: str, + algorithm: str, ) -> None: # get video basic information @@ -386,7 +384,7 @@ class Video2X: Interpolator, "interpolate", processes, - (threshold, driver), + (threshold, algorithm), ) @@ -440,11 +438,11 @@ def parse_arguments() -> argparse.Namespace: upscale.add_argument("-h", "--height", type=int, help="output height") upscale.add_argument("-n", "--noise", type=int, help="denoise level", default=3) upscale.add_argument( - "-d", - "--driver", - choices=UPSCALING_DRIVERS, - help="driver to use for upscaling", - default=UPSCALING_DRIVERS[0], + "-a", + "--algorithm", + choices=UPSCALING_ALGORITHMS, + help="algorithm to use for upscaling", + default=UPSCALING_ALGORITHMS[0], ) upscale.add_argument( "-t", @@ -462,11 +460,11 @@ def parse_arguments() -> argparse.Namespace: "--help", action="help", help="show this help message and exit" ) interpolate.add_argument( - "-d", - "--driver", - choices=UPSCALING_DRIVERS, - help="driver to use for upscaling", - default=INTERPOLATION_DRIVERS[0], + "-a", + "--algorithm", + choices=UPSCALING_ALGORITHMS, + help="algorithm to use for upscaling", + default=INTERPOLATION_ALGORITHMS[0], ) interpolate.add_argument( "-t", @@ -523,7 +521,7 @@ def main(): args.noise, args.processes, args.threshold, - args.driver, + args.algorithm, ) elif args.action == "interpolate": @@ -532,11 +530,9 @@ def main(): args.output, args.processes, args.threshold, - args.driver, + args.algorithm, ) - logger.success("Processing completed successfully") - # don't print the traceback for manual terminations except (SystemExit, KeyboardInterrupt) as e: raise SystemExit(e)