silver_pose/v1/video_source.py

"""OpenCV frame acquisition with explicit replay and reconnect states."""

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 SourceStatus(str, Enum):
    CONNECTED = "connected"
    RETRYING = "retrying"
    ERROR = "error"
    EOF = "eof"
    CLOSED = "closed"


@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],
        reconnect: bool = True,
        retry_initial_seconds: float = 1.0,
        retry_max_seconds: float = 16.0,
        capture_factory: Optional[CaptureFactory] = None,
    ) -> 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")
        self._source = str(source)
        self._reconnect = reconnect
        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 reconnect 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._release_capture()
        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._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._reconnect:
                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 _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._reconnect:
                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:
        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,
        )