forked from zakaria/chicken-counting-sukawarna-det
82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
"""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
|