From dc7d95bbcee1afbe7f70f28827bd16769de30c22 Mon Sep 17 00:00:00 2001
From: Ruben van de Ven <git@rubenvandeven.com>
Date: Tue, 30 Apr 2024 11:47:03 +0200
Subject: [PATCH] Rendered video now accepts non-constant frame interval,
 storing frames in a tmp dir

---
 trap/renderer.py | 45 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 43 insertions(+), 2 deletions(-)

diff --git a/trap/renderer.py b/trap/renderer.py
index e2c531f..bcbe221 100644
--- a/trap/renderer.py
+++ b/trap/renderer.py
@@ -9,12 +9,51 @@ 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: BaseEvent):
         self.config = config
@@ -50,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)
@@ -138,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: