src/fpgm/segmentation/sam3.py

"""SAM 3.1 video segmentation adapter.

Wraps ``sam3.model_builder.build_sam3_multiplex_video_predictor`` and its
stateful ``handle_request``/``handle_stream_request`` API behind the
:class:`~fpgm.segmentation.base.VideoSegmenter` ABC.

Two SAM 3.1 behaviours drive most of the code here and are worth stating up
front because they are easy to get subtly wrong:

* Points passed to ``add_prompt`` are *relative* ``[0, 1]`` coordinates by
  default (``rel_coordinates=True``); our ABC contract is absolute pixels, so
  every point prompt is converted using the session's known frame resolution.
* Switching to a different text prompt on a session that already has one
  produces wrong results unless ``reset_session`` runs first. This adapter
  auto-resets (loudly, via a warning log) rather than ever risking a silent
  wrong answer.

This module also carries two runtime shims (applied in :meth:`Sam3VideoSegmenter._build`)
for known upstream bugs in the pinned ``third_party/sam3`` checkout, each documented
in detail on the shim itself: :func:`_patch_init_state_kwargs` (an ``init_state``
kwarg mismatch) and :func:`_patch_build_sam2_output` (point-prompt propagation
silently returning empty masks at every frame but the prompt frame). Both are
self-test-driven so they become no-ops automatically once sam3 is fixed.
"""

from __future__ import annotations

import functools
import inspect
import logging
import types
from collections.abc import Iterator
from pathlib import Path

import numpy as np

from fpgm.config import SegmentationConfig
from fpgm.segmentation.base import VideoSegmenter
from fpgm.tracking.sampling import mask_centroid
from fpgm.types import FrameMasks, Masklet

logger = logging.getLogger(__name__)


def _patch_init_state_kwargs(predictor: object) -> None:
    """Drop ``init_state`` kwargs the underlying model does not accept.

    Upstream incompatibility in sam3 (observed at commit ``6dbb02b``):
    ``Sam3BasePredictor.start_session`` *unconditionally* forwards
    ``offload_state_to_cpu`` to ``model.init_state``, but none of the SAM 3.1
    multiplex tracking classes accept that parameter. The result is that the
    documented ``start_session`` request path raises ``TypeError`` for SAM 3.1 no
    matter what the caller passes -- even the default ``False``.

    Rather than fork the vendored repo, the model's ``init_state`` is wrapped to
    discard kwargs outside its real signature. The wrap is signature-driven, so it
    becomes a no-op automatically once upstream adds the parameter.
    """
    model = getattr(predictor, "model", None)
    init_state = getattr(model, "init_state", None)
    if model is None or init_state is None:
        return

    try:
        accepted = set(inspect.signature(init_state).parameters)
    except (TypeError, ValueError):  # pragma: no cover - exotic callables
        return
    if "kwargs" in accepted or "offload_state_to_cpu" in accepted:
        return

    @functools.wraps(init_state)
    def _init_state(*args: object, **kwargs: object):
        dropped = [k for k in kwargs if k not in accepted]
        for key in dropped:
            kwargs.pop(key)
        if dropped:
            logger.debug("dropping unsupported init_state kwargs: %s", dropped)
        return init_state(*args, **kwargs)

    model.init_state = _init_state
    logger.info(
        "applied sam3 compatibility shim: init_state does not accept %s",
        sorted({"offload_state_to_cpu"} - accepted),
    )


def _patch_build_sam2_output(predictor: object) -> None:
    """Fix the bug that makes point-prompt video propagation return empty masks.

    Empirically diagnosed (see ``tests/test_sam3_point_prompt.py`` for the
    self-contained regression case): a point prompt followed by
    ``propagate_in_video`` -- *any* direction, with or without
    ``is_instance_processing`` -- returns a real mask at the prompt frame and an
    **empty** mask at every other frame. Text/visual-grounding prompting is
    unaffected. Direct reproduction against the pinned sam3 checkpoint (bypassing
    this adapter entirely, single point prompt on a static object over a 30-frame
    clip) showed the same pattern regardless of ``propagate(direction=...)`` or a
    fresh session with only one ``propagate_in_video`` call -- which rules out this
    adapter's own request construction and rules out the action-history
    ``propagation_fetch`` dispatch in
    ``Sam3MultiplexTracking.parse_action_history_for_propagation`` (a plausible
    prior suspect: it never triggers on the very first propagate call either).

    Traced to ``sam3.model.sam3_multiplex_tracking.Sam3MultiplexTracking.
    _build_sam2_output`` (in the pinned ``third_party/sam3`` checkout)::

        def _build_sam2_output(self, inference_state, frame_idx, refined_obj_id_to_mask=None):
            if not frame_idx in inference_state["cached_frame_outputs"]:
                return {}
            ...  # merge refined_obj_id_to_mask into cached_outputs.copy()

    The early return discards the caller-supplied ``refined_obj_id_to_mask`` --
    the just-computed SAM2-propagated mask for that frame -- whenever the frame has
    never been cached before, instead of treating the cache miss as an empty *base*
    dict to merge the refined masks into. Text/VG prompting never trips this: the
    initial VG detection pass pre-caches every frame before any refinement runs. A
    point prompt never runs that pass, so nearly every frame visited by propagation
    is a cache miss, and the real mask is thrown away right before
    ``_postprocess_output`` would have kept it -- its ``out_binary_masks.any(dim=(1,
    2))`` keep-filter then reports the frame as having no objects. (Frame 0 is the
    one exception: ``add_prompt`` itself populates ``cached_frame_outputs[0]``, even
    though *its own* immediate response is emptied by this same bug, so by the time
    propagation revisits frame 0 the cache-hit path runs and merges correctly --
    which is why the prompt frame alone looks fine.)

    Verified by patching this in and rerunning the repro: propagation area went
    from "real mask at frame 0, zero at every other frame" to a stable mask area
    (+/-2%) at all 30 frames, matching a text-prompted control run on the same clip.

    Applied as a self-test-driven monkeypatch (mirrors :func:`_patch_init_state_kwargs`)
    rather than an upstream fork, so it becomes a no-op automatically once sam3 fixes
    this: it calls the *unpatched* method with a synthetic cache-miss case first and
    only replaces it if that call reproduces the bug.
    """
    model = getattr(predictor, "model", None)
    build_sam2_output = getattr(model, "_build_sam2_output", None)
    if model is None or build_sam2_output is None:
        return

    probe_state = {"cached_frame_outputs": {}}
    try:
        probed = build_sam2_output(
            probe_state, frame_idx=0, refined_obj_id_to_mask={"probe": "mask"}
        )
    except Exception:  # pragma: no cover - upstream changed the method outright
        logger.debug(
            "could not probe _build_sam2_output for the point-prompt cache-miss bug; "
            "leaving it unpatched",
            exc_info=True,
        )
        return
    if probed:
        # Already merges correctly on a cache miss -- upstream fixed it, nothing to do.
        return

    def _fixed_build_sam2_output(self, inference_state, frame_idx, refined_obj_id_to_mask=None):
        obj_id_to_mask = inference_state["cached_frame_outputs"].get(frame_idx, {}).copy()
        if refined_obj_id_to_mask is not None:
            for obj_id, refined_mask in refined_obj_id_to_mask.items():
                assert refined_mask is not None, (
                    f"Refined mask data must be provided for obj_id {obj_id}"
                )
                obj_id_to_mask[obj_id] = refined_mask
        return obj_id_to_mask

    model._build_sam2_output = types.MethodType(_fixed_build_sam2_output, model)
    logger.info(
        "applied sam3 compatibility shim: _build_sam2_output was discarding refined "
        "masks on frames with no prior cache entry, which breaks point-prompt "
        "video propagation"
    )


class Sam3VideoSegmenter(VideoSegmenter):
    """Adapts SAM 3.1's multiplex video predictor to :class:`VideoSegmenter`.

    The predictor itself (``build_sam3_multiplex_video_predictor``) is
    expensive to construct -- it loads the SAM 3.1 checkpoint -- so it is
    built once, lazily, on first use and reused across every session this
    instance opens.
    """

    def __init__(
        self,
        config: SegmentationConfig,
        device: str = "cuda",
        checkpoint_path: str | None = None,
        bpe_path: str | None = None,
    ) -> None:
        """Args:
        config: Segmentation hyperparameters (``use_fa3``, offload flags, ...).
        device: Torch device SAM 3.1 should run on. Not currently accepted by
            ``build_sam3_multiplex_video_predictor`` itself -- pin the active
            GPU (e.g. via ``CUDA_VISIBLE_DEVICES`` or ``torch.cuda.set_device``)
            before constructing this adapter if ``RuntimeConfig.gpu_index`` is
            set. Kept as a constructor argument so the adapter's interface
            does not need to change if a future sam3 release adds one.
        checkpoint_path: Local checkpoint path; ``None`` auto-downloads
            ``facebook/sam3.1`` via ``hf_hub_download``.
        bpe_path: Optional tokenizer BPE path for the text encoder.
        """
        self.config = config
        self.device = device
        self.checkpoint_path = checkpoint_path
        self.bpe_path = bpe_path

        self._predictor = None  # built lazily: importing sam3 needs torch + weights
        self._session_id: str | None = None
        self._video_path: str | None = None
        self._has_text_prompt = False
        self._resolution: tuple[int, int] | None = None  # (width, height), cached from first output

    def _build(self):
        """Construct the SAM 3.1 predictor on first use and cache it."""
        if self._predictor is None:
            # Heavy, GPU-framework-dependent import: kept out of module scope so
            # `import fpgm.segmentation.sam3` works without sam3/torch installed.
            from sam3.model_builder import build_sam3_multiplex_video_predictor

            logger.info("building SAM 3.1 predictor (use_fa3=%s)", self.config.use_fa3)
            self._predictor = build_sam3_multiplex_video_predictor(
                checkpoint_path=self.checkpoint_path,
                bpe_path=self.bpe_path,
                max_num_objects=self.config.max_num_objects,
                multiplex_count=self.config.multiplex_count,
                # FlashAttention-3 is sam3's own default (use_fa3=True) but needs a
                # separate, optional install; SegmentationConfig.use_fa3 defaults
                # to False so this adapter works with a plain sam3 install.
                # use_rope_real must track use_fa3 -- they are not independent.
                use_fa3=self.config.use_fa3,
                use_rope_real=self.config.use_fa3,
                compile=False,
                warm_up=False,
                default_output_prob_thresh=self.config.output_prob_thresh,
            )
            _patch_init_state_kwargs(self._predictor)
            _patch_build_sam2_output(self._predictor)
        return self._predictor

    def start_session(self, video_path: str) -> str:
        """Open a session on an mp4 or a folder of ``<idx>.jpg`` frames."""
        predictor = self._build()
        response = predictor.handle_request(
            request=dict(
                type="start_session",
                resource_path=video_path,
                # The session preloads every frame as a (T, 3, 1008, 1008)
                # float16 tensor, ~6.1 MB/frame. This host's GPUs are shared and
                # often have only 11-27 GB free, so offload to host RAM by
                # default; SegmentationConfig.offload_video_to_cpu=True.
                offload_video_to_cpu=self.config.offload_video_to_cpu,
                offload_state_to_cpu=self.config.offload_state_to_cpu,
            )
        )
        self._session_id = response["session_id"]
        self._video_path = video_path
        self._has_text_prompt = False
        self._resolution = None
        return self._session_id

    def add_text_prompt(self, frame_idx: int, text: str) -> FrameMasks:
        """Prompt with a natural-language phrase; returns that frame's masks."""
        self._require_session()
        if self._has_text_prompt:
            logger.warning(
                "add_text_prompt: session %s already has a text prompt applied; "
                "auto-resetting before applying %r, per sam3's requirement that "
                "a new text prompt needs a fresh session state",
                self._session_id,
                text,
            )
            self.reset_session()
        predictor = self._build()
        response = predictor.handle_request(
            request=dict(
                type="add_prompt",
                session_id=self._session_id,
                frame_index=frame_idx,
                text=text,
                output_prob_thresh=self.config.output_prob_thresh,
            )
        )
        self._has_text_prompt = True
        return self._frame_masks_from_response(response)

    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."""
        self._require_session()
        predictor = self._build()

        width, height = self._resolution_for_conversion(frame_idx)
        points_xy = np.asarray(points_xy, dtype=np.float32)
        # sam3's add_prompt wants *relative* [0, 1] coordinates
        # (rel_coordinates=True is the default); our ABC contract is absolute
        # pixels, so convert here rather than pushing the distinction onto callers.
        rel_points = points_xy / np.array([width, height], dtype=np.float32)

        request = dict(
            type="add_prompt",
            session_id=self._session_id,
            frame_index=frame_idx,
            points=rel_points.tolist(),
            point_labels=np.asarray(labels, dtype=np.int32).tolist(),
            rel_coordinates=True,
            clear_old_points=False,
            output_prob_thresh=self.config.output_prob_thresh,
        )
        if obj_id is not None:
            request["obj_id"] = obj_id
        response = predictor.handle_request(request=request)
        return self._frame_masks_from_response(response)

    def propagate(
        self,
        direction: str = "both",
        start_frame_index: int = 0,
        max_frame_num_to_track: int | None = None,
    ) -> Iterator[FrameMasks]:
        """Stream per-frame masks across the whole video.

        Args:
            direction: ``"both"``, ``"forward"``, or ``"backward"``.
            start_frame_index: Frame to start propagation from.
            max_frame_num_to_track: Optional cap on the number of frames
                propagated (in each direction); unlimited if ``None``.
        """
        self._require_session()
        predictor = self._build()
        request = dict(
            type="propagate_in_video",
            session_id=self._session_id,
            propagation_direction=direction,
            start_frame_index=start_frame_index,
            output_prob_thresh=self.config.output_prob_thresh,
        )
        if max_frame_num_to_track is not None:
            request["max_frame_num_to_track"] = max_frame_num_to_track
        for event in predictor.handle_stream_request(request=request):
            yield self._frame_masks_from_outputs(int(event["frame_index"]), event["outputs"])

    def collect_masklets(
        self,
        direction: str = "both",
        start_frame_index: int = 0,
        max_frame_num_to_track: int | None = None,
    ) -> dict[int, Masklet]:
        """Run :meth:`propagate` across the whole clip and assemble one Masklet per object.

        Returns:
            ``obj_id -> Masklet``, with every frame's mask and score for that
            object.
        """
        masklets: dict[int, Masklet] = {}
        for frame_masks in self.propagate(
            direction=direction,
            start_frame_index=start_frame_index,
            max_frame_num_to_track=max_frame_num_to_track,
        ):
            for i, obj_id in enumerate(frame_masks.obj_ids.tolist()):
                masklet = masklets.setdefault(obj_id, Masklet(obj_id=obj_id))
                masklet.frames[frame_masks.frame_idx] = frame_masks.masks[i]
                masklet.scores[frame_masks.frame_idx] = float(frame_masks.scores[i])
        return masklets

    def select_object(
        self,
        masklets: dict[int, Masklet],
        gripper_uv_per_frame: dict[int, np.ndarray] | None = None,
        prefer_largest: bool = False,
    ) -> int:
        """Disambiguate which masklet is "the" object a text prompt referred to.

        A text prompt like ``"brick"`` can match several instances in a scene.
        Preference order:

        1. If ``gripper_uv_per_frame`` is given and
           ``config.select_by_gripper_proximity`` is set, pick the masklet
           whose mask centroid is (on average, over frames shared with the
           gripper trajectory) nearest the gripper's pixel projection -- the
           manipulated object is usually the one closest to the gripper.
        2. Otherwise, or if no masklet shares a frame with the gripper
           trajectory, fall back to mean mask area (``prefer_largest=True``)
           or mean detection score (the default).

        Args:
            masklets: ``obj_id -> Masklet``, typically from
                :meth:`collect_masklets`.
            gripper_uv_per_frame: ``frame_idx -> (2,)`` gripper pixel
                projection, at the same resolution as the masks.
            prefer_largest: Use mean mask area rather than mean score as the
                score/gripper-proximity fallback criterion.

        Returns:
            The chosen ``obj_id``.

        Raises:
            ValueError: If ``masklets`` is empty.
        """
        if not masklets:
            raise ValueError("select_object: masklets is empty")
        if len(masklets) == 1:
            return next(iter(masklets))

        if gripper_uv_per_frame and self.config.select_by_gripper_proximity:
            obj_id = self._select_by_gripper_proximity(masklets, gripper_uv_per_frame)
            if obj_id is not None:
                return obj_id
            logger.warning(
                "select_object: no masklet shares a frame with the gripper "
                "trajectory; falling back to %s",
                "mean area" if prefer_largest else "mean score",
            )

        if prefer_largest:
            return max(masklets, key=lambda oid: _mean_area(masklets[oid]))
        return max(masklets, key=lambda oid: _mean_score(masklets[oid]))

    def reset_session(self) -> None:
        """Clear prompts. Required before switching to a different text prompt."""
        self._require_session()
        predictor = self._build()
        predictor.handle_request(request=dict(type="reset_session", session_id=self._session_id))
        self._has_text_prompt = False

    def close_session(self) -> None:
        """Release the session and its GPU memory. Safe to call more than once."""
        if self._session_id is None or self._predictor is None:
            return  # never opened, or already closed: double-close must be a no-op
        try:
            self._predictor.handle_request(
                request=dict(type="close_session", session_id=self._session_id, run_gc_collect=True)
            )
        finally:
            # Always clear local state, even if the server-side close failed --
            # a stuck session handle here would wedge every subsequent call.
            self._session_id = None
            self._video_path = None
            self._has_text_prompt = False
            self._resolution = None

    def _require_session(self) -> None:
        if self._session_id is None:
            raise RuntimeError(
                f"{type(self).__name__}: no active session; call start_session() first"
            )

    def _resolution_for_conversion(self, frame_idx: int) -> tuple[int, int]:
        """Resolve the (width, height) needed to turn absolute pixels into sam3's relative coords.

        Cached from the first frame's mask output; if no prompt has produced
        output yet (e.g. ``add_point_prompt`` is the very first prompt) and
        the session was opened on an mp4 file, the resolution is probed
        directly from the container instead of decoding any frames.
        """
        if self._resolution is not None:
            return self._resolution
        if self._video_path is not None and Path(self._video_path).is_file():
            import cv2  # local import: only needed for this one-off probe

            cap = cv2.VideoCapture(self._video_path)
            try:
                width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            finally:
                cap.release()
            if width > 0 and height > 0:
                self._resolution = (width, height)
                return self._resolution
        raise RuntimeError(
            f"{type(self).__name__}: frame resolution is unknown at frame {frame_idx}; "
            "call add_text_prompt first, or open the session on an mp4 file so the "
            "resolution can be probed directly"
        )

    def _frame_masks_from_response(self, response: dict) -> FrameMasks:
        return self._frame_masks_from_outputs(int(response["frame_index"]), response["outputs"])

    def _frame_masks_from_outputs(self, frame_idx: int, outputs: dict) -> FrameMasks:
        # out_binary_masks is already dense bool (N, H_orig, W_orig) at the
        # video's original resolution -- no RLE/logit decoding needed.
        masks = np.asarray(outputs["out_binary_masks"], dtype=bool)
        if self._resolution is None and masks.ndim == 3:
            self._resolution = (int(masks.shape[2]), int(masks.shape[1]))
        return FrameMasks(
            frame_idx=frame_idx,
            obj_ids=np.asarray(outputs["out_obj_ids"], dtype=np.int64),
            masks=masks,
            scores=np.asarray(outputs["out_probs"], dtype=np.float32),
            boxes_xywh=np.asarray(outputs["out_boxes_xywh"], dtype=np.float32),
        )

    def _select_by_gripper_proximity(
        self,
        masklets: dict[int, Masklet],
        gripper_uv_per_frame: dict[int, np.ndarray],
    ) -> int | None:
        best_obj_id: int | None = None
        best_dist = np.inf
        for obj_id, masklet in masklets.items():
            dists = []
            for frame_idx, gripper_uv in gripper_uv_per_frame.items():
                mask = masklet.frames.get(frame_idx)
                if mask is None or not mask.any():
                    continue
                centroid = mask_centroid(mask)
                gripper = np.asarray(gripper_uv, dtype=np.float32)
                dists.append(float(np.linalg.norm(centroid - gripper)))
            if not dists:
                continue
            mean_dist = float(np.mean(dists))
            if mean_dist < best_dist:
                best_dist = mean_dist
                best_obj_id = obj_id
        return best_obj_id


def _mean_area(masklet: Masklet) -> float:
    if not masklet.frames:
        return 0.0
    return float(np.mean([mask.sum() for mask in masklet.frames.values()]))


def _mean_score(masklet: Masklet) -> float:
    if not masklet.scores:
        return 0.0
    return float(np.mean(list(masklet.scores.values())))