Skip to content

mask_layer

sleap_nn.inference.sam.mask_layer

SAM mask inference layer — the producer that emits PredictedSegmentationMask.

:class:SamSegmentationLayer is the full-frame producer. Given a backend, a prompt mode, and the per-frame poses/centroids, it encodes each frame once, builds one :class:~sleap_nn.inference.sam.prompts.SamPrompt per instance, asks the backend for masks, and emits Outputs.pred_masks dicts at the correct full-frame offset/scale with instance=/track= populated (PLAN L8). Output collection / SLP packaging is free — it reuses Outputs.to_masks -> build_predicted_segmentation_mask -> labels.save exactly like every other seg layer (PLAN §2.5).

It is torch-light: it shells out to a :class:MaskBackend (SAM1 here, SAM3 later). The heavy SAM import lives in the backend.

Classes:

Name Description
SamSegmentationLayer

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

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 
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
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 
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
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 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
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
    ]