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: