Rendered video now accepts non-constant frame interval, storing frames in a tmp dir

trap/renderer.py

@@ -4,18 +4,58 @@ from argparse import Namespace
 import datetime
 import logging
 from multiprocessing import Event
+from multiprocessing.synchronize import Event as BaseEvent
 import cv2
 import numpy as np
 
 import zmq
+import tempfile
+from pathlib import Path
+import shutil
 
 from trap.frame_emitter import DetectionState, Frame
 
 
 logger = logging.getLogger("trap.renderer")
 
+class FrameWriter:
+    """
+    Drop-in compatible interface with cv2.VideoWriter, but support variable
+    framerate.
+    See https://video.stackexchange.com/questions/25811/ffmpeg-make-video-with-non-constant-framerate-from-image-filenames
+    """
+    def __init__(self, filename: str, fps: float, frame_size: tuple) -> None:
+        self.filename = filename
+        self.fps = fps
+        self.frame_size = frame_size
+
+        self.tmp_dir = tempfile.TemporaryDirectory(prefix="trap-output-")
+
+        self.i = 0
+    
+    def write(self, img: cv2.typing.MatLike):
+        self.i += 1
+        cv2.imwrite(self.tmp_dir.name + f"/{self.i:07d}.png", img)
+
+    def release(self):
+        """Actually write the video"""
+        # ffmpeg -f image2 -ts_from_file 2 -i %d.png out.mp4
+        logger.info(f"Write frames from {self.tmp_dir.name} to {self.filename}")
+        (
+            ffmpeg
+            # the magic here is in --ts_from_file which uses the mtime of the file for the interval
+            # this makes it possible to have non-constant intervals between frames, which is usefull
+            # since we render frames when we get them
+            .input(self.tmp_dir.name + "/%07d.png", format="image2", ts_from_file=2)
+            .output(self.filename, framerate=self.fps)
+            .run()
+        )
+        logger.info(f"Rm frame directory: {self.tmp_dir.name}")
+        self.tmp_dir.cleanup()
+
+
 class Renderer:
-    def __init__(self, config: Namespace, is_running: Event):
+    def __init__(self, config: Namespace, is_running: BaseEvent):
         self.config = config
         self.is_running = is_running
 
@@ -49,6 +89,8 @@ class Renderer:
         filename = self.config.output_dir / f"render_predictions-{date_str}-{self.config.detector}.mp4"
         logger.info(f"Write to {filename}")
 
+        return FrameWriter(str(filename), self.fps, self.frame_size)
+
         fourcc = cv2.VideoWriter_fourcc(*'vp09')
 
         return cv2.VideoWriter(str(filename), fourcc, self.fps, self.frame_size)
@@ -137,9 +179,9 @@ def decorate_frame(frame: Frame, prediction_frame: Frame, first_time: float, con
 
     overlay = np.zeros(frame.img.shape, np.uint8)
     # Fill image with red color(set each pixel to red)
-    overlay[:] = (128, 0, 128)
+    overlay[:] = (130, 0, 75)
 
-    frame.img = cv2.addWeighted(frame.img, .5, overlay, .5, 0)
+    frame.img = cv2.addWeighted(frame.img, .4, overlay, .6, 0)
     img = frame.img
 
     # all not working:
@@ -237,6 +279,6 @@ def decorate_frame(frame: Frame, prediction_frame: Frame, first_time: float, con
     return img
 
     
-def run_renderer(config: Namespace, is_running: Event):
+def run_renderer(config: Namespace, is_running: BaseEvent):
     renderer = Renderer(config, is_running)
     renderer.run()