centered_instance

sleap_nn.inference.layers.centered_instance

CenteredInstanceLayer — predicts keypoints from instance-centered crops.

Single-stage layer that runs a centered-instance model on per-instance crops and decodes keypoints. Used either standalone (testing / analysis) or composed with :class:CentroidLayer to form :class:TopDownLayer.

The use_gt_peaks=True flag skips the centered-instance model and instead matches each centroid to its nearest ground-truth instance, returning the GT keypoints. Used for top-down inference when only the centroid model is available.

The two GT fallback paths:

• :attr:CentroidLayer.use_gt_centroids — GT centroids feed cropping for a real centered_instance model.
• :attr:CenteredInstanceLayer.use_gt_peaks — GT keypoints fill stage 2 when only a centroid model is available.

Each lives on the layer that owns the role the GT data plays.

Classes:

Name	Description
CenteredInstanceLayer	Centered-instance keypoint prediction layer.

CenteredInstanceLayer

Bases: InferenceLayer

Centered-instance keypoint prediction layer.

Parameters:

Name	Type	Description	Default
backend	ModelBackend	Runtime backend wrapping the centered-instance model. Required even when use_gt_peaks=True (the layer keeps the backend interface uniform; it just doesn't call it on the GT path).	required
output_stride	int	Confmap → input-pixel stride from the head config.	required
max_stride	int	Maximum stride the model requires the input to be divisible by. Padding applied bottom-right after the preprocess input-scale resize.	1
use_gt_peaks	bool	When True, skip the model and return the GT keypoints from the nearest matched instance.	False
preprocess_config / postprocess_config		Standard knobs.	required

Methods:

Name	Description
__init__	Compose the layer with default empty configs when omitted.
postprocess	Decode confmaps → keypoints; un-scale; reshape to canonical shape.
predict	Run keypoint prediction.

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

class CenteredInstanceLayer(InferenceLayer):
    """Centered-instance keypoint prediction layer.

    Args:
        backend: Runtime backend wrapping the centered-instance model.
            Required even when ``use_gt_peaks=True`` (the layer keeps the
            backend interface uniform; it just doesn't call it on the GT
            path).
        output_stride: Confmap → input-pixel stride from the head config.
        max_stride: Maximum stride the model requires the input to be
            divisible by. Padding applied bottom-right after the
            preprocess input-scale resize.
        use_gt_peaks: When ``True``, skip the model and return the GT
            keypoints from the nearest matched instance.
        preprocess_config / postprocess_config: Standard knobs.
    """

    _HEAD_OUTPUT_KEY: str = "CenteredInstanceConfmapsHead"

    def __init__(
        self,
        backend: ModelBackend,
        output_stride: int,
        max_stride: int = 1,
        use_gt_peaks: bool = False,
        preprocess_config: Optional[PreprocessConfig] = None,
        postprocess_config: Optional[PostprocessConfig] = None,
    ) -> None:
        """Compose the layer with default empty 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.use_gt_peaks = use_gt_peaks

    # ──────────────────────────────────────────────────────────────────
    # predict(): model path or GT path
    # ──────────────────────────────────────────────────────────────────

    def predict(
        self,
        crops: ImageInput,
        centroids: Optional[torch.Tensor] = None,
        instances: Optional[torch.Tensor] = None,
        centroid_vals: Optional[torch.Tensor] = None,
    ) -> Outputs:
        """Run keypoint prediction.

        Args:
            crops: Per-instance crops, ``(N, C, cH, cW)`` or any shape
                accepted by :meth:`InferenceLayer._to_4d_float_tensor`.
                Ignored on the GT path (``use_gt_peaks=True``).
            centroids: ``(B, max_inst, 2)`` predicted centroids. Required
                on the GT path so the layer can match each one to its
                nearest GT instance.
            instances: ``(B, max_inst, n_nodes, 2)`` GT keypoints from a
                LabelsReader. Required on the GT path.
            centroid_vals: ``(B, max_inst)`` centroid confidences from the
                centroid model. Carried through as the instance score on the
                GT path (legacy parity — legacy reported ``score=centroid_val``);
                falls back to all-ones when not provided.

        Returns:
            ``Outputs`` populated with ``pred_keypoints`` and
            ``pred_peak_values`` (and optionally ``pred_confmaps``).
        """
        if self.use_gt_peaks:
            if centroids is None or instances is None:
                raise ValueError(
                    "use_gt_peaks=True requires `centroids` and `instances` "
                    "to be passed (the layer matches each centroid to its "
                    "nearest GT instance)."
                )
            return self._predict_from_gt(centroids, instances, centroid_vals)
        return super().predict(crops)

    # ──────────────────────────────────────────────────────────────────
    # GT path
    # ──────────────────────────────────────────────────────────────────

    def _predict_from_gt(
        self,
        centroids: torch.Tensor,
        instances: torch.Tensor,
        centroid_vals: Optional[torch.Tensor] = None,
    ) -> Outputs:
        """Match each centroid to its nearest GT instance; return GT keypoints.

        Mirrors the legacy ``FindInstancePeaksGroundTruth.forward`` matching:
        for each centroid, find the GT instance whose nearest keypoint to
        the centroid is closest, then emit that instance's keypoints.

        When ``centroid_vals`` is given, those confidences are reported as the
        per-instance score (and ``pred_centroid_values``), matching legacy
        ``score=centroid_val``; otherwise both fall back to all-ones.
        """
        # ``centroids``: (B, max_inst, 2) — already in image-space.
        # ``instances``: (B, max_inst, n_nodes, 2) — GT keypoints.
        B, max_inst, _ = centroids.shape
        _B, _, n_nodes, _ = instances.shape

        # Distance from each centroid to each (instance, node) pair, then
        # min over nodes → distance from centroid to its nearest keypoint
        # in each candidate GT instance. Match each centroid to argmin.
        cents = centroids.unsqueeze(2).unsqueeze(3)  # (B, max_inst, 1, 1, 2)
        insts = instances.unsqueeze(1)  # (B, 1, max_inst, n_nodes, 2)
        sq = ((cents - insts) ** 2).sum(dim=-1)  # (B, max_inst, max_inst, n_nodes)
        # NaN keypoints (missing) → infinite distance so they don't win argmin.
        sq = torch.where(torch.isnan(sq), torch.full_like(sq, float("inf")), sq)
        nearest_node_dist = sq.min(dim=-1).values  # (B, max_inst_centroid, max_inst_gt)
        match_idx = nearest_node_dist.argmin(dim=-1)  # (B, max_inst_centroid)

        # Gather matched GT instance keypoints + assign full-confidence values.
        # Allocate b_idx + matched_vals on the centroids' device so the gather
        # + ``torch.where`` below don't trip the device check on cuda / mps.
        device = centroids.device
        b_idx = torch.arange(B, device=device).view(B, 1).expand(B, max_inst)
        matched_kpts = instances[b_idx, match_idx]  # (B, max_inst, n_nodes, 2)
        matched_vals = torch.ones(B, max_inst, n_nodes, device=device)

        # Centroids that were NaN-padded shouldn't pull a real GT instance —
        # mark their matched outputs back as NaN to preserve the "no peak"
        # sentinel through the final Outputs.
        nan_centroid = torch.isnan(centroids).any(dim=-1)  # (B, max_inst)
        matched_kpts = torch.where(
            nan_centroid.unsqueeze(-1).unsqueeze(-1),
            torch.full_like(matched_kpts, float("nan")),
            matched_kpts,
        )
        matched_vals = torch.where(
            nan_centroid,
            torch.full_like(matched_vals, float("nan")),
            matched_vals,
        )

        # Report the real centroid confidence as the instance score (legacy
        # parity — legacy used score=centroid_val). Fall back to all-ones when
        # the caller didn't supply centroid_vals. NaN-padded centroid slots get
        # a NaN score so empty slots don't report a spurious value.
        if centroid_vals is not None:
            cvals = centroid_vals.to(device=device, dtype=torch.float32)
        else:
            cvals = torch.ones(B, max_inst, device=device)
        cvals = torch.where(nan_centroid, torch.full_like(cvals, float("nan")), cvals)

        return Outputs(
            pred_keypoints=matched_kpts,
            pred_peak_values=matched_vals,
            pred_centroids=centroids,
            pred_centroid_values=cvals,
            instance_scores=cvals,
        )

    # ──────────────────────────────────────────────────────────────────
    # Model path: postprocess (preprocess inherited from InferenceLayer)
    # ──────────────────────────────────────────────────────────────────

    def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
        """Decode confmaps → keypoints; un-scale; reshape to canonical shape.

        Centered-instance returns one keypoint set per crop. The Outputs
        canonical shape is ``(B, I=1, N, 2)`` where ``I=1`` because the
        crop is per-instance. Always runs the torch decode path: this layer
        is only built with a ``TorchBackend``; the exported path uses
        :class:`ExportedCenteredInstanceLayer` (no double coord ladder; #584).
        """
        confmaps = self._extract_confmaps(raw_out)
        peaks, vals = find_global_peaks(
            confmaps.detach(),
            threshold=self.postprocess_config.peak_threshold,
            refinement=self.postprocess_config.effective_refinement,
            integral_patch_size=self.postprocess_config.integral_patch_size,
        )

        peaks = undo_stride(peaks, info.output_stride)
        peaks = undo_input_scale(peaks, info.input_scale)
        peaks = undo_eff_scale(peaks, info.eff_scale)

        peaks_BIN2 = peaks.unsqueeze(1)
        vals_BIN = vals.unsqueeze(1)

        outputs = Outputs(
            pred_keypoints=peaks_BIN2,
            pred_peak_values=vals_BIN,
            preprocess_info=info,
        )
        if self.postprocess_config.return_confmaps and confmaps is not None:
            outputs = attrs.evolve(outputs, pred_confmaps=confmaps.detach())
        return outputs

__init__(backend, output_stride, max_stride=1, use_gt_peaks=False, preprocess_config=None, postprocess_config=None)

Compose the layer with default empty configs when omitted.

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

def __init__(
    self,
    backend: ModelBackend,
    output_stride: int,
    max_stride: int = 1,
    use_gt_peaks: bool = False,
    preprocess_config: Optional[PreprocessConfig] = None,
    postprocess_config: Optional[PostprocessConfig] = None,
) -> None:
    """Compose the layer with default empty 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.use_gt_peaks = use_gt_peaks

postprocess(raw_out, info)

Decode confmaps → keypoints; un-scale; reshape to canonical shape.

Centered-instance returns one keypoint set per crop. The Outputs canonical shape is (B, I=1, N, 2) where I=1 because the crop is per-instance. Always runs the torch decode path: this layer is only built with a TorchBackend; the exported path uses :class:ExportedCenteredInstanceLayer (no double coord ladder; #584).

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

def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
    """Decode confmaps → keypoints; un-scale; reshape to canonical shape.

    Centered-instance returns one keypoint set per crop. The Outputs
    canonical shape is ``(B, I=1, N, 2)`` where ``I=1`` because the
    crop is per-instance. Always runs the torch decode path: this layer
    is only built with a ``TorchBackend``; the exported path uses
    :class:`ExportedCenteredInstanceLayer` (no double coord ladder; #584).
    """
    confmaps = self._extract_confmaps(raw_out)
    peaks, vals = find_global_peaks(
        confmaps.detach(),
        threshold=self.postprocess_config.peak_threshold,
        refinement=self.postprocess_config.effective_refinement,
        integral_patch_size=self.postprocess_config.integral_patch_size,
    )

    peaks = undo_stride(peaks, info.output_stride)
    peaks = undo_input_scale(peaks, info.input_scale)
    peaks = undo_eff_scale(peaks, info.eff_scale)

    peaks_BIN2 = peaks.unsqueeze(1)
    vals_BIN = vals.unsqueeze(1)

    outputs = Outputs(
        pred_keypoints=peaks_BIN2,
        pred_peak_values=vals_BIN,
        preprocess_info=info,
    )
    if self.postprocess_config.return_confmaps and confmaps is not None:
        outputs = attrs.evolve(outputs, pred_confmaps=confmaps.detach())
    return outputs

predict(crops, centroids=None, instances=None, centroid_vals=None)

Run keypoint prediction.

Parameters:

Name	Type	Description	Default
crops	ImageInput	Per-instance crops, (N, C, cH, cW) or any shape accepted by :meth:InferenceLayer._to_4d_float_tensor. Ignored on the GT path (use_gt_peaks=True).	required
centroids	Optional[Tensor]	(B, max_inst, 2) predicted centroids. Required on the GT path so the layer can match each one to its nearest GT instance.	None
instances	Optional[Tensor]	(B, max_inst, n_nodes, 2) GT keypoints from a LabelsReader. Required on the GT path.	None
centroid_vals	Optional[Tensor]	(B, max_inst) centroid confidences from the centroid model. Carried through as the instance score on the GT path (legacy parity — legacy reported score=centroid_val); falls back to all-ones when not provided.	None

Returns:

Type	Description
Outputs	Outputs populated with pred_keypoints and pred_peak_values (and optionally pred_confmaps).

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

def predict(
    self,
    crops: ImageInput,
    centroids: Optional[torch.Tensor] = None,
    instances: Optional[torch.Tensor] = None,
    centroid_vals: Optional[torch.Tensor] = None,
) -> Outputs:
    """Run keypoint prediction.

    Args:
        crops: Per-instance crops, ``(N, C, cH, cW)`` or any shape
            accepted by :meth:`InferenceLayer._to_4d_float_tensor`.
            Ignored on the GT path (``use_gt_peaks=True``).
        centroids: ``(B, max_inst, 2)`` predicted centroids. Required
            on the GT path so the layer can match each one to its
            nearest GT instance.
        instances: ``(B, max_inst, n_nodes, 2)`` GT keypoints from a
            LabelsReader. Required on the GT path.
        centroid_vals: ``(B, max_inst)`` centroid confidences from the
            centroid model. Carried through as the instance score on the
            GT path (legacy parity — legacy reported ``score=centroid_val``);
            falls back to all-ones when not provided.

    Returns:
        ``Outputs`` populated with ``pred_keypoints`` and
        ``pred_peak_values`` (and optionally ``pred_confmaps``).
    """
    if self.use_gt_peaks:
        if centroids is None or instances is None:
            raise ValueError(
                "use_gt_peaks=True requires `centroids` and `instances` "
                "to be passed (the layer matches each centroid to its "
                "nearest GT instance)."
            )
        return self._predict_from_gt(centroids, instances, centroid_vals)
    return super().predict(crops)