backends

sleap_nn.inference.sam.backends

SAM mask backends for prompted instance segmentation (PR-A).

A backend owns the model-specific half of mask production: loading the model, preprocessing an image (grayscale -> CLAHE -> 3-channel), encoding it once, and turning a list of :class:~sleap_nn.inference.sam.prompts.SamPrompt into one boolean mask + a raw per-model score per prompt.

PR-A ships :class:SamBackend (SAM1, ViT-H, Apache-2.0, ungated; the sleap_nn[sam] extra). PR-B adds :class:Sam3Backend (Meta SAM 3, gated facebook/sam3 via transformers; the sleap_nn[sam3] extra). The model-specific recipe constants (:data:SamBackend.pred_iou_min, the candidate-rejection factor, the keypoint box margins, CLAHE) and the candidate-selection / score helpers (:func:_pick, :func:own_containment, :func:disjointify) are harvested from the closed #642 (sleap_nn/data/pseudomasks.py) and the exp-07 locked recipe, repurposed to emit a raw score rather than to drive a drop-gate (PLAN §1).

Backend selection is explicit / required (PLAN L2): there is no default mask_backend; the caller names "sam" (SAM1) or "sam3" (SAM3) and both honor the same :class:MaskBackend interface. The heavy imports (segment-anything / transformers) are lazy so the default sleap-nn install never needs either.

SAM3 specifics (NEVER shared with SAM1; harvested from the closed #643): its predicted-IoU is on a lower scale, so the per-model floor is recalibrated to :attr:Sam3Backend.pred_iou_min (0.5, not SAM1's 0.88), and its raw masks are speckly/fragmented, so each is passed through :func:_cleanup_speckle (morphological open + close + keep-keypoint-component) before it is returned.

Classes:

Name	Description
MaskBackend	Abstract prompted-mask backend (the :class:SamBackend / SAM3 interface).
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).

Functions:

Name	Description
disjointify	Make per-instance masks disjoint via keypoint-Voronoi assignment.
own_containment	Fraction of an instance's visible keypoints that fall inside mask.

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

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

disjointify(masks, kpts)

Make per-instance masks disjoint via keypoint-Voronoi assignment.

Harvested verbatim from #642 _disjointify (multi-instance only). Any pixel claimed by >=2 masks is assigned to the instance whose nearest visible keypoint is closest, so the result is exactly disjoint and each instance keeps its own keypoints (they are the Voronoi seeds).

Parameters:

Name	Type	Description	Default
masks	Sequence[ndarray]	List of (H, W) boolean masks, one per instance.	required
kpts	Sequence[ndarray]	List of (n_i, 2) visible xy keypoints, index-aligned to masks.	required

Returns:

Type	Description
List[ndarray]	List of disjoint boolean masks (a shallow copy when uncontested).

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

def disjointify(
    masks: Sequence[np.ndarray], kpts: Sequence[np.ndarray]
) -> List[np.ndarray]:
    """Make per-instance masks disjoint via keypoint-Voronoi assignment.

    Harvested verbatim from #642 ``_disjointify`` (multi-instance only). Any
    pixel claimed by >=2 masks is assigned to the instance whose nearest visible
    keypoint is closest, so the result is exactly disjoint and each instance
    keeps its own keypoints (they are the Voronoi seeds).

    Args:
        masks: List of ``(H, W)`` boolean masks, one per instance.
        kpts: List of ``(n_i, 2)`` visible xy keypoints, index-aligned to
            ``masks``.

    Returns:
        List of disjoint boolean masks (a shallow copy when uncontested).
    """
    from scipy.ndimage import distance_transform_edt

    n = len(masks)
    if n == 0:
        return []
    h, w = masks[0].shape
    stack = np.stack(masks).astype(bool)
    contested = stack.sum(0) >= 2
    if not contested.any():
        return [m.copy() for m in masks]
    dists = np.full((n, h, w), 1e9, np.float32)
    for i, ks in enumerate(kpts):
        seed = np.ones((h, w), np.uint8)
        for x, y in np.asarray(ks, dtype=np.float32).reshape(-1, 2):
            xi, yi = int(round(float(x))), int(round(float(y)))
            if 0 <= yi < h and 0 <= xi < w:
                seed[yi, xi] = 0
        if seed.min() == 0:
            dists[i] = distance_transform_edt(seed)
    owner = np.argmin(dists, 0)
    return [np.where(contested & (owner != i), False, stack[i]) for i in range(n)]

own_containment(mask, kpts, hw)

Fraction of an instance's visible keypoints that fall inside mask.

Harvested from #642 _own_containment. In the inference stack this is a score (a mask-quality signal surfaced for review), never a drop-gate.

Parameters:

Name	Type	Description	Default
mask	ndarray	(H, W) boolean mask.	required
kpts	ndarray	(n, 2) visible xy keypoints.	required
hw	Tuple[int, int]	(height, width) of mask.	required

Returns:

Type	Description
float	Containment in [0, 1] (0.0 for an empty keypoint set).

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

def own_containment(mask: np.ndarray, kpts: np.ndarray, hw: Tuple[int, int]) -> float:
    """Fraction of an instance's visible keypoints that fall inside ``mask``.

    Harvested from #642 ``_own_containment``. In the inference stack this is a
    *score* (a mask-quality signal surfaced for review), never a drop-gate.

    Args:
        mask: ``(H, W)`` boolean mask.
        kpts: ``(n, 2)`` visible xy keypoints.
        hw: ``(height, width)`` of ``mask``.

    Returns:
        Containment in ``[0, 1]`` (``0.0`` for an empty keypoint set).
    """
    kpts = np.asarray(kpts, dtype=np.float32).reshape(-1, 2)
    if len(kpts) == 0:
        return 0.0
    h, w = hw
    inside = 0
    for x, y in kpts:
        xi, yi = int(round(float(x))), int(round(float(y)))
        if 0 <= yi < h and 0 <= xi < w and mask[yi, xi]:
            inside += 1
    return inside / len(kpts)