topdown_segmentation

`sleap_nn.inference.layers.topdown_segmentation` ¶

Top-down (crop-centered) instance-segmentation inference layers (#622).

Two pieces, mirroring the keypoint top-down stack:

:class:CenteredInstanceMaskLayer — the stage-2 analog of :class:~sleap_nn.inference.layers.centered_instance.CenteredInstanceLayer. Runs the trained centered_instance_segmentation model on per-instance crops and returns one boolean foreground mask per crop (at the head's output-stride resolution) instead of keypoints. The crop-resolution masks + per-crop scores are carried on Outputs for the composed layer to place.
:class:TopDownSegmentationLayer — subclasses :class:~sleap_nn.inference.layers.topdown.TopDownLayer to reuse its stage-1 (centroid) + sizematch + crop machinery verbatim, overriding only the stage-2 emission. Each crop mask is emitted into Outputs.pred_masks with the DQ5 offset/scale contract (see :meth:_run_stage_2), so decode_mask_to_image_res upsamples the crop mask to the crop's image-space size and (offset-aware) top-left pads it to land at the correct full-frame location.

Classes:

Name	Description
`CenteredInstanceMaskLayer`	Per-crop foreground-mask prediction layer (top-down stage 2).
`TopDownSegmentationLayer`	Composed centroid + per-crop-mask two-stage segmentation layer.

`CenteredInstanceMaskLayer` ¶

Bases: InferenceLayer

Per-crop foreground-mask prediction layer (top-down stage 2).

Runs a trained centered_instance_segmentation model on per-instance crops and decodes the SegmentationHead logits into one boolean foreground mask per crop. The structural twin of :class:CenteredInstanceLayer, but returns masks (on Outputs.crops as a stacked tensor + a per-crop score) rather than keypoints — the composed :class:TopDownSegmentationLayer places each crop mask back into the frame.

Parameters:

Name	Type	Description	Default
`backend`	`ModelBackend`	Runtime backend wrapping the seg Lightning module. Its `forward` returns the `SegmentationHead` LOGITS (a single tensor, wrapped as `{"output": ...}` by `TorchBackend`).	required
`output_stride`	`int`	Head map → crop-pixel stride (default 2).	required
`max_stride`	`int`	Backbone max stride; crops are padded to a multiple of it.	`1`
`fg_threshold`	`float`	Foreground probability threshold for binarization (sigmoid(logits) > fg_threshold). Default 0.5.	`0.5`
`preprocess_config / postprocess_config`		Standard knobs. The crops are already sized (extracted from the sizematched image), so only the model's own `input_scale` is applied here.	required

Methods:

Name	Description
`__init__`	Compose the layer with default configs when omitted.
`postprocess`	Decode `SegmentationHead` logits → one bool mask per crop.

Attributes:

Name	Type	Description
`warmup_input_shape`		Tiny single-channel warmup shape.

Source code in sleap_nn/inference/layers/topdown_segmentation.py

class CenteredInstanceMaskLayer(InferenceLayer):
    """Per-crop foreground-mask prediction layer (top-down stage 2).

    Runs a trained ``centered_instance_segmentation`` model on per-instance
    crops and decodes the ``SegmentationHead`` logits into one boolean
    foreground mask per crop. The structural twin of
    :class:`CenteredInstanceLayer`, but returns masks (on ``Outputs.crops`` as
    a stacked tensor + a per-crop score) rather than keypoints — the composed
    :class:`TopDownSegmentationLayer` places each crop mask back into the frame.

    Args:
        backend: Runtime backend wrapping the seg Lightning module. Its
            ``forward`` returns the ``SegmentationHead`` LOGITS (a single
            tensor, wrapped as ``{"output": ...}`` by ``TorchBackend``).
        output_stride: Head map → crop-pixel stride (default 2).
        max_stride: Backbone max stride; crops are padded to a multiple of it.
        fg_threshold: Foreground probability threshold for binarization
            (sigmoid(logits) > fg_threshold). Default 0.5.
        preprocess_config / postprocess_config: Standard knobs. The crops are
            already sized (extracted from the sizematched image), so only the
            model's own ``input_scale`` is applied here.
    """

    _HEAD_OUTPUT_KEY: str = "SegmentationHead"

    def __init__(
        self,
        backend: ModelBackend,
        output_stride: int,
        max_stride: int = 1,
        fg_threshold: float = 0.5,
        preprocess_config: Optional[PreprocessConfig] = None,
        postprocess_config: Optional[PostprocessConfig] = None,
    ) -> None:
        """Compose the layer with default configs when omitted."""
        super().__init__(
            backend=backend,
            preprocess_config=preprocess_config or PreprocessConfig(),
            postprocess_config=postprocess_config or PostprocessConfig(),
            output_stride=output_stride,
            max_stride=max_stride,
        )
        self.fg_threshold = float(fg_threshold)
        # The parent TopDownLayer.predict inspects this on the stage-2 layer to
        # decide the GT-keypoint branch; a mask layer never takes that branch.
        self.use_gt_peaks = False

    @property
    def warmup_input_shape(self):
        """Tiny single-channel warmup shape."""
        return (1, 1, 64, 64)

    def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
        """Decode ``SegmentationHead`` logits → one bool mask per crop.

        Args:
            raw_out: Backend output dict carrying the seg-head logits.
            info: Preprocessing metadata (unused for placement — the composed
                layer owns the crop→frame mapping).

        Returns:
            ``Outputs`` with ``crops`` set to the stacked per-crop boolean
            masks ``(N, 1, h, w)`` (float) at output-stride resolution and
            ``instance_scores`` set to the per-crop mean foreground
            probability ``(N, 1)``. The composed layer reads both.
        """
        logits = self._extract_confmaps(raw_out)  # (N, 1, h, w)
        probs = torch.sigmoid(logits.detach())
        masks = (probs > self.fg_threshold).float()  # (N, 1, h, w)
        # Per-crop score: mean foreground probability over the predicted mask.
        # An empty crop scores 0.0 (and is dropped by ``Outputs.to_masks``). Kept
        # on ``instance_scores`` in the (N, 1) layout the composed layer expects.
        denom = masks.sum(dim=(2, 3)).clamp(min=1.0)
        score = (probs * masks).sum(dim=(2, 3)) / denom  # (N, 1)
        return Outputs(crops=masks, instance_scores=score, preprocess_info=info)

`warmup_input_shape` `property` ¶

Tiny single-channel warmup shape.

`init(backend, output_stride, max_stride=1, fg_threshold=0.5, preprocess_config=None, postprocess_config=None)` ¶

Compose the layer with default configs when omitted.

Source code in sleap_nn/inference/layers/topdown_segmentation.py

def __init__(
    self,
    backend: ModelBackend,
    output_stride: int,
    max_stride: int = 1,
    fg_threshold: float = 0.5,
    preprocess_config: Optional[PreprocessConfig] = None,
    postprocess_config: Optional[PostprocessConfig] = None,
) -> None:
    """Compose the layer with default configs when omitted."""
    super().__init__(
        backend=backend,
        preprocess_config=preprocess_config or PreprocessConfig(),
        postprocess_config=postprocess_config or PostprocessConfig(),
        output_stride=output_stride,
        max_stride=max_stride,
    )
    self.fg_threshold = float(fg_threshold)
    # The parent TopDownLayer.predict inspects this on the stage-2 layer to
    # decide the GT-keypoint branch; a mask layer never takes that branch.
    self.use_gt_peaks = False

`postprocess(raw_out, info)` ¶

Decode SegmentationHead logits → one bool mask per crop.

Parameters:

Name	Type	Description	Default
`raw_out`	`dict`	Backend output dict carrying the seg-head logits.	required
`info`	`PreprocInfo`	Preprocessing metadata (unused for placement — the composed layer owns the crop→frame mapping).	required

Returns:

Type	Description
`Outputs`	`Outputs` with `crops` set to the stacked per-crop boolean masks `(N, 1, h, w)` (float) at output-stride resolution and `instance_scores` set to the per-crop mean foreground probability `(N, 1)`. The composed layer reads both.

Source code in sleap_nn/inference/layers/topdown_segmentation.py

def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
    """Decode ``SegmentationHead`` logits → one bool mask per crop.

    Args:
        raw_out: Backend output dict carrying the seg-head logits.
        info: Preprocessing metadata (unused for placement — the composed
            layer owns the crop→frame mapping).

    Returns:
        ``Outputs`` with ``crops`` set to the stacked per-crop boolean
        masks ``(N, 1, h, w)`` (float) at output-stride resolution and
        ``instance_scores`` set to the per-crop mean foreground
        probability ``(N, 1)``. The composed layer reads both.
    """
    logits = self._extract_confmaps(raw_out)  # (N, 1, h, w)
    probs = torch.sigmoid(logits.detach())
    masks = (probs > self.fg_threshold).float()  # (N, 1, h, w)
    # Per-crop score: mean foreground probability over the predicted mask.
    # An empty crop scores 0.0 (and is dropped by ``Outputs.to_masks``). Kept
    # on ``instance_scores`` in the (N, 1) layout the composed layer expects.
    denom = masks.sum(dim=(2, 3)).clamp(min=1.0)
    score = (probs * masks).sum(dim=(2, 3)) / denom  # (N, 1)
    return Outputs(crops=masks, instance_scores=score, preprocess_info=info)