From cc01f2d8e24aaa8a8459505162f99ffb2343a5b2 Mon Sep 17 00:00:00 2001 From: k4yt3x Date: Thu, 25 May 2023 18:44:54 +0000 Subject: [PATCH] refactored interpolator --- video2x/interpolator.py | 109 ++++++++++++++-------------------------- video2x/video2x.py | 9 ++-- 2 files changed, 44 insertions(+), 74 deletions(-) diff --git a/video2x/interpolator.py b/video2x/interpolator.py index f414962..8404957 100755 --- a/video2x/interpolator.py +++ b/video2x/interpolator.py @@ -22,105 +22,74 @@ Date Created: May 27, 2021 Last Modified: March 20, 2022 """ -import multiprocessing -import queue -import signal import time -from multiprocessing.managers import ListProxy -from multiprocessing.sharedctypes import Synchronized from loguru import logger from PIL import ImageChops, ImageStat from rife_ncnn_vulkan_python.rife_ncnn_vulkan import Rife -ALGORITHM_CLASSES = {"rife": Rife} +from .processor import Processor -class Interpolator(multiprocessing.Process): - def __init__( - self, - processing_queue: multiprocessing.Queue, - processed_frames: ListProxy, - pause: Synchronized, - ) -> None: - multiprocessing.Process.__init__(self) - self.processing_queue = processing_queue - self.processed_frames = processed_frames - self.pause = pause +class Interpolator: + ALGORITHM_CLASSES = {"rife": Rife} - self.running = False - self.processor_objects = {} + processor_objects = {} - signal.signal(signal.SIGTERM, self._stop) - - def run(self) -> None: - self.running = True - logger.opt(colors=True).info( - f"Interpolator process {self.name} initiating" + def interpolate_image(self, image0, image1, difference_threshold, algorithm): + difference = ImageChops.difference(image0, image1) + difference_stat = ImageStat.Stat(difference) + difference_ratio = ( + sum(difference_stat.mean) / (len(difference_stat.mean) * 255) * 100 ) - while self.running is True: + + if difference_ratio < difference_threshold: + processor_object = self.processor_objects.get(algorithm) + if processor_object is None: + processor_object = self.ALGORITHM_CLASSES[algorithm](0) + self.processor_objects[algorithm] = processor_object + interpolated_image = processor_object.process(image0, image1) + + else: + interpolated_image = image0 + + return interpolated_image + + +class InterpolatorProcessor(Processor, Interpolator): + def process(self) -> None: + task = self.tasks_queue.get() + while task is not None: try: - # pause if pause flag is set - if self.pause.value is True: + if self.pause_flag.value is True: time.sleep(0.1) continue - try: - # get new job from queue - ( - frame_index, - (image0, image1), - (difference_threshold, algorithm), - ) = self.processing_queue.get(False) - except queue.Empty: - time.sleep(0.1) - continue + ( + frame_index, + image0, + image1, + (difference_threshold, algorithm), + ) = task - # if image0 is None, image1 is the first frame - # skip this round if image0 is None: + task = self.tasks_queue.get() continue - # calculate the %diff between the current frame and the previous frame - difference = ImageChops.difference(image0, image1) - difference_stat = ImageStat.Stat(difference) - difference_ratio = ( - sum(difference_stat.mean) / (len(difference_stat.mean) * 255) * 100 + interpolated_image = self.interpolate_image( + image0, image1, difference_threshold, algorithm ) - # if the difference is lower than threshold - # process the interpolation - if difference_ratio < difference_threshold: - # select a processor object with the required settings - # create a new object if none are available - processor_object = self.processor_objects.get(algorithm) - if processor_object is None: - processor_object = ALGORITHM_CLASSES[algorithm](0) - self.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 - else: - interpolated_image = image0 - if frame_index == 1: self.processed_frames[0] = image0 self.processed_frames[frame_index * 2 - 1] = interpolated_image self.processed_frames[frame_index * 2] = image1 - # send exceptions into the client connection pipe + task = self.tasks_queue.get() + except (SystemExit, KeyboardInterrupt): break except Exception as error: logger.exception(error) break - - logger.opt(colors=True).info( - f"Interpolator process {self.name} terminating" - ) - return super().run() - - def _stop(self, _signal_number, _frame) -> None: - self.running = False diff --git a/video2x/video2x.py b/video2x/video2x.py index 37efbee..5c7e9a2 100755 --- a/video2x/video2x.py +++ b/video2x/video2x.py @@ -69,7 +69,7 @@ from video2x.processor import Processor from . import __version__ from .decoder import VideoDecoder, VideoDecoderThread from .encoder import VideoEncoder -from .interpolator import Interpolator +from .interpolator import Interpolator, InterpolatorProcessor from .upscaler import Upscaler, UpscalerProcessor # for desktop environments only @@ -102,7 +102,7 @@ class ProcessingSpeedColumn(ProgressColumn): class ProcessingMode(Enum): UPSCALE = {"label": "Upscaling", "processor": UpscalerProcessor} - INTERPOLATE = {"label": "Interpolating", "processor": Interpolator} + INTERPOLATE = {"label": "Interpolating", "processor": InterpolatorProcessor} class Video2X: @@ -179,7 +179,7 @@ class Video2X: # elif mode == ProcessingMode.INTERPOLATE: else: standalone_processor: Any = Interpolator.ALGORITHM_CLASSES[ - processing_settings[2] + processing_settings[1] ] if getattr(standalone_processor, "process", None) is None: standalone_processor().process_video( @@ -222,7 +222,7 @@ class Video2X: logger.info("Starting video encoder") encoder = VideoEncoder( input_path, - frame_rate * 2 if mode == "interpolate" else frame_rate, + frame_rate * 2 if mode == ProcessingMode.INTERPOLATE else frame_rate, output_path, output_width, output_height, @@ -337,6 +337,7 @@ class Video2X: else: logger.info("Processing has completed") + logger.info("Writing video trailer") finally: # stop keyboard listener