sam

sleap_nn.inference.sam

SAM-powered prompted instance segmentation for INFERENCE.

The pivot (PLAN / README): SAM is used to predict per-instance masks for an existing pose/centroid .slp so a human can review/correct them in the GUI, then train — not to auto-generate training GT. This package is the SAM1 + SAM3 prompted producers + their backend interface, plus the torch-less reconciliation / re-tracking path (:func:retrack); the SAM3 mask-native video tracker lands in a later PR behind the same surfaces.

Public surface

• :func:get_mask_backend — explicit, no-default backend selection (PLAN L2). "sam" builds a SAM1 :class:~.backends.SamBackend; "sam3" builds a SAM3 :class:~.backends.Sam3Backend (gated facebook/sam3 via the sleap_nn[sam3] extra); an unknown or omitted name raises.
• :func:run_sam_segmentation — end-to-end orchestration: load a pose .slp, run the chosen backend with the chosen prompt mode, emit sio.PredictedSegmentationMask (raw score + instance=/track= populated, PLAN L8) onto each frame, and optionally save the .slp (backreferencing the input's images, not re-embedding) + a review overlay PNG.
• :func:retrack (+ :mod:~sleap_nn.inference.sam.reconciliation primitives) — the torch-less "refine existing tracks" path: correct an existing pose/centroid tracker's identities from identity-consistent per-frame masks. No SAM / torch / transformers dependency (numpy + scipy only).

Everything heavy (segment-anything) is imported lazily inside the backend, so importing this package on a default install is cheap and dependency-free.

Modules:

Name	Description
backends	SAM mask backends for prompted instance segmentation (PR-A).
mask_layer	SAM mask inference layer — the producer that emits PredictedSegmentationMask.
overlay	Review/debug overlay rendering for predicted segmentation masks.
prompts	Prompt builders for SAM-prompted instance segmentation (PR-A).
reconciliation	ID reconciliation for matching SAM3 masks to poses or input masks.
retrack	Mask-based re-tracking: refine existing pose/centroid track identities.

Classes:

Name	Description
IDReconciler	Matches SAM3 masks to poses and reconciles track IDs.
MaskAssignment	A single mask-to-mask assignment at a frame.
MaskBackend	Abstract prompted-mask backend (the :class:SamBackend / SAM3 interface).
MaskReconciler	Matches SAM3 masks to input masks using IoU and reconciles track IDs.
MatchContext	Context for match predicate evaluation.
RetrackResult	Result of a :func:retrack run.
Sam3Backend	SAM3 (Meta SAM 3) prompted-mask backend (the sleap_nn[sam3] extra).
SamBackend	SAM1 (ViT-H) prompted-mask backend (the sleap_nn[sam] extra).
SamPrompt	A built SAM prompt for one instance.
SamSegmentationLayer	Full-frame SAM mask producer (pose / centroid / box prompts).
SwapEvent	Detected identity swap.
TrackAssignment	A single track assignment at a frame.
TrackNameResolver	Resolves SAM3 obj_ids to GT track names via nearest-anchor flood fill.

Functions:

Name	Description
default_match_predicate	Default match predicate: require at least 1 keypoint inside mask.
get_mask_backend	Build a mask backend by explicit name (no default; PLAN L2).
require_centroid_proximity	Create predicate requiring pose centroid near mask centroid.
require_min_fraction_inside	Create predicate requiring minimum fraction of keypoints inside mask.
require_min_keypoints_inside	Create predicate requiring minimum keypoints inside mask.
require_reasonable_mask_area	Create predicate requiring mask area within bounds.
run_sam_segmentation	Predict per-instance masks for a pose .slp with a SAM backend.

IDReconciler dataclass

Matches SAM3 masks to poses and reconciles track IDs.

This class implements Hungarian algorithm matching between pose instances and SAM3 segmentation masks, using keypoints-inside-mask as the cost metric.

Attributes:

Name	Type	Description
skeleton	Skeleton	The SLEAP skeleton for node name lookups.
exclude_nodes	set[str]	Set of node names to exclude from matching.
match_predicates	list[MatchPredicate]	List of predicates that must all pass for a valid match.
ignore_gt_tracks	bool	If True, do not propagate GT track names onto assignments (track_name is set to None).

Example

reconciler = IDReconciler( ... skeleton=handler.skeleton, ... exclude_nodes={"tail0", "tail1"}, ... ) for frame_idx in gt_frame_indices: ... assignments = reconciler.match_frame( ... frame_idx=frame_idx, ... poses=lf.instances, ... masks=result.masks, ... object_ids=result.object_ids, ... ) swaps = reconciler.detect_swaps() id_map = reconciler.build_id_map()

Methods:

Name	Description
__post_init__	Add default predicate if none provided.
build_id_map	Build frame -> {sam3_id -> track_name} mapping.
clear	Clear accumulated assignments.
compute_cost_matrix	Compute cost matrix for Hungarian matching.
detect_swaps	Detect identity swaps from accumulated assignments.
get_assignments	Get all accumulated assignments.
match_frame	Match poses to masks for a single frame.

Source code in sleap_nn/inference/sam/reconciliation.py

@dataclass
class IDReconciler:
    """Matches SAM3 masks to poses and reconciles track IDs.

    This class implements Hungarian algorithm matching between pose instances
    and SAM3 segmentation masks, using keypoints-inside-mask as the cost metric.

    Attributes:
        skeleton: The SLEAP skeleton for node name lookups.
        exclude_nodes: Set of node names to exclude from matching.
        match_predicates: List of predicates that must all pass for a valid match.
        ignore_gt_tracks: If True, do not propagate GT track names onto
            assignments (track_name is set to None).

    Example:
        >>> reconciler = IDReconciler(
        ...     skeleton=handler.skeleton,
        ...     exclude_nodes={"tail0", "tail1"},
        ... )
        >>> for frame_idx in gt_frame_indices:
        ...     assignments = reconciler.match_frame(
        ...         frame_idx=frame_idx,
        ...         poses=lf.instances,
        ...         masks=result.masks,
        ...         object_ids=result.object_ids,
        ...     )
        >>> swaps = reconciler.detect_swaps()
        >>> id_map = reconciler.build_id_map()
    """

    skeleton: "sio.Skeleton"
    exclude_nodes: set[str] = field(default_factory=set)
    match_predicates: list[MatchPredicate] = field(default_factory=list)
    ignore_gt_tracks: bool = False
    _assignments: list[TrackAssignment] = field(default_factory=list, repr=False)

    def __post_init__(self):
        """Add default predicate if none provided.

        The implicit default requires at least 3 keypoints inside the mask. The
        weaker ``default_match_predicate`` (>= 1) is kept defined for direct use
        but is no longer the implicit default. ``require_min_keypoints_inside``
        is defined later in this module; it resolves at call time, so referencing
        it here is fine.
        """
        if not self.match_predicates:
            self.match_predicates = [require_min_keypoints_inside(3)]

    def compute_cost_matrix(
        self,
        poses: list["sio.Instance"],
        masks: np.ndarray,
    ) -> np.ndarray:
        """Compute cost matrix for Hungarian matching.

        The cost is the negative number of visible keypoints inside each mask.
        Lower cost = better match (more keypoints inside).

        Args:
            poses: List of pose instances to match.
            masks: Array of masks with shape (N, H, W).

        Returns:
            Cost matrix with shape (n_poses, n_masks).
        """
        n_poses = len(poses)
        n_masks = len(masks)

        if n_poses == 0 or n_masks == 0:
            return np.zeros((n_poses, n_masks))

        cost = np.zeros((n_poses, n_masks))

        # Get node names for filtering
        node_names = [n.name for n in self.skeleton.nodes]
        height, width = masks.shape[1], masks.shape[2]

        for i, pose in enumerate(poses):
            coords = pose.numpy()
            # Keypoint is visible only if BOTH x and y are finite.
            visible_mask = ~np.isnan(coords).any(axis=1)

            # Apply node exclusion filter
            if self.exclude_nodes:
                for j, name in enumerate(node_names):
                    if name in self.exclude_nodes:
                        visible_mask[j] = False

            visible_coords = coords[visible_mask].astype(int)

            if len(visible_coords) == 0:
                continue

            xs = visible_coords[:, 0]
            ys = visible_coords[:, 1]

            # Bounds check: keypoints outside the mask are not counted.
            in_bounds = (xs >= 0) & (xs < width) & (ys >= 0) & (ys < height)
            if not in_bounds.any():
                continue

            xs_in = xs[in_bounds]
            ys_in = ys[in_bounds]

            # Vectorize across masks: masks[:, ys, xs] is (n_masks, n_kpts);
            # summing over keypoints gives per-mask inside counts.
            inside_counts = masks[:, ys_in, xs_in].astype(bool).sum(axis=1)

            # Negative because Hungarian minimizes cost.
            cost[i, :] = -inside_counts

        return cost

    def match_frame(
        self,
        frame_idx: int,
        poses: list["sio.Instance"],
        masks: np.ndarray,
        object_ids: np.ndarray,
        scores: np.ndarray | None = None,
    ) -> list[TrackAssignment]:
        """Match poses to masks for a single frame.

        Uses Hungarian algorithm for optimal assignment, then filters
        matches through predicates.

        Args:
            frame_idx: Frame index for this match.
            poses: List of pose instances to match.
            masks: Array of masks with shape (N, H, W) or (N, 1, H, W).
            object_ids: Array of SAM3 object IDs corresponding to masks.
            scores: Optional SAM3 mask detection confidence scores, shape (N,).

        Returns:
            List of valid TrackAssignment objects.
        """
        if len(poses) == 0 or len(masks) == 0:
            return []

        # Default scores to 1.0 if not provided
        if scores is None:
            scores = np.ones(len(object_ids))

        # Handle (N, 1, H, W) mask format from SAM3
        if masks.ndim == 4 and masks.shape[1] == 1:
            masks = masks.squeeze(axis=1)

        # Validate per-frame lengths so a mismatch surfaces clearly here rather
        # than as a bare IndexError deeper in the loop.
        if len(object_ids) != len(masks) or len(scores) != len(masks):
            raise ValueError(
                f"match_frame: frame {frame_idx} has {len(masks)} masks but "
                f"{len(object_ids)} object_ids / {len(scores)} scores"
            )

        # Compute cost matrix and solve assignment
        cost = self.compute_cost_matrix(poses, masks)
        row_ind, col_ind = linear_sum_assignment(cost)

        # Get node names for visibility calculation
        node_names = [n.name for n in self.skeleton.nodes]

        assignments = []
        for pose_idx, mask_idx in zip(row_ind, col_ind):
            pose = poses[pose_idx]
            mask = masks[mask_idx]

            # Calculate visibility (excluding filtered nodes). A keypoint is
            # visible only if BOTH x and y are finite.
            coords = pose.numpy()
            visible_mask = ~np.isnan(coords).any(axis=1)
            if self.exclude_nodes:
                for j, name in enumerate(node_names):
                    if name in self.exclude_nodes:
                        visible_mask[j] = False
            visible_count = int(visible_mask.sum())

            # Calculate mask statistics
            ys, xs = np.where(mask)
            if len(xs) > 0:
                centroid = (float(xs.mean()), float(ys.mean()))
                mask_area = int(mask.sum())
            else:
                centroid = (0.0, 0.0)
                mask_area = 0

            # Build context for predicate evaluation
            keypoints_inside = int(-cost[pose_idx, mask_idx])
            ctx = MatchContext(
                frame_idx=frame_idx,
                sam3_obj_id=int(object_ids[mask_idx]),
                cost=float(cost[pose_idx, mask_idx]),
                keypoints_inside=keypoints_inside,
                keypoints_visible=visible_count,
                mask_area=mask_area,
                mask_centroid=centroid,
            )

            # Apply match predicates
            if all(pred(pose, mask, ctx) for pred in self.match_predicates):
                # Use None for track_name when ignoring GT tracks
                if self.ignore_gt_tracks:
                    track_name = None
                else:
                    track_name = pose.track.name if pose.track else None
                confidence = (
                    keypoints_inside / visible_count if visible_count > 0 else 0.0
                )
                assignment = TrackAssignment(
                    frame_idx=frame_idx,
                    pose_track_name=track_name,
                    pose_idx=pose_idx,
                    sam3_obj_id=ctx.sam3_obj_id,
                    confidence=confidence,
                    sam3_score=float(scores[mask_idx]),
                )
                assignments.append(assignment)

        self._assignments.extend(assignments)
        return assignments

    def detect_swaps(self) -> list[SwapEvent]:
        """Detect identity swaps from accumulated assignments.

        A swap occurs when a track name is matched to different SAM3 object IDs
        across frames.

        Returns:
            List of SwapEvent objects describing detected swaps.
        """
        swaps = []
        by_track: dict[str, list[TrackAssignment]] = defaultdict(list)

        for a in self._assignments:
            if a.pose_track_name:
                by_track[a.pose_track_name].append(a)

        for track_name, track_assignments in by_track.items():
            track_assignments.sort(key=lambda a: a.frame_idx)

            for i in range(1, len(track_assignments)):
                prev = track_assignments[i - 1]
                curr = track_assignments[i]

                if prev.sam3_obj_id != curr.sam3_obj_id:
                    swaps.append(
                        SwapEvent(
                            frame_idx=curr.frame_idx,
                            track_name=track_name,
                            old_sam3_id=prev.sam3_obj_id,
                            new_sam3_id=curr.sam3_obj_id,
                        )
                    )

        return swaps

    def build_id_map(self) -> dict[int, dict[int, str]]:
        """Build frame -> {sam3_id -> track_name} mapping.

        This can be used to remap SAM3 object IDs to consistent track names
        in output files.

        Returns:
            Dictionary mapping frame_idx to {sam3_obj_id: track_name}.
        """
        by_frame: dict[int, dict[int, str]] = defaultdict(dict)
        for a in self._assignments:
            if a.pose_track_name:
                by_frame[a.frame_idx][a.sam3_obj_id] = a.pose_track_name
        return dict(by_frame)

    def get_assignments(self) -> list[TrackAssignment]:
        """Get all accumulated assignments.

        Returns:
            List of all TrackAssignment objects from match_frame() calls.
        """
        return list(self._assignments)

    def clear(self) -> None:
        """Clear accumulated assignments."""
        self._assignments.clear()

__post_init__()

Add default predicate if none provided.

The implicit default requires at least 3 keypoints inside the mask. The weaker default_match_predicate (>= 1) is kept defined for direct use but is no longer the implicit default. require_min_keypoints_inside is defined later in this module; it resolves at call time, so referencing it here is fine.

Source code in sleap_nn/inference/sam/reconciliation.py

def __post_init__(self):
    """Add default predicate if none provided.

    The implicit default requires at least 3 keypoints inside the mask. The
    weaker ``default_match_predicate`` (>= 1) is kept defined for direct use
    but is no longer the implicit default. ``require_min_keypoints_inside``
    is defined later in this module; it resolves at call time, so referencing
    it here is fine.
    """
    if not self.match_predicates:
        self.match_predicates = [require_min_keypoints_inside(3)]

build_id_map()

Build frame -> {sam3_id -> track_name} mapping.

This can be used to remap SAM3 object IDs to consistent track names in output files.

Returns:

Type	Description
dict[int, dict[int, str]]	Dictionary mapping frame_idx to {sam3_obj_id: track_name}.

Source code in sleap_nn/inference/sam/reconciliation.py

def build_id_map(self) -> dict[int, dict[int, str]]:
    """Build frame -> {sam3_id -> track_name} mapping.

    This can be used to remap SAM3 object IDs to consistent track names
    in output files.

    Returns:
        Dictionary mapping frame_idx to {sam3_obj_id: track_name}.
    """
    by_frame: dict[int, dict[int, str]] = defaultdict(dict)
    for a in self._assignments:
        if a.pose_track_name:
            by_frame[a.frame_idx][a.sam3_obj_id] = a.pose_track_name
    return dict(by_frame)

clear()

Clear accumulated assignments.

Source code in sleap_nn/inference/sam/reconciliation.py

def clear(self) -> None:
    """Clear accumulated assignments."""
    self._assignments.clear()

compute_cost_matrix(poses, masks)

Compute cost matrix for Hungarian matching.

The cost is the negative number of visible keypoints inside each mask. Lower cost = better match (more keypoints inside).

Parameters:

Name	Type	Description	Default
poses	list[Instance]	List of pose instances to match.	required
masks	ndarray	Array of masks with shape (N, H, W).	required

Returns:

Type	Description
ndarray	Cost matrix with shape (n_poses, n_masks).

Source code in sleap_nn/inference/sam/reconciliation.py

def compute_cost_matrix(
    self,
    poses: list["sio.Instance"],
    masks: np.ndarray,
) -> np.ndarray:
    """Compute cost matrix for Hungarian matching.

    The cost is the negative number of visible keypoints inside each mask.
    Lower cost = better match (more keypoints inside).

    Args:
        poses: List of pose instances to match.
        masks: Array of masks with shape (N, H, W).

    Returns:
        Cost matrix with shape (n_poses, n_masks).
    """
    n_poses = len(poses)
    n_masks = len(masks)

    if n_poses == 0 or n_masks == 0:
        return np.zeros((n_poses, n_masks))

    cost = np.zeros((n_poses, n_masks))

    # Get node names for filtering
    node_names = [n.name for n in self.skeleton.nodes]
    height, width = masks.shape[1], masks.shape[2]

    for i, pose in enumerate(poses):
        coords = pose.numpy()
        # Keypoint is visible only if BOTH x and y are finite.
        visible_mask = ~np.isnan(coords).any(axis=1)

        # Apply node exclusion filter
        if self.exclude_nodes:
            for j, name in enumerate(node_names):
                if name in self.exclude_nodes:
                    visible_mask[j] = False

        visible_coords = coords[visible_mask].astype(int)

        if len(visible_coords) == 0:
            continue

        xs = visible_coords[:, 0]
        ys = visible_coords[:, 1]

        # Bounds check: keypoints outside the mask are not counted.
        in_bounds = (xs >= 0) & (xs < width) & (ys >= 0) & (ys < height)
        if not in_bounds.any():
            continue

        xs_in = xs[in_bounds]
        ys_in = ys[in_bounds]

        # Vectorize across masks: masks[:, ys, xs] is (n_masks, n_kpts);
        # summing over keypoints gives per-mask inside counts.
        inside_counts = masks[:, ys_in, xs_in].astype(bool).sum(axis=1)

        # Negative because Hungarian minimizes cost.
        cost[i, :] = -inside_counts

    return cost

detect_swaps()

Detect identity swaps from accumulated assignments.

A swap occurs when a track name is matched to different SAM3 object IDs across frames.

Returns:

Type	Description
list[SwapEvent]	List of SwapEvent objects describing detected swaps.

Source code in sleap_nn/inference/sam/reconciliation.py

def detect_swaps(self) -> list[SwapEvent]:
    """Detect identity swaps from accumulated assignments.

    A swap occurs when a track name is matched to different SAM3 object IDs
    across frames.

    Returns:
        List of SwapEvent objects describing detected swaps.
    """
    swaps = []
    by_track: dict[str, list[TrackAssignment]] = defaultdict(list)

    for a in self._assignments:
        if a.pose_track_name:
            by_track[a.pose_track_name].append(a)

    for track_name, track_assignments in by_track.items():
        track_assignments.sort(key=lambda a: a.frame_idx)

        for i in range(1, len(track_assignments)):
            prev = track_assignments[i - 1]
            curr = track_assignments[i]

            if prev.sam3_obj_id != curr.sam3_obj_id:
                swaps.append(
                    SwapEvent(
                        frame_idx=curr.frame_idx,
                        track_name=track_name,
                        old_sam3_id=prev.sam3_obj_id,
                        new_sam3_id=curr.sam3_obj_id,
                    )
                )

    return swaps

get_assignments()

Get all accumulated assignments.

Returns:

Type	Description
list[TrackAssignment]	List of all TrackAssignment objects from match_frame() calls.

Source code in sleap_nn/inference/sam/reconciliation.py

def get_assignments(self) -> list[TrackAssignment]:
    """Get all accumulated assignments.

    Returns:
        List of all TrackAssignment objects from match_frame() calls.
    """
    return list(self._assignments)

match_frame(frame_idx, poses, masks, object_ids, scores=None)

Match poses to masks for a single frame.

Uses Hungarian algorithm for optimal assignment, then filters matches through predicates.

Parameters:

Name	Type	Description	Default
frame_idx	int	Frame index for this match.	required
poses	list[Instance]	List of pose instances to match.	required
masks	ndarray	Array of masks with shape (N, H, W) or (N, 1, H, W).	required
object_ids	ndarray	Array of SAM3 object IDs corresponding to masks.	required
scores	ndarray | None	Optional SAM3 mask detection confidence scores, shape (N,).	None

Returns:

Type	Description
list[TrackAssignment]	List of valid TrackAssignment objects.

Source code in sleap_nn/inference/sam/reconciliation.py

def match_frame(
    self,
    frame_idx: int,
    poses: list["sio.Instance"],
    masks: np.ndarray,
    object_ids: np.ndarray,
    scores: np.ndarray | None = None,
) -> list[TrackAssignment]:
    """Match poses to masks for a single frame.

    Uses Hungarian algorithm for optimal assignment, then filters
    matches through predicates.

    Args:
        frame_idx: Frame index for this match.
        poses: List of pose instances to match.
        masks: Array of masks with shape (N, H, W) or (N, 1, H, W).
        object_ids: Array of SAM3 object IDs corresponding to masks.
        scores: Optional SAM3 mask detection confidence scores, shape (N,).

    Returns:
        List of valid TrackAssignment objects.
    """
    if len(poses) == 0 or len(masks) == 0:
        return []

    # Default scores to 1.0 if not provided
    if scores is None:
        scores = np.ones(len(object_ids))

    # Handle (N, 1, H, W) mask format from SAM3
    if masks.ndim == 4 and masks.shape[1] == 1:
        masks = masks.squeeze(axis=1)

    # Validate per-frame lengths so a mismatch surfaces clearly here rather
    # than as a bare IndexError deeper in the loop.
    if len(object_ids) != len(masks) or len(scores) != len(masks):
        raise ValueError(
            f"match_frame: frame {frame_idx} has {len(masks)} masks but "
            f"{len(object_ids)} object_ids / {len(scores)} scores"
        )

    # Compute cost matrix and solve assignment
    cost = self.compute_cost_matrix(poses, masks)
    row_ind, col_ind = linear_sum_assignment(cost)

    # Get node names for visibility calculation
    node_names = [n.name for n in self.skeleton.nodes]

    assignments = []
    for pose_idx, mask_idx in zip(row_ind, col_ind):
        pose = poses[pose_idx]
        mask = masks[mask_idx]

        # Calculate visibility (excluding filtered nodes). A keypoint is
        # visible only if BOTH x and y are finite.
        coords = pose.numpy()
        visible_mask = ~np.isnan(coords).any(axis=1)
        if self.exclude_nodes:
            for j, name in enumerate(node_names):
                if name in self.exclude_nodes:
                    visible_mask[j] = False
        visible_count = int(visible_mask.sum())

        # Calculate mask statistics
        ys, xs = np.where(mask)
        if len(xs) > 0:
            centroid = (float(xs.mean()), float(ys.mean()))
            mask_area = int(mask.sum())
        else:
            centroid = (0.0, 0.0)
            mask_area = 0

        # Build context for predicate evaluation
        keypoints_inside = int(-cost[pose_idx, mask_idx])
        ctx = MatchContext(
            frame_idx=frame_idx,
            sam3_obj_id=int(object_ids[mask_idx]),
            cost=float(cost[pose_idx, mask_idx]),
            keypoints_inside=keypoints_inside,
            keypoints_visible=visible_count,
            mask_area=mask_area,
            mask_centroid=centroid,
        )

        # Apply match predicates
        if all(pred(pose, mask, ctx) for pred in self.match_predicates):
            # Use None for track_name when ignoring GT tracks
            if self.ignore_gt_tracks:
                track_name = None
            else:
                track_name = pose.track.name if pose.track else None
            confidence = (
                keypoints_inside / visible_count if visible_count > 0 else 0.0
            )
            assignment = TrackAssignment(
                frame_idx=frame_idx,
                pose_track_name=track_name,
                pose_idx=pose_idx,
                sam3_obj_id=ctx.sam3_obj_id,
                confidence=confidence,
                sam3_score=float(scores[mask_idx]),
            )
            assignments.append(assignment)

    self._assignments.extend(assignments)
    return assignments

MaskAssignment dataclass

A single mask-to-mask assignment at a frame.

Used for matching input (anchor) masks to SAM3 output masks.

Attributes:

Name	Type	Description
frame_idx	int	Frame index where assignment was made.
input_track_id	int	Track ID from the input/anchor mask.
input_track_name	str | None	Track name from the input/anchor mask.
sam3_obj_id	int	SAM3 object ID that was matched.
iou	float	Intersection over Union score for the match.
sam3_score	float	SAM3 mask detection confidence score.

Source code in sleap_nn/inference/sam/reconciliation.py

@dataclass
class MaskAssignment:
    """A single mask-to-mask assignment at a frame.

    Used for matching input (anchor) masks to SAM3 output masks.

    Attributes:
        frame_idx: Frame index where assignment was made.
        input_track_id: Track ID from the input/anchor mask.
        input_track_name: Track name from the input/anchor mask.
        sam3_obj_id: SAM3 object ID that was matched.
        iou: Intersection over Union score for the match.
        sam3_score: SAM3 mask detection confidence score.
    """

    frame_idx: int
    input_track_id: int
    input_track_name: str | None
    sam3_obj_id: int
    iou: float
    sam3_score: float = 1.0

MaskBackend

Bases: ABC

Abstract prompted-mask backend (the :class:SamBackend / SAM3 interface).

A backend encodes one image and answers a batch of prompts on it. The composed inference layer (:mod:sleap_nn.inference.sam.mask_layer) owns the crop/frame geometry; the backend owns only the model call. Selection is explicit (PLAN L2) — see :func:sleap_nn.inference.sam.get_mask_backend.

Methods:

Name	Description
masks	Encode image once and answer each prompt with a mask + raw score.

Source code in sleap_nn/inference/sam/backends.py

class MaskBackend(ABC):
    """Abstract prompted-mask backend (the :class:`SamBackend` / SAM3 interface).

    A backend encodes one image and answers a batch of prompts on it. The
    composed inference layer (:mod:`sleap_nn.inference.sam.mask_layer`) owns the
    crop/frame geometry; the backend owns only the model call. Selection is
    explicit (PLAN L2) — see :func:`sleap_nn.inference.sam.get_mask_backend`.
    """

    #: Per-model nominal predicted-IoU floor. Defaults to SAM1's ``0.88``;
    #: :class:`Sam3Backend` overrides it with a recalibrated ``0.5`` (SAM3's
    #: predicted-IoU is on a lower scale, PLAN §2.3). Carried as a per-model
    #: attribute so SAM3 can override it; SAM1's raw predicted-IoU is reported as
    #: the mask score, not used as a gate.
    pred_iou_min: float = 0.88

    @abstractmethod
    def masks(
        self, image: np.ndarray, prompts: Sequence[SamPrompt]
    ) -> Tuple[List[np.ndarray], List[float]]:
        """Encode ``image`` once and answer each prompt with a mask + raw score.

        Args:
            image: ``(H, W)`` grayscale (or ``(H, W, C)``) image to encode.
            prompts: Per-instance prompts in image space.

        Returns:
            ``(masks, scores)``: a list of ``(H, W)`` boolean masks (one per
            prompt) and the list of raw per-model scores.
        """
        raise NotImplementedError

masks(image, prompts) abstractmethod

Encode image once and answer each prompt with a mask + raw score.

Parameters:

Name	Type	Description	Default
image	ndarray	(H, W) grayscale (or (H, W, C)) image to encode.	required
prompts	Sequence[SamPrompt]	Per-instance prompts in image space.	required

Returns:

Type	Description
Tuple[List[ndarray], List[float]]	(masks, scores): a list of (H, W) boolean masks (one per prompt) and the list of raw per-model scores.

Source code in sleap_nn/inference/sam/backends.py

@abstractmethod
def masks(
    self, image: np.ndarray, prompts: Sequence[SamPrompt]
) -> Tuple[List[np.ndarray], List[float]]:
    """Encode ``image`` once and answer each prompt with a mask + raw score.

    Args:
        image: ``(H, W)`` grayscale (or ``(H, W, C)``) image to encode.
        prompts: Per-instance prompts in image space.

    Returns:
        ``(masks, scores)``: a list of ``(H, W)`` boolean masks (one per
        prompt) and the list of raw per-model scores.
    """
    raise NotImplementedError

MaskReconciler dataclass

Matches SAM3 masks to input masks using IoU and reconciles track IDs.

This class implements Hungarian algorithm matching between input/anchor masks and SAM3 segmentation masks, using IoU (Intersection over Union) as the cost metric. This enables post-hoc identity correction using sparse ground truth mask annotations.

Unlike IDReconciler (which uses keypoints-in-mask for pose matching), this reconciler works purely with mask overlap, making it suitable for workflows where users have corrected masks at specific frames that should be used as identity anchors.

Attributes:

Name	Type	Description
min_iou	float	Minimum IoU threshold for a valid match. Matches below this threshold are rejected.
track_names	dict[int, str]	Optional mapping of input track_id -> name for naming.

Example

reconciler = MaskReconciler(min_iou=0.3) for frame_idx in anchor_frames: ... assignments = reconciler.match_frame( ... frame_idx=frame_idx, ... input_masks=reader.get_masks(frame_idx), ... input_track_ids=reader.get_track_ids(frame_idx), ... sam3_masks=result.masks, ... sam3_obj_ids=result.object_ids, ... ) swaps = reconciler.detect_swaps() id_map = reconciler.build_id_map()

Methods:

Name	Description
build_id_map	Build frame -> {sam3_id -> track_name} mapping.
clear	Clear accumulated assignments.
compute_cost_matrix	Compute cost matrix for Hungarian matching.
compute_iou	Compute Intersection over Union between two binary masks.
detect_swaps	Detect identity swaps from accumulated assignments.
get_assignments	Get all accumulated assignments.
get_iou_stats	Get IoU statistics from accumulated assignments.
match_frame	Match input masks to SAM3 masks for a single frame.

Source code in sleap_nn/inference/sam/reconciliation.py

@dataclass
class MaskReconciler:
    """Matches SAM3 masks to input masks using IoU and reconciles track IDs.

    This class implements Hungarian algorithm matching between input/anchor masks
    and SAM3 segmentation masks, using IoU (Intersection over Union) as the cost
    metric. This enables post-hoc identity correction using sparse ground truth
    mask annotations.

    Unlike IDReconciler (which uses keypoints-in-mask for pose matching), this
    reconciler works purely with mask overlap, making it suitable for workflows
    where users have corrected masks at specific frames that should be used as
    identity anchors.

    Attributes:
        min_iou: Minimum IoU threshold for a valid match. Matches below this
            threshold are rejected.
        track_names: Optional mapping of input track_id -> name for naming.

    Example:
        >>> reconciler = MaskReconciler(min_iou=0.3)
        >>> for frame_idx in anchor_frames:
        ...     assignments = reconciler.match_frame(
        ...         frame_idx=frame_idx,
        ...         input_masks=reader.get_masks(frame_idx),
        ...         input_track_ids=reader.get_track_ids(frame_idx),
        ...         sam3_masks=result.masks,
        ...         sam3_obj_ids=result.object_ids,
        ...     )
        >>> swaps = reconciler.detect_swaps()
        >>> id_map = reconciler.build_id_map()
    """

    min_iou: float = 0.3
    track_names: dict[int, str] = field(default_factory=dict)
    _assignments: list[MaskAssignment] = field(default_factory=list, repr=False)

    @staticmethod
    def compute_iou(mask1: np.ndarray, mask2: np.ndarray) -> float:
        """Compute Intersection over Union between two binary masks.

        Args:
            mask1: First binary mask as (H, W) array.
            mask2: Second binary mask as (H, W) array.

        Returns:
            IoU score between 0 and 1.
        """
        # Convert to boolean for logical operations
        m1 = mask1.astype(bool)
        m2 = mask2.astype(bool)

        intersection = np.logical_and(m1, m2).sum()
        union = np.logical_or(m1, m2).sum()

        if union == 0:
            return 0.0
        return float(intersection / union)

    def compute_cost_matrix(
        self,
        input_masks: np.ndarray,
        sam3_masks: np.ndarray,
    ) -> np.ndarray:
        """Compute cost matrix for Hungarian matching.

        The cost is the negative IoU (because Hungarian minimizes cost).
        Lower cost = better match (higher IoU).

        Args:
            input_masks: Input/anchor masks with shape (N, H, W).
            sam3_masks: SAM3 output masks with shape (M, H, W) or (M, 1, H, W).

        Returns:
            Cost matrix with shape (n_input, n_sam3).
        """
        # Handle (M, 1, H, W) mask format from SAM3
        if sam3_masks.ndim == 4 and sam3_masks.shape[1] == 1:
            sam3_masks = sam3_masks.squeeze(axis=1)

        n_input = len(input_masks)
        n_sam3 = len(sam3_masks)

        if n_input == 0 or n_sam3 == 0:
            return np.zeros((n_input, n_sam3))

        cost = np.zeros((n_input, n_sam3))

        for i, input_mask in enumerate(input_masks):
            for j, sam3_mask in enumerate(sam3_masks):
                iou = self.compute_iou(input_mask, sam3_mask)
                # Negative because Hungarian minimizes cost
                cost[i, j] = -iou

        return cost

    def match_frame(
        self,
        frame_idx: int,
        input_masks: np.ndarray,
        input_track_ids: np.ndarray,
        sam3_masks: np.ndarray,
        sam3_obj_ids: np.ndarray,
        scores: np.ndarray | None = None,
    ) -> list[MaskAssignment]:
        """Match input masks to SAM3 masks for a single frame.

        Uses Hungarian algorithm for optimal assignment, then filters
        matches by IoU threshold.

        Args:
            frame_idx: Frame index for this match.
            input_masks: Input/anchor masks with shape (N, H, W).
            input_track_ids: Track IDs corresponding to input masks.
            sam3_masks: SAM3 output masks with shape (M, H, W) or (M, 1, H, W).
            sam3_obj_ids: SAM3 object IDs corresponding to SAM3 masks.
            scores: Optional SAM3 mask detection confidence scores, shape (M,).

        Returns:
            List of valid MaskAssignment objects.
        """
        if len(input_masks) == 0 or len(sam3_masks) == 0:
            return []

        # Default scores to 1.0 if not provided
        if scores is None:
            scores = np.ones(len(sam3_obj_ids))

        # Compute cost matrix and solve assignment
        cost = self.compute_cost_matrix(input_masks, sam3_masks)
        row_ind, col_ind = linear_sum_assignment(cost)

        assignments = []
        for input_idx, sam3_idx in zip(row_ind, col_ind):
            iou = -cost[input_idx, sam3_idx]  # Convert back from negative

            # Apply IoU threshold
            if iou < self.min_iou:
                continue

            input_track_id = int(input_track_ids[input_idx])
            track_name = self.track_names.get(input_track_id)

            assignment = MaskAssignment(
                frame_idx=frame_idx,
                input_track_id=input_track_id,
                input_track_name=track_name,
                sam3_obj_id=int(sam3_obj_ids[sam3_idx]),
                iou=iou,
                sam3_score=float(scores[sam3_idx]),
            )
            assignments.append(assignment)

        self._assignments.extend(assignments)
        return assignments

    def detect_swaps(self) -> list[SwapEvent]:
        """Detect identity swaps from accumulated assignments.

        A swap occurs when an input track is matched to different SAM3 object IDs
        across frames.

        Returns:
            List of SwapEvent objects describing detected swaps.
        """
        swaps = []
        by_track: dict[int, list[MaskAssignment]] = defaultdict(list)

        for a in self._assignments:
            by_track[a.input_track_id].append(a)

        for input_track_id, track_assignments in by_track.items():
            track_assignments.sort(key=lambda a: a.frame_idx)

            for i in range(1, len(track_assignments)):
                prev = track_assignments[i - 1]
                curr = track_assignments[i]

                if prev.sam3_obj_id != curr.sam3_obj_id:
                    # Use track name if available, otherwise use "track_{id}"
                    track_name = (
                        curr.input_track_name
                        or self.track_names.get(input_track_id)
                        or f"track_{input_track_id}"
                    )
                    swaps.append(
                        SwapEvent(
                            frame_idx=curr.frame_idx,
                            track_name=track_name,
                            old_sam3_id=prev.sam3_obj_id,
                            new_sam3_id=curr.sam3_obj_id,
                        )
                    )

        return swaps

    def build_id_map(self) -> dict[int, dict[int, str]]:
        """Build frame -> {sam3_id -> track_name} mapping.

        This can be used to remap SAM3 object IDs to consistent track names
        in output files.

        Returns:
            Dictionary mapping frame_idx to {sam3_obj_id: track_name}.
        """
        by_frame: dict[int, dict[int, str]] = defaultdict(dict)
        for a in self._assignments:
            name = (
                a.input_track_name
                or self.track_names.get(a.input_track_id)
                or f"track_{a.input_track_id}"
            )
            by_frame[a.frame_idx][a.sam3_obj_id] = name
        return dict(by_frame)

    def get_assignments(self) -> list[MaskAssignment]:
        """Get all accumulated assignments.

        Returns:
            List of all MaskAssignment objects from match_frame() calls.
        """
        return list(self._assignments)

    def get_iou_stats(self) -> dict[str, float]:
        """Get IoU statistics from accumulated assignments.

        Returns:
            Dictionary with 'min', 'max', 'mean', 'median' IoU values.
        """
        if not self._assignments:
            return {"min": 0.0, "max": 0.0, "mean": 0.0, "median": 0.0}

        ious = [a.iou for a in self._assignments]
        return {
            "min": float(min(ious)),
            "max": float(max(ious)),
            "mean": float(np.mean(ious)),
            "median": float(np.median(ious)),
        }

    def clear(self) -> None:
        """Clear accumulated assignments."""
        self._assignments.clear()

build_id_map()

Build frame -> {sam3_id -> track_name} mapping.

This can be used to remap SAM3 object IDs to consistent track names in output files.

Returns:

Type	Description
dict[int, dict[int, str]]	Dictionary mapping frame_idx to {sam3_obj_id: track_name}.

Source code in sleap_nn/inference/sam/reconciliation.py

def build_id_map(self) -> dict[int, dict[int, str]]:
    """Build frame -> {sam3_id -> track_name} mapping.

    This can be used to remap SAM3 object IDs to consistent track names
    in output files.

    Returns:
        Dictionary mapping frame_idx to {sam3_obj_id: track_name}.
    """
    by_frame: dict[int, dict[int, str]] = defaultdict(dict)
    for a in self._assignments:
        name = (
            a.input_track_name
            or self.track_names.get(a.input_track_id)
            or f"track_{a.input_track_id}"
        )
        by_frame[a.frame_idx][a.sam3_obj_id] = name
    return dict(by_frame)

clear()

Clear accumulated assignments.

Source code in sleap_nn/inference/sam/reconciliation.py

def clear(self) -> None:
    """Clear accumulated assignments."""
    self._assignments.clear()

compute_cost_matrix(input_masks, sam3_masks)

Compute cost matrix for Hungarian matching.

The cost is the negative IoU (because Hungarian minimizes cost). Lower cost = better match (higher IoU).

Parameters:

Name	Type	Description	Default
input_masks	ndarray	Input/anchor masks with shape (N, H, W).	required
sam3_masks	ndarray	SAM3 output masks with shape (M, H, W) or (M, 1, H, W).	required

Returns:

Type	Description
ndarray	Cost matrix with shape (n_input, n_sam3).

Source code in sleap_nn/inference/sam/reconciliation.py

def compute_cost_matrix(
    self,
    input_masks: np.ndarray,
    sam3_masks: np.ndarray,
) -> np.ndarray:
    """Compute cost matrix for Hungarian matching.

    The cost is the negative IoU (because Hungarian minimizes cost).
    Lower cost = better match (higher IoU).

    Args:
        input_masks: Input/anchor masks with shape (N, H, W).
        sam3_masks: SAM3 output masks with shape (M, H, W) or (M, 1, H, W).

    Returns:
        Cost matrix with shape (n_input, n_sam3).
    """
    # Handle (M, 1, H, W) mask format from SAM3
    if sam3_masks.ndim == 4 and sam3_masks.shape[1] == 1:
        sam3_masks = sam3_masks.squeeze(axis=1)

    n_input = len(input_masks)
    n_sam3 = len(sam3_masks)

    if n_input == 0 or n_sam3 == 0:
        return np.zeros((n_input, n_sam3))

    cost = np.zeros((n_input, n_sam3))

    for i, input_mask in enumerate(input_masks):
        for j, sam3_mask in enumerate(sam3_masks):
            iou = self.compute_iou(input_mask, sam3_mask)
            # Negative because Hungarian minimizes cost
            cost[i, j] = -iou

    return cost

compute_iou(mask1, mask2) staticmethod

Compute Intersection over Union between two binary masks.

Parameters:

Name	Type	Description	Default
mask1	ndarray	First binary mask as (H, W) array.	required
mask2	ndarray	Second binary mask as (H, W) array.	required

Returns:

Type	Description
float	IoU score between 0 and 1.

Source code in sleap_nn/inference/sam/reconciliation.py

@staticmethod
def compute_iou(mask1: np.ndarray, mask2: np.ndarray) -> float:
    """Compute Intersection over Union between two binary masks.

    Args:
        mask1: First binary mask as (H, W) array.
        mask2: Second binary mask as (H, W) array.

    Returns:
        IoU score between 0 and 1.
    """
    # Convert to boolean for logical operations
    m1 = mask1.astype(bool)
    m2 = mask2.astype(bool)

    intersection = np.logical_and(m1, m2).sum()
    union = np.logical_or(m1, m2).sum()

    if union == 0:
        return 0.0
    return float(intersection / union)

detect_swaps()

Detect identity swaps from accumulated assignments.

A swap occurs when an input track is matched to different SAM3 object IDs across frames.

Returns:

Type	Description
list[SwapEvent]	List of SwapEvent objects describing detected swaps.

Source code in sleap_nn/inference/sam/reconciliation.py

def detect_swaps(self) -> list[SwapEvent]:
    """Detect identity swaps from accumulated assignments.

    A swap occurs when an input track is matched to different SAM3 object IDs
    across frames.

    Returns:
        List of SwapEvent objects describing detected swaps.
    """
    swaps = []
    by_track: dict[int, list[MaskAssignment]] = defaultdict(list)

    for a in self._assignments:
        by_track[a.input_track_id].append(a)

    for input_track_id, track_assignments in by_track.items():
        track_assignments.sort(key=lambda a: a.frame_idx)

        for i in range(1, len(track_assignments)):
            prev = track_assignments[i - 1]
            curr = track_assignments[i]

            if prev.sam3_obj_id != curr.sam3_obj_id:
                # Use track name if available, otherwise use "track_{id}"
                track_name = (
                    curr.input_track_name
                    or self.track_names.get(input_track_id)
                    or f"track_{input_track_id}"
                )
                swaps.append(
                    SwapEvent(
                        frame_idx=curr.frame_idx,
                        track_name=track_name,
                        old_sam3_id=prev.sam3_obj_id,
                        new_sam3_id=curr.sam3_obj_id,
                    )
                )

    return swaps

get_assignments()

Get all accumulated assignments.

Returns:

Type	Description
list[MaskAssignment]	List of all MaskAssignment objects from match_frame() calls.

Source code in sleap_nn/inference/sam/reconciliation.py

def get_assignments(self) -> list[MaskAssignment]:
    """Get all accumulated assignments.

    Returns:
        List of all MaskAssignment objects from match_frame() calls.
    """
    return list(self._assignments)

get_iou_stats()

Get IoU statistics from accumulated assignments.

Returns:

Type	Description
dict[str, float]	Dictionary with 'min', 'max', 'mean', 'median' IoU values.

Source code in sleap_nn/inference/sam/reconciliation.py

def get_iou_stats(self) -> dict[str, float]:
    """Get IoU statistics from accumulated assignments.

    Returns:
        Dictionary with 'min', 'max', 'mean', 'median' IoU values.
    """
    if not self._assignments:
        return {"min": 0.0, "max": 0.0, "mean": 0.0, "median": 0.0}

    ious = [a.iou for a in self._assignments]
    return {
        "min": float(min(ious)),
        "max": float(max(ious)),
        "mean": float(np.mean(ious)),
        "median": float(np.median(ious)),
    }

match_frame(frame_idx, input_masks, input_track_ids, sam3_masks, sam3_obj_ids, scores=None)

Match input masks to SAM3 masks for a single frame.

Uses Hungarian algorithm for optimal assignment, then filters matches by IoU threshold.

Parameters:

Name	Type	Description	Default
frame_idx	int	Frame index for this match.	required
input_masks	ndarray	Input/anchor masks with shape (N, H, W).	required
input_track_ids	ndarray	Track IDs corresponding to input masks.	required
sam3_masks	ndarray	SAM3 output masks with shape (M, H, W) or (M, 1, H, W).	required
sam3_obj_ids	ndarray	SAM3 object IDs corresponding to SAM3 masks.	required
scores	ndarray | None	Optional SAM3 mask detection confidence scores, shape (M,).	None

Returns:

Type	Description
list[MaskAssignment]	List of valid MaskAssignment objects.

Source code in sleap_nn/inference/sam/reconciliation.py

def match_frame(
    self,
    frame_idx: int,
    input_masks: np.ndarray,
    input_track_ids: np.ndarray,
    sam3_masks: np.ndarray,
    sam3_obj_ids: np.ndarray,
    scores: np.ndarray | None = None,
) -> list[MaskAssignment]:
    """Match input masks to SAM3 masks for a single frame.

    Uses Hungarian algorithm for optimal assignment, then filters
    matches by IoU threshold.

    Args:
        frame_idx: Frame index for this match.
        input_masks: Input/anchor masks with shape (N, H, W).
        input_track_ids: Track IDs corresponding to input masks.
        sam3_masks: SAM3 output masks with shape (M, H, W) or (M, 1, H, W).
        sam3_obj_ids: SAM3 object IDs corresponding to SAM3 masks.
        scores: Optional SAM3 mask detection confidence scores, shape (M,).

    Returns:
        List of valid MaskAssignment objects.
    """
    if len(input_masks) == 0 or len(sam3_masks) == 0:
        return []

    # Default scores to 1.0 if not provided
    if scores is None:
        scores = np.ones(len(sam3_obj_ids))

    # Compute cost matrix and solve assignment
    cost = self.compute_cost_matrix(input_masks, sam3_masks)
    row_ind, col_ind = linear_sum_assignment(cost)

    assignments = []
    for input_idx, sam3_idx in zip(row_ind, col_ind):
        iou = -cost[input_idx, sam3_idx]  # Convert back from negative

        # Apply IoU threshold
        if iou < self.min_iou:
            continue

        input_track_id = int(input_track_ids[input_idx])
        track_name = self.track_names.get(input_track_id)

        assignment = MaskAssignment(
            frame_idx=frame_idx,
            input_track_id=input_track_id,
            input_track_name=track_name,
            sam3_obj_id=int(sam3_obj_ids[sam3_idx]),
            iou=iou,
            sam3_score=float(scores[sam3_idx]),
        )
        assignments.append(assignment)

    self._assignments.extend(assignments)
    return assignments

MatchContext dataclass

Context for match predicate evaluation.

Attributes:

Name	Type	Description
frame_idx	int	Frame index where the match was made.
sam3_obj_id	int	SAM3 object ID of the matched mask.
cost	float	Raw cost from the cost matrix (negative keypoints inside).
keypoints_inside	int	Number of visible keypoints inside the mask.
keypoints_visible	int	Total number of visible keypoints in the pose.
mask_area	int	Area of the mask in pixels.
mask_centroid	tuple[float, float]	Centroid of the mask as (x, y).

Source code in sleap_nn/inference/sam/reconciliation.py

@dataclass
class MatchContext:
    """Context for match predicate evaluation.

    Attributes:
        frame_idx: Frame index where the match was made.
        sam3_obj_id: SAM3 object ID of the matched mask.
        cost: Raw cost from the cost matrix (negative keypoints inside).
        keypoints_inside: Number of visible keypoints inside the mask.
        keypoints_visible: Total number of visible keypoints in the pose.
        mask_area: Area of the mask in pixels.
        mask_centroid: Centroid of the mask as (x, y).
    """

    frame_idx: int
    sam3_obj_id: int
    cost: float
    keypoints_inside: int
    keypoints_visible: int
    mask_area: int
    mask_centroid: tuple[float, float]

RetrackResult dataclass

Result of a :func:retrack run.

Attributes:

Name	Type	Description
labeled_frames	list['sio.LabeledFrame']	The relabeled frames. Same objects as the input when in_place=True; the corrected deep copies when in_place=False.
assignments	list[TrackAssignment]	All :class:TrackAssignment objects produced across frames (pose<->mask Hungarian matches, after predicate filtering).
id_map	dict[int, dict[int, str]]	Sparse anchor map frame_idx -> {mask_obj_id: track_name} built from trusted (anchor) frames only.
canonical_map	dict[int, str]	The global mask_obj_id -> track_name map used to relabel instances. Each obj_id is named by majority vote across anchor frames; obj_ids with no clear majority (an exact tie) are omitted here and resolved per-frame via the nearest anchor.
resolver	TrackNameResolver | None	The :class:TrackNameResolver used for nearest-anchor fallback, exposed for inspection / debugging.
num_relabeled	int	Number of instances whose track was changed.
num_matched	int	Number of instances that received a mask match.
anchor_frames	list[int]	Sorted frame indices used as identity anchors.

Source code in sleap_nn/inference/sam/retrack.py

@dataclass
class RetrackResult:
    """Result of a :func:`retrack` run.

    Attributes:
        labeled_frames: The relabeled frames. Same objects as the input when
            ``in_place=True``; the corrected deep copies when ``in_place=False``.
        assignments: All :class:`TrackAssignment` objects produced across frames
            (pose<->mask Hungarian matches, after predicate filtering).
        id_map: Sparse anchor map ``frame_idx -> {mask_obj_id: track_name}``
            built from trusted (anchor) frames only.
        canonical_map: The global ``mask_obj_id -> track_name`` map used to
            relabel instances. Each obj_id is named by majority vote across
            anchor frames; obj_ids with no clear majority (an exact tie) are
            omitted here and resolved per-frame via the nearest anchor.
        resolver: The :class:`TrackNameResolver` used for nearest-anchor
            fallback, exposed for inspection / debugging.
        num_relabeled: Number of instances whose ``track`` was changed.
        num_matched: Number of instances that received a mask match.
        anchor_frames: Sorted frame indices used as identity anchors.
    """

    labeled_frames: list["sio.LabeledFrame"] = field(default_factory=list)
    assignments: list[TrackAssignment] = field(default_factory=list)
    id_map: dict[int, dict[int, str]] = field(default_factory=dict)
    canonical_map: dict[int, str] = field(default_factory=dict)
    resolver: TrackNameResolver | None = None
    num_relabeled: int = 0
    num_matched: int = 0
    anchor_frames: list[int] = field(default_factory=list)

Sam3Backend

Bases: MaskBackend

SAM3 (Meta SAM 3) prompted-mask backend (the sleap_nn[sam3] extra).

Wraps a lazily loaded transformers Sam3TrackerModel + Sam3TrackerProcessor image visual-prompt pair. Honors the same :class:MaskBackend surface as :class:SamBackend, but two SAM3 specifics are mandatory and NEVER shared with SAM1 (PLAN §2.3, harvested from #643):

• Recalibrated floor. SAM3's iou_scores (predicted-IoU) are on a LOWER scale than SAM1 (median ~0.68 vs SAM1's ~0.95). SAM1's 0.88 floor applied verbatim would drop ~100% of SAM3 masks as a pure calibration artifact, so the per-model :attr:pred_iou_min defaults to 0.5 (~SAM1's 0.88 in percentile terms), never SAM1's 0.88. As with SAM1 the raw chosen-candidate score is reported, not gated on.
• Speckle cleanup. Raw SAM3 masks are speckly/fragmented (median ~14 connected components per mask vs SAM1's 1), with ~97% of the area in the keypoint-connected component. The speckle is cosmetic, so each chosen mask is passed through :func:_cleanup_speckle (morphological open + close + keep-keypoint-component, -> median 1 component, ~97% area retained) before it is returned. Mandatory for SAM3; SAM1 masks are already solid.

Unlike SAM1's per-prompt loop, SAM3 runs all prompts for the frame in a single batched forward pass (each prompt is one object), matching #643's _sam3_instance_masks. The candidate selection (:func:_pick) and the raw-score contract are identical to SAM1.

Parameters:

Name	Type	Description	Default
model		A ready Sam3TrackerModel (e.g. from :func:_load_sam3 or injected for testing).	required
processor		The matching Sam3TrackerProcessor.	required
device	str	Torch device the prompt tensors are moved to.	'cuda'
clahe	bool	Whether to CLAHE-equalize before encoding.	True
max_box_area_factor	float	Candidate-rejection factor (:func:_pick).	1.5
clahe_clip_limit	float	CLAHE clip limit.	3.0
clahe_tile_grid	Tuple[int, int]	CLAHE tile grid.	(8, 8)
cleanup_radius	int	Speckle-cleanup morphological radius (px).	3
pred_iou_min	float	Per-model nominal predicted-IoU floor (default 0.5, recalibrated; NEVER SAM1's 0.88); reported, not gated.	0.5

Methods:

Name	Description
__init__	Stash the model/processor and the (SAM3-specific) recipe knobs.
from_pretrained	Build a backend by lazily loading the gated SAM3 model + processor.
masks	Encode image once, run all prompts batched, return masks + scores.

Source code in sleap_nn/inference/sam/backends.py

class Sam3Backend(MaskBackend):
    """SAM3 (Meta SAM 3) prompted-mask backend (the ``sleap_nn[sam3]`` extra).

    Wraps a lazily loaded transformers ``Sam3TrackerModel`` + ``Sam3TrackerProcessor``
    image visual-prompt pair. Honors the same :class:`MaskBackend` surface as
    :class:`SamBackend`, but two SAM3 specifics are mandatory and NEVER shared
    with SAM1 (PLAN §2.3, harvested from #643):

    * **Recalibrated floor.** SAM3's ``iou_scores`` (predicted-IoU) are on a
      LOWER scale than SAM1 (median ~0.68 vs SAM1's ~0.95). SAM1's ``0.88`` floor
      applied verbatim would drop ~100% of SAM3 masks as a pure calibration
      artifact, so the per-model :attr:`pred_iou_min` defaults to ``0.5``
      (~SAM1's ``0.88`` in percentile terms), never SAM1's ``0.88``. As with SAM1
      the raw chosen-candidate score is reported, not gated on.
    * **Speckle cleanup.** Raw SAM3 masks are speckly/fragmented (median ~14
      connected components per mask vs SAM1's 1), with ~97% of the area in the
      keypoint-connected component. The speckle is cosmetic, so each chosen mask
      is passed through :func:`_cleanup_speckle` (morphological open + close +
      keep-keypoint-component, -> median 1 component, ~97% area retained) before
      it is returned. Mandatory for SAM3; SAM1 masks are already solid.

    Unlike SAM1's per-prompt loop, SAM3 runs **all prompts for the frame in a
    single batched forward pass** (each prompt is one object), matching #643's
    ``_sam3_instance_masks``. The candidate selection (:func:`_pick`) and the
    raw-score contract are identical to SAM1.

    Args:
        model: A ready ``Sam3TrackerModel`` (e.g. from :func:`_load_sam3` or
            injected for testing).
        processor: The matching ``Sam3TrackerProcessor``.
        device: Torch device the prompt tensors are moved to.
        clahe: Whether to CLAHE-equalize before encoding.
        max_box_area_factor: Candidate-rejection factor (:func:`_pick`).
        clahe_clip_limit: CLAHE clip limit.
        clahe_tile_grid: CLAHE tile grid.
        cleanup_radius: Speckle-cleanup morphological radius (px).
        pred_iou_min: Per-model nominal predicted-IoU floor (default ``0.5``,
            recalibrated; NEVER SAM1's ``0.88``); reported, not gated.
    """

    #: SAM3's recalibrated predicted-IoU floor. NEVER SAM1's ``0.88`` (SAM3's
    #: predicted-IoU is on a lower scale; the ``0.5`` value is ~SAM1's ``0.88`` in
    #: percentile terms). Reported as the per-model score, not gated on.
    pred_iou_min: float = 0.5

    def __init__(
        self,
        model,
        processor,
        device: str = "cuda",
        clahe: bool = True,
        max_box_area_factor: float = 1.5,
        clahe_clip_limit: float = 3.0,
        clahe_tile_grid: Tuple[int, int] = (8, 8),
        cleanup_radius: int = 3,
        pred_iou_min: float = 0.5,
    ) -> None:
        """Stash the model/processor and the (SAM3-specific) recipe knobs.

        The SAM1-shared recipe defaults match :class:`SamBackend`
        (``max_box_area_factor=1.5``, ``clahe_clip_limit=3.0``,
        ``clahe_tile_grid=(8, 8)``). The SAM3-specific defaults are
        ``cleanup_radius=3`` (the morphological open + close radius (px) for the
        mandatory speckle cleanup) and ``pred_iou_min=0.5`` (the recalibrated
        floor; NEVER SAM1's ``0.88``, since SAM3's predicted-IoU is on a lower
        scale).
        """
        self.model = model
        self.processor = processor
        self.device = str(device)
        self.clahe = bool(clahe)
        self.max_box_area_factor = float(max_box_area_factor)
        self.clahe_clip_limit = float(clahe_clip_limit)
        self.clahe_tile_grid = tuple(clahe_tile_grid)
        self.cleanup_radius = int(cleanup_radius)
        self.pred_iou_min = float(pred_iou_min)

    @classmethod
    def from_pretrained(
        cls,
        model_id: str = "facebook/sam3",
        device: str = "cuda",
        **kwargs,
    ) -> "Sam3Backend":
        """Build a backend by lazily loading the gated SAM3 model + processor.

        Args:
            model_id: Hugging Face model id (default ``"facebook/sam3"``).
            device: Torch device for the model.
            **kwargs: Forwarded to :class:`Sam3Backend` (e.g. ``clahe``).

        Returns:
            A ready :class:`Sam3Backend`.

        Raises:
            ImportError: If ``transformers`` (with SAM3 support) is absent.
        """
        model, processor = _load_sam3(model_id=model_id, device=device)
        return cls(model, processor, device=device, **kwargs)

    def masks(
        self, image: np.ndarray, prompts: Sequence[SamPrompt]
    ) -> Tuple[List[np.ndarray], List[float]]:
        """Encode ``image`` once, run all prompts batched, return masks + scores.

        Mirrors #643's ``_sam3_instance_masks``: one batched forward pass over all
        prompts (each prompt is an object), :func:`_pick` to choose a candidate,
        :func:`_cleanup_speckle` to de-fragment, and the raw chosen predicted-IoU
        as the per-mask score.

        Args:
            image: ``(H, W)`` grayscale (or ``(H, W, C)``) image / crop.
            prompts: Per-instance :class:`SamPrompt` in ``image`` pixel space.

        Returns:
            ``(masks, scores)`` with one ``(H, W)`` boolean mask + raw
            predicted-IoU per prompt (on SAM3's lower scale). An empty prompt list
            returns ``([], [])``.
        """
        import torch

        prompts = list(prompts)
        img = np.asarray(image)
        if img.ndim == 3:
            img = img[..., 0]
        img = np.ascontiguousarray(img).astype(np.uint8)
        h, w = img.shape[:2]

        out_masks: List[np.ndarray] = [np.zeros((h, w), bool) for _ in prompts]
        out_scores: List[float] = [0.0 for _ in prompts]
        if not prompts:
            return out_masks, out_scores

        rgb = _to_3ch_clahe(
            img,
            clahe=self.clahe,
            clahe_clip_limit=self.clahe_clip_limit,
            clahe_tile_grid=self.clahe_tile_grid,
        )

        # Build batched per-object prompts (one image, each prompt an object).
        # Mirror SAM1 (``SamBackend.masks``): forward only a prompt's real
        # ``box`` — never ``reject_box``, which exists solely for the
        # candidate-rejection heuristic (:func:`_pick`). Point-only prompts
        # (e.g. ``centroid`` mode, ``box is None``) carry no box; feeding them
        # ``reject_box`` would hand SAM3 a whole-frame "segment everything" box
        # and make SAM3 diverge from SAM1 on identical input.
        obj_points: List[List[List[float]]] = []
        obj_labels: List[List[int]] = []
        obj_boxes: List[List[float]] = []
        any_box = False
        for prompt in prompts:
            pc = prompt.point_coords
            pl = prompt.point_labels
            if pc is not None and len(pc):
                obj_points.append([[float(x), float(y)] for x, y in pc])
                labels = [int(v) for v in pl] if pl is not None else [1] * len(pc)
                obj_labels.append(labels)
            else:
                obj_points.append([])
                obj_labels.append([])
            if prompt.box is not None:
                obj_boxes.append([float(v) for v in np.asarray(prompt.box).reshape(4)])
                any_box = True
            else:
                obj_boxes.append([])

        processor_kwargs = dict(
            images=rgb,
            input_points=[obj_points],
            input_labels=[obj_labels],
            return_tensors="pt",
        )
        # Only forward boxes when a prompt actually has one (pose / box modes);
        # a frame of point-only prompts forwards no boxes at all.
        if any_box:
            processor_kwargs["input_boxes"] = [obj_boxes]
        inputs = self.processor(**processor_kwargs).to(self.device)
        with torch.no_grad():
            out = self.model(**inputs, multimask_output=True)
        post = self.processor.post_process_masks(
            out.pred_masks, original_sizes=inputs["original_sizes"], binarize=True
        )[
            0
        ]  # (n_obj, n_cand, H, W) bool
        post = np.asarray(post.cpu().numpy()).astype(bool)
        scores = np.asarray(out.iou_scores.float().cpu().numpy()[0])  # (n_obj, n_cand)

        for j, prompt in enumerate(prompts):
            cand_masks = post[j]
            cand_scores = scores[j]
            b = _pick(
                cand_masks, cand_scores, prompt.reject_box, self.max_box_area_factor
            )
            mask = _cleanup_speckle(
                cand_masks[b], _cleanup_seed(prompt), self.cleanup_radius
            )
            out_masks[j] = mask.astype(bool)
            out_scores[j] = float(cand_scores[b])

        for m in out_masks:
            if m.shape[:2] != (h, w):
                raise ValueError(
                    f"SAM3 returned a {m.shape} mask for a {(h, w)} image."
                )
        return out_masks, out_scores

__init__(model, processor, device='cuda', clahe=True, max_box_area_factor=1.5, clahe_clip_limit=3.0, clahe_tile_grid=(8, 8), cleanup_radius=3, pred_iou_min=0.5)

Stash the model/processor and the (SAM3-specific) recipe knobs.

The SAM1-shared recipe defaults match :class:SamBackend (max_box_area_factor=1.5, clahe_clip_limit=3.0, clahe_tile_grid=(8, 8)). The SAM3-specific defaults are cleanup_radius=3 (the morphological open + close radius (px) for the mandatory speckle cleanup) and pred_iou_min=0.5 (the recalibrated floor; NEVER SAM1's 0.88, since SAM3's predicted-IoU is on a lower scale).

Source code in sleap_nn/inference/sam/backends.py

def __init__(
    self,
    model,
    processor,
    device: str = "cuda",
    clahe: bool = True,
    max_box_area_factor: float = 1.5,
    clahe_clip_limit: float = 3.0,
    clahe_tile_grid: Tuple[int, int] = (8, 8),
    cleanup_radius: int = 3,
    pred_iou_min: float = 0.5,
) -> None:
    """Stash the model/processor and the (SAM3-specific) recipe knobs.

    The SAM1-shared recipe defaults match :class:`SamBackend`
    (``max_box_area_factor=1.5``, ``clahe_clip_limit=3.0``,
    ``clahe_tile_grid=(8, 8)``). The SAM3-specific defaults are
    ``cleanup_radius=3`` (the morphological open + close radius (px) for the
    mandatory speckle cleanup) and ``pred_iou_min=0.5`` (the recalibrated
    floor; NEVER SAM1's ``0.88``, since SAM3's predicted-IoU is on a lower
    scale).
    """
    self.model = model
    self.processor = processor
    self.device = str(device)
    self.clahe = bool(clahe)
    self.max_box_area_factor = float(max_box_area_factor)
    self.clahe_clip_limit = float(clahe_clip_limit)
    self.clahe_tile_grid = tuple(clahe_tile_grid)
    self.cleanup_radius = int(cleanup_radius)
    self.pred_iou_min = float(pred_iou_min)

from_pretrained(model_id='facebook/sam3', device='cuda', **kwargs) classmethod

Build a backend by lazily loading the gated SAM3 model + processor.

Parameters:

Name	Type	Description	Default
model_id	str	Hugging Face model id (default "facebook/sam3").	'facebook/sam3'
device	str	Torch device for the model.	'cuda'
**kwargs		Forwarded to :class:Sam3Backend (e.g. clahe).	{}

Returns:

Type	Description
'Sam3Backend'	A ready :class:Sam3Backend.

Raises:

Type	Description
ImportError	If transformers (with SAM3 support) is absent.

Source code in sleap_nn/inference/sam/backends.py

@classmethod
def from_pretrained(
    cls,
    model_id: str = "facebook/sam3",
    device: str = "cuda",
    **kwargs,
) -> "Sam3Backend":
    """Build a backend by lazily loading the gated SAM3 model + processor.

    Args:
        model_id: Hugging Face model id (default ``"facebook/sam3"``).
        device: Torch device for the model.
        **kwargs: Forwarded to :class:`Sam3Backend` (e.g. ``clahe``).

    Returns:
        A ready :class:`Sam3Backend`.

    Raises:
        ImportError: If ``transformers`` (with SAM3 support) is absent.
    """
    model, processor = _load_sam3(model_id=model_id, device=device)
    return cls(model, processor, device=device, **kwargs)

masks(image, prompts)

Encode image once, run all prompts batched, return masks + scores.

Mirrors #643's _sam3_instance_masks: one batched forward pass over all prompts (each prompt is an object), :func:_pick to choose a candidate, :func:_cleanup_speckle to de-fragment, and the raw chosen predicted-IoU as the per-mask score.

Parameters:

Name	Type	Description	Default
image	ndarray	(H, W) grayscale (or (H, W, C)) image / crop.	required
prompts	Sequence[SamPrompt]	Per-instance :class:SamPrompt in image pixel space.	required

Returns:

Type	Description
Tuple[List[ndarray], List[float]]	(masks, scores) with one (H, W) boolean mask + raw predicted-IoU per prompt (on SAM3's lower scale). An empty prompt list returns ([], []).

Source code in sleap_nn/inference/sam/backends.py

def masks(
    self, image: np.ndarray, prompts: Sequence[SamPrompt]
) -> Tuple[List[np.ndarray], List[float]]:
    """Encode ``image`` once, run all prompts batched, return masks + scores.

    Mirrors #643's ``_sam3_instance_masks``: one batched forward pass over all
    prompts (each prompt is an object), :func:`_pick` to choose a candidate,
    :func:`_cleanup_speckle` to de-fragment, and the raw chosen predicted-IoU
    as the per-mask score.

    Args:
        image: ``(H, W)`` grayscale (or ``(H, W, C)``) image / crop.
        prompts: Per-instance :class:`SamPrompt` in ``image`` pixel space.

    Returns:
        ``(masks, scores)`` with one ``(H, W)`` boolean mask + raw
        predicted-IoU per prompt (on SAM3's lower scale). An empty prompt list
        returns ``([], [])``.
    """
    import torch

    prompts = list(prompts)
    img = np.asarray(image)
    if img.ndim == 3:
        img = img[..., 0]
    img = np.ascontiguousarray(img).astype(np.uint8)
    h, w = img.shape[:2]

    out_masks: List[np.ndarray] = [np.zeros((h, w), bool) for _ in prompts]
    out_scores: List[float] = [0.0 for _ in prompts]
    if not prompts:
        return out_masks, out_scores

    rgb = _to_3ch_clahe(
        img,
        clahe=self.clahe,
        clahe_clip_limit=self.clahe_clip_limit,
        clahe_tile_grid=self.clahe_tile_grid,
    )

    # Build batched per-object prompts (one image, each prompt an object).
    # Mirror SAM1 (``SamBackend.masks``): forward only a prompt's real
    # ``box`` — never ``reject_box``, which exists solely for the
    # candidate-rejection heuristic (:func:`_pick`). Point-only prompts
    # (e.g. ``centroid`` mode, ``box is None``) carry no box; feeding them
    # ``reject_box`` would hand SAM3 a whole-frame "segment everything" box
    # and make SAM3 diverge from SAM1 on identical input.
    obj_points: List[List[List[float]]] = []
    obj_labels: List[List[int]] = []
    obj_boxes: List[List[float]] = []
    any_box = False
    for prompt in prompts:
        pc = prompt.point_coords
        pl = prompt.point_labels
        if pc is not None and len(pc):
            obj_points.append([[float(x), float(y)] for x, y in pc])
            labels = [int(v) for v in pl] if pl is not None else [1] * len(pc)
            obj_labels.append(labels)
        else:
            obj_points.append([])
            obj_labels.append([])
        if prompt.box is not None:
            obj_boxes.append([float(v) for v in np.asarray(prompt.box).reshape(4)])
            any_box = True
        else:
            obj_boxes.append([])

    processor_kwargs = dict(
        images=rgb,
        input_points=[obj_points],
        input_labels=[obj_labels],
        return_tensors="pt",
    )
    # Only forward boxes when a prompt actually has one (pose / box modes);
    # a frame of point-only prompts forwards no boxes at all.
    if any_box:
        processor_kwargs["input_boxes"] = [obj_boxes]
    inputs = self.processor(**processor_kwargs).to(self.device)
    with torch.no_grad():
        out = self.model(**inputs, multimask_output=True)
    post = self.processor.post_process_masks(
        out.pred_masks, original_sizes=inputs["original_sizes"], binarize=True
    )[
        0
    ]  # (n_obj, n_cand, H, W) bool
    post = np.asarray(post.cpu().numpy()).astype(bool)
    scores = np.asarray(out.iou_scores.float().cpu().numpy()[0])  # (n_obj, n_cand)

    for j, prompt in enumerate(prompts):
        cand_masks = post[j]
        cand_scores = scores[j]
        b = _pick(
            cand_masks, cand_scores, prompt.reject_box, self.max_box_area_factor
        )
        mask = _cleanup_speckle(
            cand_masks[b], _cleanup_seed(prompt), self.cleanup_radius
        )
        out_masks[j] = mask.astype(bool)
        out_scores[j] = float(cand_scores[b])

    for m in out_masks:
        if m.shape[:2] != (h, w):
            raise ValueError(
                f"SAM3 returned a {m.shape} mask for a {(h, w)} image."
            )
    return out_masks, out_scores

SamBackend

Bases: MaskBackend

SAM1 (ViT-H) prompted-mask backend (the sleap_nn[sam] extra).

Wraps a lazily loaded segment_anything.SamPredictor. For one frame: CLAHE-equalize + 3-channel replicate, set_image once, then per prompt call predict(..., multimask_output=True) and select via :func:_pick. The raw SAM predicted-IoU of the chosen candidate is the mask score (PLAN §2.3 — store the raw per-model score; no drop-gate).

Parameters:

Name	Type	Description	Default
predictor		A ready SamPredictor (e.g. from :func:_load_sam_predictor or injected for testing). When None, :meth:from_checkpoint builds one.	required
clahe	bool	Whether to CLAHE-equalize before encoding.	True
max_box_area_factor	float	Candidate-rejection factor (:func:_pick).	1.5
clahe_clip_limit	float	CLAHE clip limit.	3.0
clahe_tile_grid	Tuple[int, int]	CLAHE tile grid.	(8, 8)
pred_iou_min	float	Nominal predicted-IoU floor carried for parity with SAM3; SAM1 reports the raw score and does not gate on it.	0.88

Methods:

Name	Description
__init__	Stash the predictor and the (model-specific) recipe knobs.
from_checkpoint	Build a backend by lazily loading a SAM checkpoint.
masks	Encode image once, run each prompt, return masks + raw scores.

Source code in sleap_nn/inference/sam/backends.py

class SamBackend(MaskBackend):
    """SAM1 (ViT-H) prompted-mask backend (the ``sleap_nn[sam]`` extra).

    Wraps a lazily loaded ``segment_anything.SamPredictor``. For one frame:
    CLAHE-equalize + 3-channel replicate, ``set_image`` once, then per prompt
    call ``predict(..., multimask_output=True)`` and select via :func:`_pick`.
    The raw SAM predicted-IoU of the chosen candidate is the mask score (PLAN
    §2.3 — store the raw per-model score; no drop-gate).

    Args:
        predictor: A ready ``SamPredictor`` (e.g. from :func:`_load_sam_predictor`
            or injected for testing). When ``None``, :meth:`from_checkpoint`
            builds one.
        clahe: Whether to CLAHE-equalize before encoding.
        max_box_area_factor: Candidate-rejection factor (:func:`_pick`).
        clahe_clip_limit: CLAHE clip limit.
        clahe_tile_grid: CLAHE tile grid.
        pred_iou_min: Nominal predicted-IoU floor carried for parity with SAM3;
            SAM1 reports the raw score and does not gate on it.
    """

    def __init__(
        self,
        predictor,
        clahe: bool = True,
        max_box_area_factor: float = 1.5,
        clahe_clip_limit: float = 3.0,
        clahe_tile_grid: Tuple[int, int] = (8, 8),
        pred_iou_min: float = 0.88,
    ) -> None:
        """Stash the predictor and the (model-specific) recipe knobs.

        The recipe defaults are the locked SAM1 values (harvested from #642 /
        exp-07; PLAN §1): ``max_box_area_factor=1.5`` drops candidates whose area
        exceeds ``1.5 * box-area`` (kills SAM's over-confident whole-arena
        candidate, see :func:`_pick`); ``clahe_clip_limit=3.0`` /
        ``clahe_tile_grid=(8, 8)`` are the CLAHE parameters applied to the
        grayscale image before encoding; ``pred_iou_min=0.88`` is SAM1's nominal
        predicted-IoU floor, reported (not gated) and carried for SAM3 parity.
        """
        self.predictor = predictor
        self.clahe = bool(clahe)
        self.max_box_area_factor = float(max_box_area_factor)
        self.clahe_clip_limit = float(clahe_clip_limit)
        self.clahe_tile_grid = tuple(clahe_tile_grid)
        self.pred_iou_min = float(pred_iou_min)

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint: str,
        model_type: str = "vit_h",
        device: str = "cuda",
        **kwargs,
    ) -> "SamBackend":
        """Build a backend by lazily loading a SAM checkpoint.

        Args:
            checkpoint: Path to the SAM checkpoint.
            model_type: SAM model registry key.
            device: Torch device for the model.
            **kwargs: Forwarded to :class:`SamBackend` (e.g. ``clahe``).

        Returns:
            A ready :class:`SamBackend`.
        """
        predictor = _load_sam_predictor(
            checkpoint, model_type=model_type, device=device
        )
        return cls(predictor, **kwargs)

    def masks(
        self, image: np.ndarray, prompts: Sequence[SamPrompt]
    ) -> Tuple[List[np.ndarray], List[float]]:
        """Encode ``image`` once, run each prompt, return masks + raw scores.

        Args:
            image: ``(H, W)`` grayscale (or ``(H, W, C)``) image / crop.
            prompts: Per-instance :class:`SamPrompt` in ``image`` pixel space.

        Returns:
            ``(masks, scores)`` with one ``(H, W)`` boolean mask + raw
            predicted-IoU per prompt. An empty prompt list returns ``([], [])``.
        """
        img = np.asarray(image)
        if img.ndim == 3:
            img = img[..., 0]
        h, w = img.shape[:2]
        rgb = _to_3ch_clahe(
            img,
            clahe=self.clahe,
            clahe_clip_limit=self.clahe_clip_limit,
            clahe_tile_grid=self.clahe_tile_grid,
        )
        self.predictor.set_image(rgb)

        out_masks: List[np.ndarray] = []
        out_scores: List[float] = []
        for prompt in prompts:
            point_coords = (
                prompt.point_coords.astype(np.float32)
                if prompt.point_coords is not None
                else None
            )
            point_labels = (
                prompt.point_labels.astype(np.int32)
                if prompt.point_labels is not None
                else None
            )
            box = prompt.box.astype(np.float32) if prompt.box is not None else None
            ms, sc, _ = self.predictor.predict(
                point_coords=point_coords,
                point_labels=point_labels,
                box=box,
                multimask_output=True,
            )
            b = _pick(ms, sc, prompt.reject_box, self.max_box_area_factor)
            out_masks.append(ms[b].astype(bool))
            out_scores.append(float(sc[b]))
        # Defensive: a degenerate single-px prompt could yield a (h, w) mismatch
        # only if SAM is given a wrong-size image; guard the contract here.
        for m in out_masks:
            if m.shape[:2] != (h, w):
                raise ValueError(f"SAM returned a {m.shape} mask for a {(h, w)} image.")
        return out_masks, out_scores

__init__(predictor, clahe=True, max_box_area_factor=1.5, clahe_clip_limit=3.0, clahe_tile_grid=(8, 8), pred_iou_min=0.88)

Stash the predictor and the (model-specific) recipe knobs.

The recipe defaults are the locked SAM1 values (harvested from #642 / exp-07; PLAN §1): max_box_area_factor=1.5 drops candidates whose area exceeds 1.5 * box-area (kills SAM's over-confident whole-arena candidate, see :func:_pick); clahe_clip_limit=3.0 / clahe_tile_grid=(8, 8) are the CLAHE parameters applied to the grayscale image before encoding; pred_iou_min=0.88 is SAM1's nominal predicted-IoU floor, reported (not gated) and carried for SAM3 parity.

Source code in sleap_nn/inference/sam/backends.py

def __init__(
    self,
    predictor,
    clahe: bool = True,
    max_box_area_factor: float = 1.5,
    clahe_clip_limit: float = 3.0,
    clahe_tile_grid: Tuple[int, int] = (8, 8),
    pred_iou_min: float = 0.88,
) -> None:
    """Stash the predictor and the (model-specific) recipe knobs.

    The recipe defaults are the locked SAM1 values (harvested from #642 /
    exp-07; PLAN §1): ``max_box_area_factor=1.5`` drops candidates whose area
    exceeds ``1.5 * box-area`` (kills SAM's over-confident whole-arena
    candidate, see :func:`_pick`); ``clahe_clip_limit=3.0`` /
    ``clahe_tile_grid=(8, 8)`` are the CLAHE parameters applied to the
    grayscale image before encoding; ``pred_iou_min=0.88`` is SAM1's nominal
    predicted-IoU floor, reported (not gated) and carried for SAM3 parity.
    """
    self.predictor = predictor
    self.clahe = bool(clahe)
    self.max_box_area_factor = float(max_box_area_factor)
    self.clahe_clip_limit = float(clahe_clip_limit)
    self.clahe_tile_grid = tuple(clahe_tile_grid)
    self.pred_iou_min = float(pred_iou_min)

from_checkpoint(checkpoint, model_type='vit_h', device='cuda', **kwargs) classmethod

Build a backend by lazily loading a SAM checkpoint.

Parameters:

Name	Type	Description	Default
checkpoint	str	Path to the SAM checkpoint.	required
model_type	str	SAM model registry key.	'vit_h'
device	str	Torch device for the model.	'cuda'
**kwargs		Forwarded to :class:SamBackend (e.g. clahe).	{}

Returns:

Type	Description
'SamBackend'	A ready :class:SamBackend.

Source code in sleap_nn/inference/sam/backends.py

@classmethod
def from_checkpoint(
    cls,
    checkpoint: str,
    model_type: str = "vit_h",
    device: str = "cuda",
    **kwargs,
) -> "SamBackend":
    """Build a backend by lazily loading a SAM checkpoint.

    Args:
        checkpoint: Path to the SAM checkpoint.
        model_type: SAM model registry key.
        device: Torch device for the model.
        **kwargs: Forwarded to :class:`SamBackend` (e.g. ``clahe``).

    Returns:
        A ready :class:`SamBackend`.
    """
    predictor = _load_sam_predictor(
        checkpoint, model_type=model_type, device=device
    )
    return cls(predictor, **kwargs)

masks(image, prompts)

Encode image once, run each prompt, return masks + raw scores.

Parameters:

Name	Type	Description	Default
image	ndarray	(H, W) grayscale (or (H, W, C)) image / crop.	required
prompts	Sequence[SamPrompt]	Per-instance :class:SamPrompt in image pixel space.	required

Returns:

Type	Description
Tuple[List[ndarray], List[float]]	(masks, scores) with one (H, W) boolean mask + raw predicted-IoU per prompt. An empty prompt list returns ([], []).

Source code in sleap_nn/inference/sam/backends.py

def masks(
    self, image: np.ndarray, prompts: Sequence[SamPrompt]
) -> Tuple[List[np.ndarray], List[float]]:
    """Encode ``image`` once, run each prompt, return masks + raw scores.

    Args:
        image: ``(H, W)`` grayscale (or ``(H, W, C)``) image / crop.
        prompts: Per-instance :class:`SamPrompt` in ``image`` pixel space.

    Returns:
        ``(masks, scores)`` with one ``(H, W)`` boolean mask + raw
        predicted-IoU per prompt. An empty prompt list returns ``([], [])``.
    """
    img = np.asarray(image)
    if img.ndim == 3:
        img = img[..., 0]
    h, w = img.shape[:2]
    rgb = _to_3ch_clahe(
        img,
        clahe=self.clahe,
        clahe_clip_limit=self.clahe_clip_limit,
        clahe_tile_grid=self.clahe_tile_grid,
    )
    self.predictor.set_image(rgb)

    out_masks: List[np.ndarray] = []
    out_scores: List[float] = []
    for prompt in prompts:
        point_coords = (
            prompt.point_coords.astype(np.float32)
            if prompt.point_coords is not None
            else None
        )
        point_labels = (
            prompt.point_labels.astype(np.int32)
            if prompt.point_labels is not None
            else None
        )
        box = prompt.box.astype(np.float32) if prompt.box is not None else None
        ms, sc, _ = self.predictor.predict(
            point_coords=point_coords,
            point_labels=point_labels,
            box=box,
            multimask_output=True,
        )
        b = _pick(ms, sc, prompt.reject_box, self.max_box_area_factor)
        out_masks.append(ms[b].astype(bool))
        out_scores.append(float(sc[b]))
    # Defensive: a degenerate single-px prompt could yield a (h, w) mismatch
    # only if SAM is given a wrong-size image; guard the contract here.
    for m in out_masks:
        if m.shape[:2] != (h, w):
            raise ValueError(f"SAM returned a {m.shape} mask for a {(h, w)} image.")
    return out_masks, out_scores

SamPrompt dataclass

A built SAM prompt for one instance.

Attributes:

Name	Type	Description
point_coords	Optional[ndarray]	(n, 2) float32 positive-point xy, or None when the prompt is box-only.
point_labels	Optional[ndarray]	(n,) int32 labels (all 1 — positive; automatic prompting uses no negatives, PLAN §2.2). None iff point_coords is None.
box	Optional[ndarray]	[x0, y0, x1, y1] float32 box prompt, or None when the prompt is point-only.
reject_box	ndarray	[x0, y0, x1, y1] float32 box used only by the candidate-rejection heuristic (:func:backends._pick) — never passed to SAM. Always populated so :func:backends._pick can size-reject the whole-arena candidate even in point-only modes.
mode	str	The originating mode tag ("pose" / "centroid" / "box"), carried for diagnostics / overlays.

Source code in sleap_nn/inference/sam/prompts.py

@dataclass
class SamPrompt:
    """A built SAM prompt for one instance.

    Attributes:
        point_coords: ``(n, 2)`` float32 positive-point xy, or ``None`` when the
            prompt is box-only.
        point_labels: ``(n,)`` int32 labels (all ``1`` — positive; automatic
            prompting uses no negatives, PLAN §2.2). ``None`` iff
            ``point_coords`` is ``None``.
        box: ``[x0, y0, x1, y1]`` float32 box prompt, or ``None`` when the prompt
            is point-only.
        reject_box: ``[x0, y0, x1, y1]`` float32 box used **only** by the
            candidate-rejection heuristic (:func:`backends._pick`) — never passed
            to SAM. Always populated so :func:`backends._pick` can size-reject the
            whole-arena candidate even in point-only modes.
        mode: The originating mode tag (``"pose"`` / ``"centroid"`` / ``"box"``),
            carried for diagnostics / overlays.
    """

    point_coords: Optional[np.ndarray]
    point_labels: Optional[np.ndarray]
    box: Optional[np.ndarray]
    reject_box: np.ndarray
    mode: str

SamSegmentationLayer

Full-frame SAM mask producer (pose / centroid / box prompts).

Operates on in-memory sio.LabeledFrame content (image + pose/centroid instances), not on a torch model — there is no trained net here. For each frame it encodes the image once via the backend, builds one prompt per instance, and emits per-frame Outputs.pred_masks dicts that the standard Outputs.to_masks path packages into sio.PredictedSegmentationMask. Full-frame masks use identity scale/offset (the whole-frame representation the P1 prototype produced).

Parameters:

Name	Type	Description	Default
backend	MaskBackend	A :class:MaskBackend (SAM1 here; SAM3 later).	required
prompt_mode	str	One of "pose" / "centroid" / "box". "pose" applies the L3 product rule (pose-if-visible-else-centroid-point).	'pose'
anchor_ind	Optional[int]	Optional skeleton node index used as the centroid anchor for prompt_mode="centroid"; None uses the mean of visible keypoints.	None
disjointify_masks	bool	When True and a frame has >=2 instances, make the per-frame masks disjoint via keypoint-Voronoi (harvested #642). Default False (single-instance is the common case; disjointify is a multi-instance refinement).	False

Methods:

Name	Description
__init__	Stash the backend and prompt knobs.
masks_for_frame	Produce one pred_masks dict per posed instance for a frame.
predict_labels	Build pred_masks for every labeled frame of a sio.Labels.

Source code in sleap_nn/inference/sam/mask_layer.py

class SamSegmentationLayer:
    """Full-frame SAM mask producer (pose / centroid / box prompts).

    Operates on in-memory ``sio.LabeledFrame`` content (image + pose/centroid
    instances), not on a torch model — there is no trained net here. For each
    frame it encodes the image once via the backend, builds one prompt per
    instance, and emits per-frame ``Outputs.pred_masks`` dicts that the standard
    ``Outputs.to_masks`` path packages into ``sio.PredictedSegmentationMask``.
    Full-frame masks use identity ``scale``/``offset`` (the whole-frame
    representation the P1 prototype produced).

    Args:
        backend: A :class:`MaskBackend` (SAM1 here; SAM3 later).
        prompt_mode: One of ``"pose"`` / ``"centroid"`` / ``"box"``. ``"pose"``
            applies the L3 product rule (pose-if-visible-else-centroid-point).
        anchor_ind: Optional skeleton node index used as the centroid anchor for
            ``prompt_mode="centroid"``; ``None`` uses the mean of visible
            keypoints.
        disjointify_masks: When ``True`` and a frame has >=2 instances, make the
            per-frame masks disjoint via keypoint-Voronoi (harvested #642).
            Default ``False`` (single-instance is the common case; disjointify is
            a multi-instance refinement).
    """

    def __init__(
        self,
        backend: MaskBackend,
        prompt_mode: str = "pose",
        anchor_ind: Optional[int] = None,
        disjointify_masks: bool = False,
    ) -> None:
        """Stash the backend and prompt knobs."""
        if prompt_mode not in ("pose", "centroid", "box"):
            raise ValueError(
                f"SamSegmentationLayer prompt_mode must be 'pose'/'centroid'/'box', "
                f"got {prompt_mode!r}."
            )
        self.backend = backend
        self.prompt_mode = prompt_mode
        self.anchor_ind = anchor_ind
        self.disjointify_masks = bool(disjointify_masks)

    def _instance_centroid(self, kpts_vis: np.ndarray, inst) -> Optional[np.ndarray]:
        """Anchor point for an instance: anchor node if set/visible, else mean."""
        if self.anchor_ind is not None:
            pts = np.asarray(inst.numpy()[:, :2], dtype=np.float32)
            if 0 <= self.anchor_ind < len(pts):
                a = pts[self.anchor_ind]
                if np.isfinite(a).all():
                    return a.astype(np.float32)
        if len(kpts_vis) > 0:
            return kpts_vis.mean(0).astype(np.float32)
        return None

    def masks_for_frame(self, image, instances: Sequence) -> List[dict]:
        """Produce one ``pred_masks`` dict per posed instance for a frame.

        Args:
            image: The frame image (``(H, W)`` / ``(H, W, C)`` / ``(C, H, W)``).
            instances: The frame's ``sio.PredictedInstance`` (or ``sio.Instance``)
                pose/centroid instances. Instances with no visible keypoints (and
                no usable centroid) are skipped.

        Returns:
            A list of ``pred_masks`` dicts ``{"mask", "score", "scale",
            "offset", "instance", "track", "tracking_score"}`` — full-frame masks
            with identity scale/offset and ``instance``/``track`` populated when
            the source instance carries them (PLAN L8).
        """
        gray = _frame_gray(image)
        h, w = gray.shape
        prompts: List[SamPrompt] = []
        kept = []  # (instance, kpts_vis)
        for inst in instances:
            kpts = np.asarray(inst.numpy()[:, :2], dtype=np.float32)
            kpts_vis = visible_keypoints(kpts)
            centroid = self._instance_centroid(kpts_vis, inst)
            try:
                prompt = prompt_for_instance(
                    self.prompt_mode,
                    (h, w),
                    keypoints=kpts_vis if len(kpts_vis) else None,
                    centroid=centroid,
                )
            except ValueError:
                # No usable prompt source for this instance — skip it.
                continue
            prompts.append(prompt)
            kept.append((inst, kpts_vis))

        if not prompts:
            return []

        masks, scores = self.backend.masks(gray, prompts)

        if self.disjointify_masks and len(masks) >= 2:
            from sleap_nn.inference.sam.backends import disjointify

            masks = disjointify(masks, [kv[1] for kv in kept])

        out: List[dict] = []
        for (inst, _kpts), mask, score in zip(kept, masks, scores):
            if mask is None or not mask.any():
                continue
            out.append(
                {
                    "mask": np.ascontiguousarray(mask, dtype=bool),
                    "score": float(score),
                    "scale": (1.0, 1.0),
                    "offset": (0.0, 0.0),
                    "instance": inst if _is_predicted(inst) else None,
                    "track": getattr(inst, "track", None),
                    "tracking_score": _tracking_score(inst),
                }
            )
        return out

    def predict_labels(self, labels) -> "List[List[dict]]":
        """Build ``pred_masks`` for every labeled frame of a ``sio.Labels``.

        Args:
            labels: The source ``sio.Labels`` with pose/centroid instances + image
                data (used as the prompt source).

        Returns:
            A list (one entry per labeled frame) of the frame's ``pred_masks``
            dicts; frames are index-aligned to ``labels.labeled_frames``.
        """
        return [
            self.masks_for_frame(lf.image, lf.instances) for lf in labels.labeled_frames
        ]

__init__(backend, prompt_mode='pose', anchor_ind=None, disjointify_masks=False)

Stash the backend and prompt knobs.

Source code in sleap_nn/inference/sam/mask_layer.py

def __init__(
    self,
    backend: MaskBackend,
    prompt_mode: str = "pose",
    anchor_ind: Optional[int] = None,
    disjointify_masks: bool = False,
) -> None:
    """Stash the backend and prompt knobs."""
    if prompt_mode not in ("pose", "centroid", "box"):
        raise ValueError(
            f"SamSegmentationLayer prompt_mode must be 'pose'/'centroid'/'box', "
            f"got {prompt_mode!r}."
        )
    self.backend = backend
    self.prompt_mode = prompt_mode
    self.anchor_ind = anchor_ind
    self.disjointify_masks = bool(disjointify_masks)

masks_for_frame(image, instances)

Produce one pred_masks dict per posed instance for a frame.

Parameters:

Name	Type	Description	Default
image		The frame image ((H, W) / (H, W, C) / (C, H, W)).	required
instances	Sequence	The frame's sio.PredictedInstance (or sio.Instance) pose/centroid instances. Instances with no visible keypoints (and no usable centroid) are skipped.	required

Returns:

Type	Description
List[dict]	A list of pred_masks dicts {"mask", "score", "scale", "offset", "instance", "track", "tracking_score"} — full-frame masks with identity scale/offset and instance/track populated when the source instance carries them (PLAN L8).

Source code in sleap_nn/inference/sam/mask_layer.py

def masks_for_frame(self, image, instances: Sequence) -> List[dict]:
    """Produce one ``pred_masks`` dict per posed instance for a frame.

    Args:
        image: The frame image (``(H, W)`` / ``(H, W, C)`` / ``(C, H, W)``).
        instances: The frame's ``sio.PredictedInstance`` (or ``sio.Instance``)
            pose/centroid instances. Instances with no visible keypoints (and
            no usable centroid) are skipped.

    Returns:
        A list of ``pred_masks`` dicts ``{"mask", "score", "scale",
        "offset", "instance", "track", "tracking_score"}`` — full-frame masks
        with identity scale/offset and ``instance``/``track`` populated when
        the source instance carries them (PLAN L8).
    """
    gray = _frame_gray(image)
    h, w = gray.shape
    prompts: List[SamPrompt] = []
    kept = []  # (instance, kpts_vis)
    for inst in instances:
        kpts = np.asarray(inst.numpy()[:, :2], dtype=np.float32)
        kpts_vis = visible_keypoints(kpts)
        centroid = self._instance_centroid(kpts_vis, inst)
        try:
            prompt = prompt_for_instance(
                self.prompt_mode,
                (h, w),
                keypoints=kpts_vis if len(kpts_vis) else None,
                centroid=centroid,
            )
        except ValueError:
            # No usable prompt source for this instance — skip it.
            continue
        prompts.append(prompt)
        kept.append((inst, kpts_vis))

    if not prompts:
        return []

    masks, scores = self.backend.masks(gray, prompts)

    if self.disjointify_masks and len(masks) >= 2:
        from sleap_nn.inference.sam.backends import disjointify

        masks = disjointify(masks, [kv[1] for kv in kept])

    out: List[dict] = []
    for (inst, _kpts), mask, score in zip(kept, masks, scores):
        if mask is None or not mask.any():
            continue
        out.append(
            {
                "mask": np.ascontiguousarray(mask, dtype=bool),
                "score": float(score),
                "scale": (1.0, 1.0),
                "offset": (0.0, 0.0),
                "instance": inst if _is_predicted(inst) else None,
                "track": getattr(inst, "track", None),
                "tracking_score": _tracking_score(inst),
            }
        )
    return out

predict_labels(labels)

Build pred_masks for every labeled frame of a sio.Labels.

Parameters:

Name	Type	Description	Default
labels		The source sio.Labels with pose/centroid instances + image data (used as the prompt source).	required

Returns:

Type	Description
'List[List[dict]]'	A list (one entry per labeled frame) of the frame's pred_masks dicts; frames are index-aligned to labels.labeled_frames.

Source code in sleap_nn/inference/sam/mask_layer.py

def predict_labels(self, labels) -> "List[List[dict]]":
    """Build ``pred_masks`` for every labeled frame of a ``sio.Labels``.

    Args:
        labels: The source ``sio.Labels`` with pose/centroid instances + image
            data (used as the prompt source).

    Returns:
        A list (one entry per labeled frame) of the frame's ``pred_masks``
        dicts; frames are index-aligned to ``labels.labeled_frames``.
    """
    return [
        self.masks_for_frame(lf.image, lf.instances) for lf in labels.labeled_frames
    ]

SwapEvent dataclass

Detected identity swap.

Attributes:

Name	Type	Description
frame_idx	int	Frame where the swap was detected.
track_name	str	Name of the track that swapped.
old_sam3_id	int	Previous SAM3 object ID.
new_sam3_id	int	New SAM3 object ID after swap.

Source code in sleap_nn/inference/sam/reconciliation.py

@dataclass
class SwapEvent:
    """Detected identity swap.

    Attributes:
        frame_idx: Frame where the swap was detected.
        track_name: Name of the track that swapped.
        old_sam3_id: Previous SAM3 object ID.
        new_sam3_id: New SAM3 object ID after swap.
    """

    frame_idx: int
    track_name: str
    old_sam3_id: int
    new_sam3_id: int

TrackAssignment dataclass

A single track assignment at a frame.

Attributes:

Name	Type	Description
frame_idx	int	Frame index where assignment was made.
pose_track_name	str | None	Name of the pose's track (None if untracked).
pose_idx	int	Index of the pose in the frame's instance list.
sam3_obj_id	int	SAM3 object ID that was matched.
confidence	float	Match quality score (0-1, higher is better).
sam3_score	float	SAM3 mask detection confidence score.

Source code in sleap_nn/inference/sam/reconciliation.py

@dataclass
class TrackAssignment:
    """A single track assignment at a frame.

    Attributes:
        frame_idx: Frame index where assignment was made.
        pose_track_name: Name of the pose's track (None if untracked).
        pose_idx: Index of the pose in the frame's instance list.
        sam3_obj_id: SAM3 object ID that was matched.
        confidence: Match quality score (0-1, higher is better).
        sam3_score: SAM3 mask detection confidence score.
    """

    frame_idx: int
    pose_track_name: str | None
    pose_idx: int
    sam3_obj_id: int
    confidence: float
    sam3_score: float = 1.0

TrackNameResolver dataclass

Resolves SAM3 obj_ids to GT track names via nearest-anchor flood fill.

This class takes the sparse ID mappings from GT anchor frames and propagates them to all frames using a nearest-anchor approach. Each frame uses the mapping from its closest GT anchor frame.

Attributes:

Name	Type	Description
gt_anchors	dict[int, dict[int, str]]	Mapping of frame_idx -> {sam3_obj_id: track_name} at GT frames.
fallback_names	dict[int, str]	Optional mapping of sam3_obj_id -> name for objects without GT matches (e.g., from initial prompt).

Example

resolver = TrackNameResolver.from_reconciler(reconciler)

Get track name for a specific frame and object

name = resolver.get_track_name(frame_idx=150, sam3_obj_id=1)

Get all mappings for batch processing

all_mappings = resolver.resolve_all_frames(total_frames=1000)

Methods:

Name	Description
__post_init__	Cache sorted anchor frames for efficient lookup.
from_id_map	Create resolver from an existing ID map.
from_reconciler	Create resolver from an IDReconciler with accumulated assignments.
get_all_sam3_obj_ids	Get all unique SAM3 object IDs from GT anchors.
get_all_track_names	Get all unique track names from GT anchors.
get_anchor_frames	Get sorted list of GT anchor frame indices.
get_anchor_source	Get the anchor frame and propagation direction for a frame.
get_canonical_mapping	Get a canonical sam3_obj_id -> track_name mapping.
get_mapping_at_frame	Get the sam3_obj_id -> track_name mapping for a frame.
get_track_name	Get track name for a SAM3 obj_id at a given frame.
resolve_all_frames	Get resolved mappings for all frames.

Source code in sleap_nn/inference/sam/reconciliation.py

@dataclass
class TrackNameResolver:
    """Resolves SAM3 obj_ids to GT track names via nearest-anchor flood fill.

    This class takes the sparse ID mappings from GT anchor frames and propagates
    them to all frames using a nearest-anchor approach. Each frame uses the
    mapping from its closest GT anchor frame.

    Attributes:
        gt_anchors: Mapping of frame_idx -> {sam3_obj_id: track_name} at GT frames.
        fallback_names: Optional mapping of sam3_obj_id -> name for objects without
            GT matches (e.g., from initial prompt).

    Example:
        >>> resolver = TrackNameResolver.from_reconciler(reconciler)
        >>> # Get track name for a specific frame and object
        >>> name = resolver.get_track_name(frame_idx=150, sam3_obj_id=1)
        >>> # Get all mappings for batch processing
        >>> all_mappings = resolver.resolve_all_frames(total_frames=1000)
    """

    gt_anchors: dict[int, dict[int, str]] = field(default_factory=dict)
    fallback_names: dict[int, str] = field(default_factory=dict)
    _anchor_frames: list[int] = field(default_factory=list, repr=False)

    def __post_init__(self):
        """Cache sorted anchor frames for efficient lookup."""
        self._anchor_frames = sorted(self.gt_anchors.keys())

    @classmethod
    def from_reconciler(
        cls,
        reconciler: IDReconciler,
        fallback_names: dict[int, str] | None = None,
    ) -> "TrackNameResolver":
        """Create resolver from an IDReconciler with accumulated assignments.

        Args:
            reconciler: IDReconciler that has processed GT frames.
            fallback_names: Optional mapping for objects without GT matches.

        Returns:
            TrackNameResolver initialized with the reconciler's ID map.
        """
        return cls(
            gt_anchors=reconciler.build_id_map(),
            fallback_names=fallback_names or {},
        )

    @classmethod
    def from_id_map(
        cls,
        id_map: dict[int, dict[int, str]],
        fallback_names: dict[int, str] | None = None,
    ) -> "TrackNameResolver":
        """Create resolver from an existing ID map.

        Args:
            id_map: Mapping of frame_idx -> {sam3_obj_id: track_name}.
            fallback_names: Optional mapping for objects without GT matches.

        Returns:
            TrackNameResolver initialized with the ID map.
        """
        return cls(
            gt_anchors=id_map,
            fallback_names=fallback_names or {},
        )

    def _find_nearest_anchor(self, frame_idx: int) -> int | None:
        """Find the nearest GT anchor frame to a given frame.

        Args:
            frame_idx: The frame index to find nearest anchor for.

        Returns:
            The frame index of the nearest anchor, or None if no anchors exist.
        """
        if not self._anchor_frames:
            return None

        # Binary search would be faster for large anchor lists,
        # but linear is fine for typical use cases (< 100 anchors)
        return min(self._anchor_frames, key=lambda a: abs(frame_idx - a))

    def get_mapping_at_frame(self, frame_idx: int) -> dict[int, str]:
        """Get the sam3_obj_id -> track_name mapping for a frame.

        Uses the mapping from the nearest GT anchor frame.

        Args:
            frame_idx: The frame index to get mapping for.

        Returns:
            Dictionary mapping sam3_obj_id to track_name.
            Returns empty dict if no GT anchors exist.
        """
        nearest = self._find_nearest_anchor(frame_idx)
        if nearest is None:
            return {}
        return self.gt_anchors[nearest]

    def get_track_name(
        self,
        frame_idx: int,
        sam3_obj_id: int,
        default: str | None = None,
    ) -> str:
        """Get track name for a SAM3 obj_id at a given frame.

        Uses the mapping from the nearest GT anchor frame. Falls back to
        fallback_names, then to a generated name.

        Args:
            frame_idx: The frame index.
            sam3_obj_id: The SAM3 object ID.
            default: Optional default name if not found. If None, generates
                a name like "track_{sam3_obj_id}".

        Returns:
            The resolved track name.
        """
        mapping = self.get_mapping_at_frame(frame_idx)

        if sam3_obj_id in mapping:
            return mapping[sam3_obj_id]

        if sam3_obj_id in self.fallback_names:
            return self.fallback_names[sam3_obj_id]

        if default is not None:
            return default

        return f"track_{sam3_obj_id}"

    def resolve_all_frames(
        self,
        total_frames: int,
    ) -> dict[int, dict[int, str]]:
        """Get resolved mappings for all frames.

        Args:
            total_frames: Total number of frames in the video.

        Returns:
            Dictionary mapping frame_idx -> {sam3_obj_id: track_name}.
            Empty frames (no mapping) are not included in the result.
        """
        if not self._anchor_frames:
            return {}

        result: dict[int, dict[int, str]] = {}
        for frame_idx in range(total_frames):
            nearest = self._find_nearest_anchor(frame_idx)
            if nearest is not None:
                result[frame_idx] = self.gt_anchors[nearest]

        return result

    def get_anchor_frames(self) -> list[int]:
        """Get sorted list of GT anchor frame indices.

        Returns:
            List of frame indices where GT anchors exist.
        """
        return list(self._anchor_frames)

    def get_all_track_names(self) -> set[str]:
        """Get all unique track names from GT anchors.

        Returns:
            Set of all track names found in GT mappings.
        """
        names: set[str] = set()
        for mapping in self.gt_anchors.values():
            names.update(mapping.values())
        return names

    def get_all_sam3_obj_ids(self) -> set[int]:
        """Get all unique SAM3 object IDs from GT anchors.

        Returns:
            Set of all SAM3 object IDs found in GT mappings.
        """
        obj_ids: set[int] = set()
        for mapping in self.gt_anchors.values():
            obj_ids.update(mapping.keys())
        return obj_ids

    def get_canonical_mapping(self) -> dict[int, str]:
        """Get a canonical sam3_obj_id -> track_name mapping.

        Returns a single global mapping from SAM3 object IDs to track names.
        For objects that appear in multiple anchors, uses the name from the
        first anchor frame.

        This is useful for writers that need a single consistent mapping
        (like BBoxWriter and SegmentationWriter) rather than per-frame mappings.

        Returns:
            Dictionary mapping sam3_obj_id to track_name.
        """
        canonical: dict[int, str] = {}

        # Iterate anchors in frame order to get consistent "first seen" names
        for frame_idx in self._anchor_frames:
            mapping = self.gt_anchors[frame_idx]
            for obj_id, name in mapping.items():
                if obj_id not in canonical:
                    canonical[obj_id] = name

        return canonical

    def get_anchor_source(self, frame_idx: int) -> tuple[int | None, str]:
        """Get the anchor frame and propagation direction for a frame.

        Useful for debugging and visualization.

        Args:
            frame_idx: The frame index to check.

        Returns:
            Tuple of (anchor_frame_idx, direction) where direction is one of:
            - "anchor": frame_idx is a GT anchor
            - "forward": propagated forward from an earlier anchor
            - "backward": propagated backward from a later anchor
            - "none": no anchors exist
        """
        if not self._anchor_frames:
            return (None, "none")

        nearest = self._find_nearest_anchor(frame_idx)
        if nearest is None:
            return (None, "none")

        if frame_idx == nearest:
            return (nearest, "anchor")
        elif frame_idx > nearest:
            return (nearest, "forward")
        else:
            return (nearest, "backward")

__post_init__()

Cache sorted anchor frames for efficient lookup.

Source code in sleap_nn/inference/sam/reconciliation.py

def __post_init__(self):
    """Cache sorted anchor frames for efficient lookup."""
    self._anchor_frames = sorted(self.gt_anchors.keys())

from_id_map(id_map, fallback_names=None) classmethod

Create resolver from an existing ID map.

Parameters:

Name	Type	Description	Default
id_map	dict[int, dict[int, str]]	Mapping of frame_idx -> {sam3_obj_id: track_name}.	required
fallback_names	dict[int, str] | None	Optional mapping for objects without GT matches.	None

Returns:

Type	Description
TrackNameResolver	TrackNameResolver initialized with the ID map.

Source code in sleap_nn/inference/sam/reconciliation.py

@classmethod
def from_id_map(
    cls,
    id_map: dict[int, dict[int, str]],
    fallback_names: dict[int, str] | None = None,
) -> "TrackNameResolver":
    """Create resolver from an existing ID map.

    Args:
        id_map: Mapping of frame_idx -> {sam3_obj_id: track_name}.
        fallback_names: Optional mapping for objects without GT matches.

    Returns:
        TrackNameResolver initialized with the ID map.
    """
    return cls(
        gt_anchors=id_map,
        fallback_names=fallback_names or {},
    )

from_reconciler(reconciler, fallback_names=None) classmethod

Create resolver from an IDReconciler with accumulated assignments.

Parameters:

Name	Type	Description	Default
reconciler	IDReconciler	IDReconciler that has processed GT frames.	required
fallback_names	dict[int, str] | None	Optional mapping for objects without GT matches.	None

Returns:

Type	Description
TrackNameResolver	TrackNameResolver initialized with the reconciler's ID map.

Source code in sleap_nn/inference/sam/reconciliation.py

@classmethod
def from_reconciler(
    cls,
    reconciler: IDReconciler,
    fallback_names: dict[int, str] | None = None,
) -> "TrackNameResolver":
    """Create resolver from an IDReconciler with accumulated assignments.

    Args:
        reconciler: IDReconciler that has processed GT frames.
        fallback_names: Optional mapping for objects without GT matches.

    Returns:
        TrackNameResolver initialized with the reconciler's ID map.
    """
    return cls(
        gt_anchors=reconciler.build_id_map(),
        fallback_names=fallback_names or {},
    )

get_all_sam3_obj_ids()

Get all unique SAM3 object IDs from GT anchors.

Returns:

Type	Description
set[int]	Set of all SAM3 object IDs found in GT mappings.

Source code in sleap_nn/inference/sam/reconciliation.py

def get_all_sam3_obj_ids(self) -> set[int]:
    """Get all unique SAM3 object IDs from GT anchors.

    Returns:
        Set of all SAM3 object IDs found in GT mappings.
    """
    obj_ids: set[int] = set()
    for mapping in self.gt_anchors.values():
        obj_ids.update(mapping.keys())
    return obj_ids

get_all_track_names()

Get all unique track names from GT anchors.

Returns:

Type	Description
set[str]	Set of all track names found in GT mappings.

Source code in sleap_nn/inference/sam/reconciliation.py

def get_all_track_names(self) -> set[str]:
    """Get all unique track names from GT anchors.

    Returns:
        Set of all track names found in GT mappings.
    """
    names: set[str] = set()
    for mapping in self.gt_anchors.values():
        names.update(mapping.values())
    return names

get_anchor_frames()

Get sorted list of GT anchor frame indices.

Returns:

Type	Description
list[int]	List of frame indices where GT anchors exist.

Source code in sleap_nn/inference/sam/reconciliation.py

def get_anchor_frames(self) -> list[int]:
    """Get sorted list of GT anchor frame indices.

    Returns:
        List of frame indices where GT anchors exist.
    """
    return list(self._anchor_frames)

get_anchor_source(frame_idx)

Get the anchor frame and propagation direction for a frame.

Useful for debugging and visualization.

Parameters:

Name	Type	Description	Default
frame_idx	int	The frame index to check.	required

Returns:

Type	Description
tuple[int | None, str]	Tuple of (anchor_frame_idx, direction) where direction is one of: - "anchor": frame_idx is a GT anchor - "forward": propagated forward from an earlier anchor - "backward": propagated backward from a later anchor - "none": no anchors exist

Source code in sleap_nn/inference/sam/reconciliation.py

def get_anchor_source(self, frame_idx: int) -> tuple[int | None, str]:
    """Get the anchor frame and propagation direction for a frame.

    Useful for debugging and visualization.

    Args:
        frame_idx: The frame index to check.

    Returns:
        Tuple of (anchor_frame_idx, direction) where direction is one of:
        - "anchor": frame_idx is a GT anchor
        - "forward": propagated forward from an earlier anchor
        - "backward": propagated backward from a later anchor
        - "none": no anchors exist
    """
    if not self._anchor_frames:
        return (None, "none")

    nearest = self._find_nearest_anchor(frame_idx)
    if nearest is None:
        return (None, "none")

    if frame_idx == nearest:
        return (nearest, "anchor")
    elif frame_idx > nearest:
        return (nearest, "forward")
    else:
        return (nearest, "backward")

get_canonical_mapping()

Get a canonical sam3_obj_id -> track_name mapping.

Returns a single global mapping from SAM3 object IDs to track names. For objects that appear in multiple anchors, uses the name from the first anchor frame.

This is useful for writers that need a single consistent mapping (like BBoxWriter and SegmentationWriter) rather than per-frame mappings.

Returns:

Type	Description
dict[int, str]	Dictionary mapping sam3_obj_id to track_name.

Source code in sleap_nn/inference/sam/reconciliation.py

def get_canonical_mapping(self) -> dict[int, str]:
    """Get a canonical sam3_obj_id -> track_name mapping.

    Returns a single global mapping from SAM3 object IDs to track names.
    For objects that appear in multiple anchors, uses the name from the
    first anchor frame.

    This is useful for writers that need a single consistent mapping
    (like BBoxWriter and SegmentationWriter) rather than per-frame mappings.

    Returns:
        Dictionary mapping sam3_obj_id to track_name.
    """
    canonical: dict[int, str] = {}

    # Iterate anchors in frame order to get consistent "first seen" names
    for frame_idx in self._anchor_frames:
        mapping = self.gt_anchors[frame_idx]
        for obj_id, name in mapping.items():
            if obj_id not in canonical:
                canonical[obj_id] = name

    return canonical

get_mapping_at_frame(frame_idx)

Get the sam3_obj_id -> track_name mapping for a frame.

Uses the mapping from the nearest GT anchor frame.

Parameters:

Name	Type	Description	Default
frame_idx	int	The frame index to get mapping for.	required

Returns:

Type	Description
dict[int, str]	Dictionary mapping sam3_obj_id to track_name. Returns empty dict if no GT anchors exist.

Source code in sleap_nn/inference/sam/reconciliation.py

def get_mapping_at_frame(self, frame_idx: int) -> dict[int, str]:
    """Get the sam3_obj_id -> track_name mapping for a frame.

    Uses the mapping from the nearest GT anchor frame.

    Args:
        frame_idx: The frame index to get mapping for.

    Returns:
        Dictionary mapping sam3_obj_id to track_name.
        Returns empty dict if no GT anchors exist.
    """
    nearest = self._find_nearest_anchor(frame_idx)
    if nearest is None:
        return {}
    return self.gt_anchors[nearest]

get_track_name(frame_idx, sam3_obj_id, default=None)

Get track name for a SAM3 obj_id at a given frame.

Uses the mapping from the nearest GT anchor frame. Falls back to fallback_names, then to a generated name.

Parameters:

Name	Type	Description	Default
frame_idx	int	The frame index.	required
sam3_obj_id	int	The SAM3 object ID.	required
default	str | None	Optional default name if not found. If None, generates a name like "track_{sam3_obj_id}".	None

Returns:

Type	Description
str	The resolved track name.

Source code in sleap_nn/inference/sam/reconciliation.py

def get_track_name(
    self,
    frame_idx: int,
    sam3_obj_id: int,
    default: str | None = None,
) -> str:
    """Get track name for a SAM3 obj_id at a given frame.

    Uses the mapping from the nearest GT anchor frame. Falls back to
    fallback_names, then to a generated name.

    Args:
        frame_idx: The frame index.
        sam3_obj_id: The SAM3 object ID.
        default: Optional default name if not found. If None, generates
            a name like "track_{sam3_obj_id}".

    Returns:
        The resolved track name.
    """
    mapping = self.get_mapping_at_frame(frame_idx)

    if sam3_obj_id in mapping:
        return mapping[sam3_obj_id]

    if sam3_obj_id in self.fallback_names:
        return self.fallback_names[sam3_obj_id]

    if default is not None:
        return default

    return f"track_{sam3_obj_id}"

resolve_all_frames(total_frames)

Get resolved mappings for all frames.

Parameters:

Name	Type	Description	Default
total_frames	int	Total number of frames in the video.	required

Returns:

Type	Description
dict[int, dict[int, str]]	Dictionary mapping frame_idx -> {sam3_obj_id: track_name}. Empty frames (no mapping) are not included in the result.

Source code in sleap_nn/inference/sam/reconciliation.py

def resolve_all_frames(
    self,
    total_frames: int,
) -> dict[int, dict[int, str]]:
    """Get resolved mappings for all frames.

    Args:
        total_frames: Total number of frames in the video.

    Returns:
        Dictionary mapping frame_idx -> {sam3_obj_id: track_name}.
        Empty frames (no mapping) are not included in the result.
    """
    if not self._anchor_frames:
        return {}

    result: dict[int, dict[int, str]] = {}
    for frame_idx in range(total_frames):
        nearest = self._find_nearest_anchor(frame_idx)
        if nearest is not None:
            result[frame_idx] = self.gt_anchors[nearest]

    return result

default_match_predicate(pose, mask, ctx)

Default match predicate: require at least 1 keypoint inside mask.

Source code in sleap_nn/inference/sam/reconciliation.py

def default_match_predicate(
    pose: "sio.Instance", mask: np.ndarray, ctx: MatchContext
) -> bool:
    """Default match predicate: require at least 1 keypoint inside mask."""
    return ctx.keypoints_inside >= 1

get_mask_backend(mask_backend, *, sam_checkpoint=None, sam_model_type='vit_h', sam3_model_id='facebook/sam3', device='cuda', **kwargs)

Build a mask backend by explicit name (no default; PLAN L2).

Parameters:

Name	Type	Description	Default
mask_backend	Optional[str]	The backend name. "sam" builds a SAM1 :class:~.backends.SamBackend; "sam3" builds a SAM3 :class:~.backends.Sam3Backend. There is no default — the caller must name one.	required
sam_checkpoint	Optional[str]	Path to the SAM1 checkpoint (required for "sam").	None
sam_model_type	str	SAM1 model registry key.	'vit_h'
sam3_model_id	str	Hugging Face model id for the SAM3 path (gated; the sleap_nn[sam3] extra). Used only for "sam3".	'facebook/sam3'
device	str	Torch device for the model.	'cuda'
**kwargs		Forwarded to the backend constructor (e.g. clahe).	{}

Returns:

Type	Description
MaskBackend	A ready :class:~.backends.MaskBackend.

Raises:

Type	Description
ValueError	If mask_backend is not a registered name.
ImportError	If mask_backend == "sam3" and transformers (with SAM3 support) is not installed — with an actionable install/auth message.

Source code in sleap_nn/inference/sam/__init__.py

def get_mask_backend(
    mask_backend: Optional[str],
    *,
    sam_checkpoint: Optional[str] = None,
    sam_model_type: str = "vit_h",
    sam3_model_id: str = "facebook/sam3",
    device: str = "cuda",
    **kwargs,
) -> MaskBackend:
    """Build a mask backend by **explicit** name (no default; PLAN L2).

    Args:
        mask_backend: The backend name. ``"sam"`` builds a SAM1
            :class:`~.backends.SamBackend`; ``"sam3"`` builds a SAM3
            :class:`~.backends.Sam3Backend`. There is no default — the caller must
            name one.
        sam_checkpoint: Path to the SAM1 checkpoint (required for ``"sam"``).
        sam_model_type: SAM1 model registry key.
        sam3_model_id: Hugging Face model id for the SAM3 path (gated; the
            ``sleap_nn[sam3]`` extra). Used only for ``"sam3"``.
        device: Torch device for the model.
        **kwargs: Forwarded to the backend constructor (e.g. ``clahe``).

    Returns:
        A ready :class:`~.backends.MaskBackend`.

    Raises:
        ValueError: If ``mask_backend`` is not a registered name.
        ImportError: If ``mask_backend == "sam3"`` and ``transformers`` (with SAM3
            support) is not installed — with an actionable install/auth message.
    """
    if mask_backend is None:
        raise ValueError(
            "mask_backend is required and has no default; pass one of "
            f"{MASK_BACKENDS} (PLAN L2)."
        )
    name = str(mask_backend).lower()
    if name == "sam":
        return SamBackend.from_checkpoint(
            sam_checkpoint,
            model_type=sam_model_type,
            device=device,
            **kwargs,
        )
    if name == "sam3":
        return Sam3Backend.from_pretrained(
            model_id=sam3_model_id,
            device=device,
            **kwargs,
        )
    raise ValueError(
        f"Unknown mask_backend {mask_backend!r}; expected one of {MASK_BACKENDS}."
    )

require_centroid_proximity(max_dist=100.0)

Create predicate requiring pose centroid near mask centroid.

Parameters:

Name	Type	Description	Default
max_dist	float	Maximum allowed distance between centroids in pixels.	100.0

Returns:

Type	Description
MatchPredicate	A MatchPredicate function.

Source code in sleap_nn/inference/sam/reconciliation.py

def require_centroid_proximity(max_dist: float = 100.0) -> MatchPredicate:
    """Create predicate requiring pose centroid near mask centroid.

    Args:
        max_dist: Maximum allowed distance between centroids in pixels.

    Returns:
        A MatchPredicate function.
    """

    def predicate(pose: "sio.Instance", mask: np.ndarray, ctx: MatchContext) -> bool:
        coords = pose.numpy()
        pose_centroid = np.nanmean(coords, axis=0)
        if np.any(np.isnan(pose_centroid)):
            return False
        dist = np.linalg.norm(pose_centroid - np.array(ctx.mask_centroid))
        return float(dist) <= max_dist

    return predicate

require_min_fraction_inside(min_frac=0.5)

Create predicate requiring minimum fraction of keypoints inside mask.

Parameters:

Name	Type	Description	Default
min_frac	float	Minimum fraction (0-1) of visible keypoints inside mask.	0.5

Returns:

Type	Description
MatchPredicate	A MatchPredicate function.

Source code in sleap_nn/inference/sam/reconciliation.py

def require_min_fraction_inside(min_frac: float = 0.5) -> MatchPredicate:
    """Create predicate requiring minimum fraction of keypoints inside mask.

    Args:
        min_frac: Minimum fraction (0-1) of visible keypoints inside mask.

    Returns:
        A MatchPredicate function.
    """

    def predicate(pose: "sio.Instance", mask: np.ndarray, ctx: MatchContext) -> bool:
        if ctx.keypoints_visible == 0:
            return False
        return ctx.keypoints_inside / ctx.keypoints_visible >= min_frac

    return predicate

require_min_keypoints_inside(min_count=3)

Create predicate requiring minimum keypoints inside mask.

Parameters:

Name	Type	Description	Default
min_count	int	Minimum number of keypoints required inside mask.	3

Returns:

Type	Description
MatchPredicate	A MatchPredicate function.

Source code in sleap_nn/inference/sam/reconciliation.py

def require_min_keypoints_inside(min_count: int = 3) -> MatchPredicate:
    """Create predicate requiring minimum keypoints inside mask.

    Args:
        min_count: Minimum number of keypoints required inside mask.

    Returns:
        A MatchPredicate function.
    """

    def predicate(pose: "sio.Instance", mask: np.ndarray, ctx: MatchContext) -> bool:
        return ctx.keypoints_inside >= min_count

    return predicate

require_reasonable_mask_area(min_area=1000, max_area=500000)

Create predicate requiring mask area within bounds.

Parameters:

Name	Type	Description	Default
min_area	int	Minimum mask area in pixels.	1000
max_area	int	Maximum mask area in pixels.	500000

Returns:

Type	Description
MatchPredicate	A MatchPredicate function.

Source code in sleap_nn/inference/sam/reconciliation.py

def require_reasonable_mask_area(
    min_area: int = 1000, max_area: int = 500000
) -> MatchPredicate:
    """Create predicate requiring mask area within bounds.

    Args:
        min_area: Minimum mask area in pixels.
        max_area: Maximum mask area in pixels.

    Returns:
        A MatchPredicate function.
    """

    def predicate(pose: "sio.Instance", mask: np.ndarray, ctx: MatchContext) -> bool:
        return min_area <= ctx.mask_area <= max_area

    return predicate

run_sam_segmentation(source, mask_backend, *, prompt_mode='pose', sam_checkpoint=None, sam_model_type='vit_h', sam3_model_id='facebook/sam3', device='cuda', anchor_ind=None, disjointify_masks=False, backend=None, output_path=None, overlay_path=None, frames=None, clean_empty_frames=False, embed='false', restore_source_videos=False)

Predict per-instance masks for a pose .slp with a SAM backend.

Loads (or accepts) a sio.Labels whose frames carry pose/centroid instances, runs the chosen backend with the chosen prompt mode, attaches one sio.PredictedSegmentationMask per instance (raw score + instance= / track= populated, PLAN L8), and returns a new sio.Labels. The instances are retained alongside the masks (correction needs the pose).

Parameters:

Name	Type	Description	Default
source		A path to a pose .slp/.pkg.slp (with image data) or an in-memory sio.Labels.	required
mask_backend	str	Explicit backend name (PLAN L2): "sam" / "sam3".	required
prompt_mode	str	"pose" / "centroid" / "box" (full-frame).	'pose'
sam_checkpoint	Optional[str]	SAM1 checkpoint path (required for "sam" unless a pre-built backend is passed).	None
sam_model_type	str	SAM1 model registry key.	'vit_h'
sam3_model_id	str	Hugging Face model id for the gated SAM3 path ("sam3").	'facebook/sam3'
device	str	Torch device for the model.	'cuda'
anchor_ind	Optional[int]	Optional centroid anchor node index for "centroid".	None
disjointify_masks	bool	Make per-frame masks disjoint when >=2 instances.	False
backend	Optional[MaskBackend]	A pre-built :class:~.backends.MaskBackend to use directly (skips loading); when given, mask_backend is still validated for the name but the checkpoint/device args are ignored.	None
output_path	Optional[str]	Optional .slp path to save the result to. Saved like the regular prediction path (labels.save): the embedding policy is configurable via embed and defaults to "false" — images are not re-embedded, and (with restore_source_videos also defaulting to False) the output backreferences the input file itself, so a .pkg.slp input stays matchable without depending on a pre-embedding source video that's often not on disk.	None
overlay_path	Optional[str]	Optional path to write a review overlay PNG of the first frame.	None
frames	Optional[Sequence[int]]	Optional frame indices (matched against lf.frame_idx) to restrict masking to; None masks every labeled frame. SAM encoding is the slow step, so subsetting here avoids unrequested compute.	None
clean_empty_frames	bool	If True, drop fully-empty output frames (no instances and no masks) before saving/returning, mirroring the regular prediction path's --no_empty_frames. A frame that has poses but no mask is NOT empty (its instances are retained) and is kept.	False
embed	Union[str, bool]	Image-embedding policy for the .slp output, one of "false" (the default; never embed, backreference source media), "true" (embed images into a self-contained .pkg.slp-style file), or "auto" (embed iff the input was itself an embedded .pkg.slp). A bool passes through unchanged.	'false'
restore_source_videos	bool	On a non-embedding save, False (the default) keeps references to the input .pkg.slp file(s) — the pixels are already there, and the pre-embedding source video is often unavailable. True instead restores references to the original pre-embedding source video files, when recorded. Maps to sleap-io's restore_original_videos and is ignored when embedding.	False

Returns:

Type	Description
	A new sio.Labels with per-frame PredictedSegmentationMask (and the original pose instances retained).

Source code in sleap_nn/inference/sam/__init__.py

def run_sam_segmentation(
    source,
    mask_backend: str,
    *,
    prompt_mode: str = "pose",
    sam_checkpoint: Optional[str] = None,
    sam_model_type: str = "vit_h",
    sam3_model_id: str = "facebook/sam3",
    device: str = "cuda",
    anchor_ind: Optional[int] = None,
    disjointify_masks: bool = False,
    backend: Optional[MaskBackend] = None,
    output_path: Optional[str] = None,
    overlay_path: Optional[str] = None,
    frames: Optional[Sequence[int]] = None,
    clean_empty_frames: bool = False,
    embed: Union[str, bool] = "false",
    restore_source_videos: bool = False,
):
    """Predict per-instance masks for a pose ``.slp`` with a SAM backend.

    Loads (or accepts) a ``sio.Labels`` whose frames carry pose/centroid
    instances, runs the chosen backend with the chosen prompt mode, attaches one
    ``sio.PredictedSegmentationMask`` per instance (raw score + ``instance=`` /
    ``track=`` populated, PLAN L8), and returns a new ``sio.Labels``. The
    instances are retained alongside the masks (correction needs the pose).

    Args:
        source: A path to a pose ``.slp``/``.pkg.slp`` (with image data) or an
            in-memory ``sio.Labels``.
        mask_backend: **Explicit** backend name (PLAN L2): ``"sam"`` / ``"sam3"``.
        prompt_mode: ``"pose"`` / ``"centroid"`` / ``"box"`` (full-frame).
        sam_checkpoint: SAM1 checkpoint path (required for ``"sam"`` unless a
            pre-built ``backend`` is passed).
        sam_model_type: SAM1 model registry key.
        sam3_model_id: Hugging Face model id for the gated SAM3 path (``"sam3"``).
        device: Torch device for the model.
        anchor_ind: Optional centroid anchor node index for ``"centroid"``.
        disjointify_masks: Make per-frame masks disjoint when >=2 instances.
        backend: A pre-built :class:`~.backends.MaskBackend` to use directly
            (skips loading); when given, ``mask_backend`` is still validated for
            the name but the checkpoint/device args are ignored.
        output_path: Optional ``.slp`` path to save the result to. Saved like the
            regular prediction path (``labels.save``): the embedding policy is
            **configurable** via ``embed`` and defaults to ``"false"`` — images
            are not re-embedded, and (with ``restore_source_videos`` also
            defaulting to ``False``) the output backreferences the input file
            itself, so a ``.pkg.slp`` input stays matchable without depending
            on a pre-embedding source video that's often not on disk.
        overlay_path: Optional path to write a review overlay PNG of the first
            frame.
        frames: Optional frame indices (matched against ``lf.frame_idx``) to
            restrict masking to; ``None`` masks every labeled frame. SAM encoding
            is the slow step, so subsetting here avoids unrequested compute.
        clean_empty_frames: If ``True``, drop fully-empty output frames (no
            instances and no masks) before saving/returning, mirroring the
            regular prediction path's ``--no_empty_frames``. A frame that has
            poses but no mask is NOT empty (its instances are retained) and is
            kept.
        embed: Image-embedding policy for the ``.slp`` output, one of
            ``"false"`` (the default; never embed, backreference source media),
            ``"true"`` (embed images into a self-contained ``.pkg.slp``-style
            file), or ``"auto"`` (embed iff the input was itself an embedded
            ``.pkg.slp``). A bool passes through unchanged.
        restore_source_videos: On a non-embedding save, ``False`` (the default)
            keeps references to the input ``.pkg.slp`` file(s) — the pixels are
            already there, and the pre-embedding source video is often
            unavailable. ``True`` instead restores references to the original
            pre-embedding source video files, when recorded. Maps to
            sleap-io's ``restore_original_videos`` and is ignored when embedding.

    Returns:
        A new ``sio.Labels`` with per-frame ``PredictedSegmentationMask`` (and the
        original pose instances retained).
    """
    import sleap_io as sio

    from sleap_nn.inference.outputs import Outputs
    from sleap_nn.inference.sam.overlay import save_mask_overlay

    if isinstance(source, sio.Labels):
        labels = source
    else:
        labels = sio.load_slp(Path(source).expanduser().as_posix())

    if backend is None:
        backend = get_mask_backend(
            mask_backend,
            sam_checkpoint=sam_checkpoint,
            sam_model_type=sam_model_type,
            sam3_model_id=sam3_model_id,
            device=device,
        )
    elif mask_backend not in MASK_BACKENDS:
        raise ValueError(
            f"Unknown mask_backend {mask_backend!r}; expected one of {MASK_BACKENDS}."
        )

    layer = SamSegmentationLayer(
        backend,
        prompt_mode=prompt_mode,
        anchor_ind=anchor_ind,
        disjointify_masks=disjointify_masks,
    )

    if frames is not None:
        wanted = {int(f) for f in frames}
        source_lfs = [lf for lf in labels.labeled_frames if int(lf.frame_idx) in wanted]
    else:
        source_lfs = list(labels.labeled_frames)

    new_lfs = []
    for lf in source_lfs:
        frame_masks = layer.masks_for_frame(lf.image, lf.instances)
        # Always emit the frame for review continuity (§3.5b): when SAM yields no
        # mask, keep the frame (video/frame_idx/instances unchanged) with an empty
        # ``masks=[]`` rather than dropping it — the poses must survive for
        # correction, and a missing frame would silently disappear from review.
        if frame_masks:
            # Reuse the standard packaging path (build_predicted_segmentation_mask).
            masks = Outputs(pred_masks=[frame_masks]).to_masks(0)
        else:
            masks = []
        new_lfs.append(
            sio.LabeledFrame(
                video=lf.video,
                frame_idx=lf.frame_idx,
                instances=list(lf.instances),  # retain poses for correction
                masks=masks,
            )
        )

    out = sio.Labels(
        videos=list(labels.videos),
        skeletons=list(labels.skeletons),
        labeled_frames=new_lfs,
    )

    if clean_empty_frames:
        # Mirror the regular path's --no_empty_frames: drop frames with no
        # annotations. A posed-but-mask-less frame keeps its instances and is
        # NOT dropped; only fully-empty source frames are removed.
        out.clean(frames=True, skeletons=False)

    if output_path is not None:
        out_path = Path(output_path).expanduser()
        out_path.parent.mkdir(parents=True, exist_ok=True)
        # Save like the regular prediction path (``labels.save(path)``). The
        # embedding policy is configurable and defaults to ``embed="false"``: do
        # NOT re-embed images — that is large and wasteful. By default the output
        # backreferences the input file itself (``restore_source_videos=False``,
        # PRESERVE_SOURCE): a ``.pkg.slp`` input's frames stay matchable to that
        # same ``.pkg.slp``, not to a pre-embedding source video that's often not
        # on disk. The masks always serialize into the ``.slp`` regardless.
        # ``out.videos`` are the input videos, so ``source_video`` provenance is
        # intact for ``embed="auto"`` detection.
        from sleap_nn.inference.run import _resolve_embed

        out.save(
            out_path.as_posix(),
            embed=_resolve_embed(embed, out),
            restore_original_videos=restore_source_videos,
        )
    if overlay_path is not None:
        # Flag masks below the backend's per-model nominal predicted-IoU floor
        # (SAM1 0.88 / SAM3 0.5) so the review overlay surfaces low-confidence
        # masks a human should scrutinize (§3.3 — consume pred_iou_min).
        save_mask_overlay(out, overlay_path, low_score_threshold=backend.pred_iou_min)

    return out