src/fpgm/segmentation/base.py

"""Abstract interface for video segmentation backends."""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Iterator

import numpy as np

from fpgm.types import FrameMasks


class VideoSegmenter(ABC):
    """Stateful promptable video segmentation.

    Lifecycle: :meth:`start_session` -> :meth:`add_text_prompt` /
    :meth:`add_point_prompt` -> :meth:`propagate` -> :meth:`close_session`.
    Implementations must be usable as a context manager so sessions are always
    released, which matters on a shared GPU.
    """

    @abstractmethod
    def start_session(self, video_path: str) -> str:
        """Open a session on a JPEG folder or mp4; returns a session id."""

    @abstractmethod
    def add_text_prompt(self, frame_idx: int, text: str) -> FrameMasks:
        """Prompt with a natural-language phrase; returns that frame's masks."""

    @abstractmethod
    def add_point_prompt(
        self,
        frame_idx: int,
        points_xy: np.ndarray,
        labels: np.ndarray,
        obj_id: int | None = None,
    ) -> FrameMasks:
        """Prompt with positive/negative clicks in absolute pixel coordinates."""

    @abstractmethod
    def propagate(self) -> Iterator[FrameMasks]:
        """Stream per-frame masks across the whole video."""

    @abstractmethod
    def reset_session(self) -> None:
        """Clear prompts. Required before switching to a different text prompt."""

    @abstractmethod
    def close_session(self) -> None:
        """Release the session and its GPU memory."""

    def __enter__(self) -> "VideoSegmenter":
        return self

    def __exit__(self, *exc_info: object) -> None:
        self.close_session()