silver_pose/v1/video_source.py

"""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,
        )