54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
"""Annotated video writer — wraps OpenCV VideoWriter for output."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
class AnnotatedVideoWriter:
|
|
"""Writes annotated frames to an MP4 file.
|
|
|
|
Args:
|
|
output_path: Destination .mp4 file path.
|
|
fps: Frames per second for the output video.
|
|
frame_size: (width, height) tuple.
|
|
codec: FourCC codec string (default "mp4v").
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
output_path: str,
|
|
fps: float,
|
|
frame_size: tuple[int, int],
|
|
codec: str = "mp4v",
|
|
) -> None:
|
|
if fps <= 0:
|
|
raise ValueError(f"fps must be > 0, got {fps}")
|
|
self._path = Path(output_path)
|
|
self._path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
w, h = frame_size
|
|
fourcc = cv2.VideoWriter_fourcc(*codec)
|
|
self._writer = cv2.VideoWriter(str(self._path), fourcc, fps, (w, h))
|
|
self._frame_count = 0
|
|
|
|
if not self._writer.isOpened():
|
|
raise RuntimeError(f"Failed to open VideoWriter for {self._path}")
|
|
|
|
def write_frame(self, frame: np.ndarray) -> None:
|
|
"""Write one frame. Frame size must match constructor frame_size."""
|
|
self._writer.write(frame)
|
|
self._frame_count += 1
|
|
|
|
def finish(self) -> None:
|
|
"""Release the writer. Idempotent — safe to call multiple times."""
|
|
if self._writer is not None and self._writer.isOpened():
|
|
self._writer.release()
|
|
|
|
@property
|
|
def frame_count(self) -> int:
|
|
return self._frame_count
|