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)