intial commit

src/chicken_counter/video_writer.py

@@ -0,0 +1,81 @@
+"""Write annotated output videos via GStreamer or OpenCV codecs."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import cv2
+
+
+def make_video_writer(
+    path: str,
+    frame_size: tuple[int, int],
+    fps: float,
+    *,
+    encoder: str = "auto",
+    output_bitrate_kbps: int = 4000,
+    codec_preference: list[str] | None = None,
+) -> cv2.VideoWriter:
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+    width, height = frame_size
+    bitrate_bps = max(1, output_bitrate_kbps) * 1000
+    codecs = codec_preference or ["avc1", "mp4v", "H264"]
+
+    if encoder in ("auto", "gstreamer"):
+        writer = _try_gstreamer_writer(path, width, height, fps, bitrate_bps)
+        if writer is not None:
+            print(f"[video] opened GStreamer hardware encoder (bitrate={output_bitrate_kbps} kbps)")
+            return writer
+        if encoder == "gstreamer":
+            raise RuntimeError(
+                f"GStreamer video writer failed for path: {path}. "
+                "Ensure OpenCV was built with GStreamer and Jetson encoder plugins are available."
+            )
+
+    writer = _try_opencv_writer(path, frame_size, fps, codecs)
+    if writer is not None:
+        print(f"[video] opened OpenCV encoder (codecs tried: {codecs})")
+        return writer
+
+    raise RuntimeError(
+        "Unable to open video writer for path: "
+        f"{path}. Tried encoder={encoder}, codecs={codecs}"
+    )
+
+
+def _try_gstreamer_writer(
+    path: str,
+    width: int,
+    height: int,
+    fps: float,
+    bitrate_bps: int,
+) -> cv2.VideoWriter | None:
+    fps_int = max(1, int(round(fps)))
+    pipeline = (
+        f"appsrc ! video/x-raw, format=BGR ! "
+        f"video/x-raw,width={width},height={height},framerate={fps_int}/1 ! "
+        f"videoconvert ! nvvidconv ! "
+        f"video/x-raw(memory:NVMM),format=NV12 ! "
+        f"nvv4l2h264enc bitrate={bitrate_bps} insert-sps-pps=true ! "
+        f"h264parse ! mp4mux ! filesink location={path}"
+    )
+    writer = cv2.VideoWriter(pipeline, cv2.CAP_GSTREAMER, 0, fps, (width, height), True)
+    if writer.isOpened():
+        return writer
+    writer.release()
+    return None
+
+
+def _try_opencv_writer(
+    path: str,
+    frame_size: tuple[int, int],
+    fps: float,
+    codec_candidates: list[str],
+) -> cv2.VideoWriter | None:
+    for codec in codec_candidates:
+        fourcc = cv2.VideoWriter_fourcc(*codec)
+        writer = cv2.VideoWriter(path, fourcc, fps, frame_size)
+        if writer.isOpened():
+            return writer
+        writer.release()
+    return None