"""OpenCV frame acquisition with explicit replay and reconnect states.""" import threading import time from dataclasses import dataclass from enum import Enum from pathlib import Path from typing import Any, Callable, Optional, Union import cv2 import numpy as np class _FrameGrabber(threading.Thread): """Continuously read a capture, keeping only the newest frame. Decoding runs here so a slow consumer never builds a backlog: it always gets the latest decoded frame and older frames are dropped. This bounds live RTSP latency to roughly one inference regardless of stream vs processing rate. """ def __init__(self, capture: Any) -> None: super().__init__(daemon=True) self._capture = capture self._lock = threading.Lock() self._latest: Optional[np.ndarray] = None self._failed = False self._stopped = threading.Event() def _pump(self) -> bool: try: success, image = self._capture.read() except Exception: # noqa: BLE001 - a released capture must not crash the thread success, image = False, None if not success or image is None: with self._lock: self._failed = True return False with self._lock: self._latest = image return True def run(self) -> None: while not self._stopped.is_set(): if not self._pump(): break def take_latest(self): with self._lock: return self._latest, self._failed def stop(self) -> None: self._stopped.set() class SourceStatus(str, Enum): CONNECTED = "connected" RETRYING = "retrying" ERROR = "error" EOF = "eof" CLOSED = "closed" class SourceMode(str, Enum): REPLAY = "replay" STREAM = "stream" @dataclass(frozen=True) class FramePacket: image: Optional[np.ndarray] timestamp_monotonic: float status: SourceStatus error: Optional[str] = None CaptureFactory = Callable[[str], Any] class VideoSource: """Read replay or RTSP frames without treating source failures as evidence.""" def __init__( self, source: Union[str, Path], mode: SourceMode, retry_initial_seconds: float = 1.0, retry_max_seconds: float = 16.0, capture_factory: Optional[CaptureFactory] = None, drop_stale: bool = False, ) -> None: if retry_initial_seconds <= 0: raise ValueError("retry_initial_seconds must be positive") if retry_max_seconds < retry_initial_seconds: raise ValueError("retry_max_seconds must not be smaller than retry_initial_seconds") if not isinstance(mode, SourceMode): raise ValueError("mode must be a SourceMode") self._source = str(source) self._mode = mode self._drop_stale = bool(drop_stale) and mode is SourceMode.STREAM self._grabber: Optional[_FrameGrabber] = None self._retry_initial_seconds = retry_initial_seconds self._retry_max_seconds = retry_max_seconds self._capture_factory = capture_factory or cv2.VideoCapture self._capture = None self._status = SourceStatus.RETRYING if mode is SourceMode.STREAM else SourceStatus.ERROR self._retry_delay_seconds = retry_initial_seconds self._next_retry_at = 0.0 self._last_timestamp: Optional[float] = None self._last_error: Optional[str] = None self._closed = False @property def status(self) -> SourceStatus: return self._status def close(self) -> None: self._stop_grabber() self._closed = True self._status = SourceStatus.CLOSED def read(self, now: Optional[float] = None) -> FramePacket: """Return one frame or an explicit non-frame state. ``now`` is injectable for deterministic reconnect tests. It is never interpreted as video time; replay timestamps come from OpenCV. """ timestamp = time.monotonic() if now is None else float(now) if self._closed: return self._packet(timestamp, SourceStatus.CLOSED, "source is closed") if self._drop_stale: return self._read_latest(timestamp) if self._mode is SourceMode.REPLAY and self._status is SourceStatus.EOF: return self._packet(timestamp, SourceStatus.EOF, self._last_error) if self._capture is None: if self._status is SourceStatus.RETRYING and timestamp < self._next_retry_at: return self._packet(timestamp, SourceStatus.RETRYING, self._last_error) if not self._open_capture(timestamp): return self._packet(timestamp, self._status, self._last_error) success, image = self._capture.read() if not success or image is None: self._release_capture() if self._mode is SourceMode.STREAM: self._schedule_retry(timestamp, "frame read failed; retry scheduled") return self._packet(timestamp, SourceStatus.RETRYING, self._last_error) self._status = SourceStatus.EOF self._last_error = "frame read reached end of source" return self._packet(timestamp, SourceStatus.EOF, self._last_error) self._status = SourceStatus.CONNECTED self._last_error = None self._retry_delay_seconds = self._retry_initial_seconds frame_timestamp = self._frame_timestamp(timestamp) return FramePacket(image=image, timestamp_monotonic=frame_timestamp, status=SourceStatus.CONNECTED) def _read_latest(self, timestamp: float) -> FramePacket: """STREAM path: return the newest grabbed frame, dropping any backlog.""" if self._capture is None: if self._status is SourceStatus.RETRYING and timestamp < self._next_retry_at: return self._packet(timestamp, SourceStatus.RETRYING, self._last_error) if not self._open_capture(timestamp): return self._packet(timestamp, self._status, self._last_error) self._grabber = _FrameGrabber(self._capture) self._grabber.start() image, failed = self._grabber.take_latest() if failed: self._stop_grabber() self._schedule_retry(timestamp, "stream read failed; retry scheduled") return self._packet(timestamp, SourceStatus.RETRYING, self._last_error) if image is None: return self._packet(timestamp, SourceStatus.RETRYING, "connecting") self._status = SourceStatus.CONNECTED self._last_error = None self._retry_delay_seconds = self._retry_initial_seconds self._last_timestamp = timestamp return FramePacket(image=image, timestamp_monotonic=timestamp, status=SourceStatus.CONNECTED) def _stop_grabber(self) -> None: if self._grabber is not None: self._grabber.stop() self._grabber = None self._release_capture() def _open_capture(self, now: float) -> bool: capture = self._capture_factory(self._source) if capture is None or not capture.isOpened(): if capture is not None: capture.release() if self._mode is SourceMode.STREAM: self._schedule_retry(now, "unable to open source; retry scheduled") else: self._status = SourceStatus.ERROR self._last_error = "unable to open source" return False self._capture = capture self._status = SourceStatus.CONNECTED return True def _frame_timestamp(self, fallback_now: float) -> float: if self._mode is SourceMode.STREAM: self._last_timestamp = fallback_now return fallback_now source_seconds = float(self._capture.get(cv2.CAP_PROP_POS_MSEC)) / 1000.0 frame_seconds = self._timestamp_from_frame_index() if source_seconds < 0: source_seconds = frame_seconds if frame_seconds is not None else fallback_now elif ( self._last_timestamp is not None and source_seconds <= self._last_timestamp and frame_seconds is not None and frame_seconds > self._last_timestamp ): source_seconds = frame_seconds if self._last_timestamp is not None and source_seconds <= self._last_timestamp: source_seconds = self._last_timestamp + 0.000001 self._last_timestamp = source_seconds return source_seconds def _timestamp_from_frame_index(self) -> Optional[float]: fps = float(self._capture.get(cv2.CAP_PROP_FPS)) frame_index = float(self._capture.get(cv2.CAP_PROP_POS_FRAMES)) if fps <= 0 or frame_index < 1: return None return (frame_index - 1.0) / fps def _schedule_retry(self, now: float, reason: str) -> None: self._release_capture() self._status = SourceStatus.RETRYING self._last_error = reason self._next_retry_at = now + self._retry_delay_seconds self._retry_delay_seconds = min( self._retry_delay_seconds * 2.0, self._retry_max_seconds ) def _release_capture(self) -> None: if self._capture is not None: self._capture.release() self._capture = None def _packet( self, timestamp: float, status: SourceStatus, error: Optional[str] ) -> FramePacket: return FramePacket( image=None, timestamp_monotonic=timestamp, status=status, error=error, )