From 7b600415298d32cf40756e16b52b77c61f07af44 Mon Sep 17 00:00:00 2001 From: k4yt3x Date: Sat, 12 Feb 2022 09:08:28 +0000 Subject: [PATCH] replaced tqdm with rich.progress; enhanced error handling --- video2x/decoder.py | 38 ++++++++++++----- video2x/encoder.py | 2 +- video2x/upscaler.py | 8 +++- video2x/video2x.py | 102 ++++++++++++++++++++++++++++++++------------ 4 files changed, 108 insertions(+), 42 deletions(-) diff --git a/video2x/decoder.py b/video2x/decoder.py index bbec441..4b4a0f0 100755 --- a/video2x/decoder.py +++ b/video2x/decoder.py @@ -19,13 +19,15 @@ along with this program. If not, see . Name: Video Decoder Author: K4YT3X Date Created: June 17, 2021 -Last Modified: June 17, 2021 +Last Modified: February 12, 2022 """ # built-in imports +import contextlib import os import pathlib import queue +import signal import subprocess import threading @@ -119,20 +121,29 @@ class VideoDecoder(threading.Thread): "RGB", (self.input_width, self.input_height), buffer ) - self.processing_queue.put( - ( - frame_index, - (previous_image, image), - self.processing_settings, - ) - ) - previous_image = image + # keep checking if the running flag is set to False + # while waiting to put the next image into the queue + while self.running: + with contextlib.suppress(queue.Full): + self.processing_queue.put( + ( + frame_index, + (previous_image, image), + self.processing_settings, + ), + timeout=0.1, + ) + break + previous_image = image frame_index += 1 # most likely "not enough image data" except ValueError as e: - logger.exception(e) + + # ignore queue closed + if not "is closed" in str(e): + logger.exception(e) break # send exceptions into the client connection pipe @@ -141,9 +152,14 @@ class VideoDecoder(threading.Thread): logger.exception(e) break + # send SIGINT (2) to FFmpeg + # this instructs it to finalize and exit + if self.decoder.poll() is None: + self.decoder.send_signal(signal.SIGTERM) + # ensure the decoder has exited self.decoder.wait() - logger.debug("Decoder thread exiting") + logger.info("Decoder thread exiting") self.running = False return super().run() diff --git a/video2x/encoder.py b/video2x/encoder.py index 5109c7e..925569f 100755 --- a/video2x/encoder.py +++ b/video2x/encoder.py @@ -156,7 +156,7 @@ class VideoEncoder(threading.Thread): # wait for process to terminate self.encoder.wait() - logger.debug("Encoder thread exiting") + logger.info("Encoder thread exiting") self.running = False return super().run() diff --git a/video2x/upscaler.py b/video2x/upscaler.py index 2c5a543..f6c91d6 100755 --- a/video2x/upscaler.py +++ b/video2x/upscaler.py @@ -66,7 +66,9 @@ class Upscaler(multiprocessing.Process): def run(self): self.running = True - logger.info(f"Upscaler process {self.name} initiating") + logger.opt(colors=True).info( + f"Upscaler process {self.name} initiating" + ) driver_objects = {} while self.running: try: @@ -185,7 +187,9 @@ class Upscaler(multiprocessing.Process): logger.exception(e) break - logger.info(f"Upscaler process {self.name} terminating") + logger.opt(colors=True).info( + f"Upscaler process {self.name} terminating" + ) self.running = False return super().run() diff --git a/video2x/video2x.py b/video2x/video2x.py index 89fa861..441e5bb 100755 --- a/video2x/video2x.py +++ b/video2x/video2x.py @@ -58,7 +58,15 @@ import time # third-party imports from loguru import logger from rich import print -from tqdm import tqdm +from rich.progress import ( + BarColumn, + Progress, + ProgressColumn, + Task, + TimeElapsedColumn, + TimeRemainingColumn, +) +from rich.text import Text import cv2 import ffmpeg @@ -87,6 +95,20 @@ DRIVER_FIXED_SCALING_RATIOS = { "realsr": [4], } +# progress bar labels for different modes +MODE_LABELS = {"upscale": "Upscaling", "interpolate": "Interpolating"} + + +class ProcessingSpeedColumn(ProgressColumn): + """Custom progress bar column that displays the processing speed""" + + def render(self, task: Task) -> Text: + speed = task.finished_speed or task.speed + return Text( + f"{round(speed, 2) if isinstance(speed, float) else '?'} FPS", + style="progress.data.speed", + ) + class Video2X: """ @@ -165,7 +187,7 @@ class Video2X: logger.info("Starting video encoder") self.encoder = VideoEncoder( input_path, - frame_rate, + frame_rate * 2 if mode == "interpolate" else frame_rate, output_path, output_width, output_height, @@ -183,39 +205,53 @@ class Video2X: process.start() self.processor_processes.append(process) - # set progress bar values based on mode - if mode == "upscale": - label = "UPSC" - - # in interpolate mode, the frame rate is doubled - elif mode == "interpolate": - frame_rate *= 2 - label = "INTERP" - - else: - raise ValueError(f"unknown mode {mode}") - - # create progress bar - progress = tqdm(total=total_frames, desc=f"{label} {input_path.name}") + # a temporary variable that stores the exception + exception = [] try: - # wait for jobs in queue to deplete - while self.encoder.is_alive() is True: - for process in self.processor_processes: - if not process.is_alive(): - raise Exception("process died unexpectedly") - progress.n = self.processed.value - progress.refresh() - time.sleep(0.1) + # create progress bar + with Progress( + "[progress.description]{task.description}", + BarColumn(), + "[progress.percentage]{task.percentage:>3.0f}%", + "[color(240)]({task.completed}/{task.total})", + ProcessingSpeedColumn(), + TimeElapsedColumn(), + "<", + TimeRemainingColumn(), + disable=True, + ) as progress: + task = progress.add_task( + f"[cyan]{MODE_LABELS.get(mode, 'Unknown')}", total=total_frames + ) - logger.info("Encoding has completed") - progress.n = self.processed.value - progress.refresh() + # wait for jobs in queue to deplete + while self.encoder.is_alive() is True: + for process in self.processor_processes: + if not process.is_alive(): + raise Exception("process died unexpectedly") + + # show progress bar when upscale starts + if progress.disable is True and self.processed.value > 0: + progress.disable = False + progress.start() + + # update progress + progress.update(task, completed=self.processed.value) + time.sleep(0.5) + + logger.info("Encoding has completed") + progress.update(task, completed=self.processed.value) # if SIGTERM is received or ^C is pressed # TODO: pause and continue here - except (SystemExit, KeyboardInterrupt): + except (SystemExit, KeyboardInterrupt) as e: logger.warning("Exit signal received, terminating") + exception.append(e) + + except Exception as e: + logger.exception(e) + exception.append(e) finally: # mark processing queue as closed @@ -236,6 +272,10 @@ class Video2X: self.decoder.join() self.encoder.join() + # raise the error if there is any + if len(exception) > 0: + raise exception[0] + def upscale( self, input_path: pathlib.Path, @@ -468,5 +508,11 @@ def main(): args.threshold, args.driver, ) + + # don't print the traceback for manual terminations + except (SystemExit, KeyboardInterrupt) as e: + raise SystemExit(e) + except Exception as e: logger.exception(e) + raise SystemExit(e)