replaced tqdm with rich.progress; enhanced error handling

video2x/decoder.py

@@ -19,13 +19,15 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
 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,19 +121,28 @@ class VideoDecoder(threading.Thread):
                     "RGB", (self.input_width, self.input_height), buffer
                 )
 
+                # 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
 
+                previous_image = image
                 frame_index += 1
 
             # most likely "not enough image data"
             except ValueError as e:
+
+                # ignore queue closed
+                if not "is closed" in str(e):
                     logger.exception(e)
                 break
 
@@ -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()

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()

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 <blue>{self.name}</blue> 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 <blue>{self.name}</blue> terminating"
+        )
         self.running = False
         return super().run()
 

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
+        # a temporary variable that stores the exception
-        if mode == "upscale":
+        exception = []
-            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}")
 
         try:
+            # 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
+                )
+
                 # 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()
+                    # show progress bar when upscale starts
-                time.sleep(0.1)
+                    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.n = self.processed.value
+                progress.update(task, completed=self.processed.value)
-            progress.refresh()
 
         # 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)