centroid

sleap_nn.inference.layers.centroid

CentroidLayer — predicts instance centroids from a confmap model.

Single-stage layer used either standalone (centroid-only inference) or composed with :class:CenteredInstanceLayer to form :class:TopDownLayer.

The use_gt_centroids=True flag skips the centroid model and reads ground-truth centroids directly from a LabelsReader batch's "instances" field. Used for top-down inference when only the centered_instance model is available.

The two GT fallback paths live on different layers:

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

Each is independently configurable on the layer that owns the role the GT data plays.

Classes:

Name	Description
CentroidLayer	Centroid prediction layer.

CentroidLayer

Bases: InferenceLayer

Centroid prediction layer.

Parameters:

Name	Type	Description	Default
backend	ModelBackend	Runtime backend for the centroid model. Required even when use_gt_centroids=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_instances	Optional[int]	Cap on returned centroids per frame. Below-cap results are NaN-padded; above-cap are truncated by topk confidence.	None
max_stride	int	Maximum stride the model requires the input to be divisible by. Padding is applied bottom-right after the preprocess input-scale resize.	1
centroid_method	Optional[str]	How a GT centroid is derived from the instance's points when use_gt_centroids=True -- one of "center_of_mass", "bbox_center", "geometric_median", "anchor". None (default) infers it from anchor_ind, i.e. the historical behavior. Must match the training config; the loaders read it off the checkpoint's head config (#586).	None
centroid_fallback	Optional[str]	Reduce method used when the anchor node is not visible.	None
anchor_ind	Optional[int]	Skeleton-node index to use as the centroid anchor when use_gt_centroids=True. None falls back to the NaN-ignoring mean of all visible nodes for each instance.	None
use_gt_centroids	bool	When True, skip the model and read centroids from a batch's "instances" field (the LabelsReader path).	False
preprocess_config / postprocess_config		Standard knobs.	required

Methods:

Name	Description
__init__	Compose the layer with default empty configs when omitted.
postprocess	Decode confmaps → centroids; coord-unscale; topk + NaN-pad.
predict	Run centroid prediction on image.

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

class CentroidLayer(InferenceLayer):
    """Centroid prediction layer.

    Args:
        backend: Runtime backend for the centroid model. Required even when
            ``use_gt_centroids=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_instances: Cap on returned centroids per frame. Below-cap results
            are NaN-padded; above-cap are truncated by ``topk`` confidence.
        max_stride: Maximum stride the model requires the input to be
            divisible by. Padding is applied bottom-right after the
            preprocess input-scale resize.
        centroid_method: How a GT centroid is derived from the instance's points
            when ``use_gt_centroids=True`` -- one of ``"center_of_mass"``,
            ``"bbox_center"``, ``"geometric_median"``, ``"anchor"``. ``None``
            (default) infers it from ``anchor_ind``, i.e. the historical behavior.
            Must match the training config; the loaders read it off the
            checkpoint's head config (#586).
        centroid_fallback: Reduce method used when the anchor node is not visible.
        anchor_ind: Skeleton-node index to use as the centroid anchor when
            ``use_gt_centroids=True``. ``None`` falls back to the NaN-ignoring
            mean of all visible nodes for each instance.
        use_gt_centroids: When ``True``, skip the model and read centroids
            from a batch's ``"instances"`` field (the LabelsReader path).
        preprocess_config / postprocess_config: Standard knobs.
    """

    _HEAD_OUTPUT_KEY: str = "CentroidConfmapsHead"

    def __init__(
        self,
        backend: ModelBackend,
        output_stride: int,
        max_instances: Optional[int] = None,
        max_stride: int = 1,
        anchor_ind: Optional[int] = None,
        centroid_method: Optional[str] = None,
        centroid_fallback: Optional[str] = None,
        use_gt_centroids: 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(
                max_instances=max_instances,
            ),
            output_stride=output_stride,
            max_stride=max_stride,
        )
        self.max_instances = max_instances
        self.anchor_ind = anchor_ind
        self.centroid_method, self.centroid_fallback = degrade_anchor_if_unresolved(
            centroid_method
            or ("anchor" if anchor_ind is not None else "center_of_mass"),
            centroid_fallback,
            anchor_ind,
        )
        self.use_gt_centroids = use_gt_centroids

    # ──────────────────────────────────────────────────────────────────
    # predict(): override to handle the use_gt_centroids branch
    # ──────────────────────────────────────────────────────────────────

    def predict(
        self,
        image: ImageInput,
        instances: Optional[torch.Tensor] = None,
    ) -> Outputs:
        """Run centroid prediction on ``image``.

        Args:
            image: ``np.ndarray`` or ``torch.Tensor`` in any of the shapes
                accepted by :meth:`InferenceLayer._to_4d_float_tensor`.
            instances: ``(B, max_instances, n_nodes, 2)`` GT instance
                keypoints. Required when ``use_gt_centroids=True``;
                ignored otherwise.

        Returns:
            ``Outputs`` populated with ``pred_centroids`` and
            ``pred_centroid_values`` (and optionally ``pred_confmaps`` if
            the postprocess config asks for it).
        """
        if self.use_gt_centroids:
            if instances is None:
                raise ValueError(
                    "use_gt_centroids=True requires `instances` to be passed "
                    "(typically from a LabelsReader's ground-truth field)."
                )
            return self._predict_from_gt(image, instances)
        return super().predict(image)

    def _predict_from_gt(self, image: ImageInput, instances: torch.Tensor) -> Outputs:
        """Compute centroids from GT instances, no model forward.

        Mirrors the legacy ``CentroidCrop(use_gt_centroids=True)`` branch:
        ``generate_centroids`` reduces ``(B, max_inst, n_nodes, 2)`` GT
        keypoints to ``(B, max_inst, 2)`` centroids. NaN-padded instance
        slots stay NaN; corresponding centroid_values are NaN-masked.
        Truncated/padded to ``self.max_instances`` if set.
        """
        x = self._to_4d_float_tensor(image)
        B = x.shape[0]
        H, W = x.shape[-2], x.shape[-1]

        centroids = generate_centroids(
            instances,
            anchor_ind=self.anchor_ind,
            method=self.centroid_method,
            fallback=self.centroid_fallback,
        )
        # ``centroids`` shape: ``(B, max_inst, 2)`` (3D — same rank as
        # ``Outputs.pred_centroids``).
        device = centroids.device
        n_valid = centroids.shape[1]

        # Confidence = 1.0 where centroid is valid, NaN where padded.
        nan_mask = torch.isnan(centroids).any(dim=-1)  # (B, max_inst)
        centroid_vals = torch.where(
            nan_mask,
            torch.full((B, n_valid), float("nan"), device=device),
            torch.ones((B, n_valid), device=device),
        )

        # Honor max_instances cap; pad-with-NaN or truncate to it.
        # Prefer predict-time override (postprocess_config) over build-time value.
        max_inst = (
            getattr(self.postprocess_config, "max_instances", None)
            or self.max_instances
            or n_valid
        )
        if max_inst > n_valid:
            pad_n = max_inst - n_valid
            centroids = torch.cat(
                [
                    centroids,
                    torch.full((B, pad_n, 2), float("nan"), device=device),
                ],
                dim=1,
            )
            centroid_vals = torch.cat(
                [
                    centroid_vals,
                    torch.full((B, pad_n), float("nan"), device=device),
                ],
                dim=1,
            )
        elif max_inst < n_valid:
            centroids = centroids[:, :max_inst]
            centroid_vals = centroid_vals[:, :max_inst]

        info = PreprocInfo(
            original_size=(H, W),
            processed_size=(H, W),
            eff_scale=torch.ones(B, device=device),
            input_scale=1.0,
            output_stride=1,
        )
        return Outputs(
            pred_centroids=centroids,
            pred_centroid_values=centroid_vals,
            preprocess_info=info,
        )

    # ──────────────────────────────────────────────────────────────────
    # postprocess(): find_local_peaks + coord ladder + topk + NaN pad
    # ──────────────────────────────────────────────────────────────────

    def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
        """Decode confmaps → centroids; coord-unscale; topk + NaN-pad.

        Mirrors the legacy ``CentroidCrop.forward()`` shape contract:
        returns ``(B, max_instances, 2)`` centroids and ``(B, max_instances)``
        values, NaN-padded where no detection.
        """
        # Always the torch decode path: this layer is only built with a
        # ``TorchBackend``; the exported path uses ``ExportedCentroidLayer``
        # (so no double coord ladder, #584).
        confmaps = self._extract_confmaps(raw_out)
        peaks, peak_vals, sample_inds, _channel_inds = find_local_peaks(
            confmaps.detach(),
            threshold=self.postprocess_config.peak_threshold,
            refinement=self.postprocess_config.effective_refinement,
            integral_patch_size=self.postprocess_config.integral_patch_size,
        )

        # Coord ladder: confmap → input pixels → original-image pixels.
        peaks = undo_stride(peaks, info.output_stride)
        peaks = undo_input_scale(peaks, info.input_scale)

        # Batch size from confmaps shape (always available on the torch path).
        B = int(confmaps.shape[0])

        max_instances = (
            getattr(self.postprocess_config, "max_instances", None)
            or self.max_instances
            or self._infer_max_instances(sample_inds)
        )
        if max_instances == 0:
            max_instances = 1  # always emit at least one slot for shape stability

        # Allocate the padded outputs on the same device as the peaks so the
        # scatter below doesn't trip the device check on cuda / mps. Falling
        # back to CPU produces correct results on CPU but silently routes
        # cuda / mps results through CPU (or errors on a downstream scatter).
        device = peaks.device
        padded_peaks = torch.full((B, max_instances, 2), float("nan"), device=device)
        padded_vals = torch.full((B, max_instances), float("nan"), device=device)

        for b in range(B):
            mask = sample_inds == b
            sample_peaks = peaks[mask]  # (n_b, 2)
            sample_vals = peak_vals[mask]  # (n_b,)
            if sample_peaks.numel() == 0:
                continue
            if sample_peaks.shape[0] > max_instances:
                sample_vals, idx = torch.topk(sample_vals, max_instances)
                sample_peaks = sample_peaks[idx]
            n = sample_peaks.shape[0]
            padded_peaks[b, :n] = sample_peaks
            padded_vals[b, :n] = sample_vals

        # Reverse the per-sample sizematcher last (matches legacy ordering).
        padded_peaks = undo_eff_scale(padded_peaks, info.eff_scale)

        outputs = Outputs(
            pred_centroids=padded_peaks,
            pred_centroid_values=padded_vals,
            preprocess_info=info,
        )
        if self.postprocess_config.return_confmaps and confmaps is not None:
            outputs = attrs.evolve(outputs, pred_confmaps=confmaps.detach())
        return outputs

    # ──────────────────────────────────────────────────────────────────
    # Helpers
    # ──────────────────────────────────────────────────────────────────

    @staticmethod
    def _infer_max_instances(sample_inds: torch.Tensor) -> int:
        """Find the busiest sample's peak count."""
        if sample_inds.numel() == 0:
            return 0
        counts = torch.bincount(sample_inds.long())
        return int(counts.max().item())

__init__(backend, output_stride, max_instances=None, max_stride=1, anchor_ind=None, centroid_method=None, centroid_fallback=None, use_gt_centroids=False, preprocess_config=None, postprocess_config=None)

Compose the layer with default empty configs when omitted.

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

def __init__(
    self,
    backend: ModelBackend,
    output_stride: int,
    max_instances: Optional[int] = None,
    max_stride: int = 1,
    anchor_ind: Optional[int] = None,
    centroid_method: Optional[str] = None,
    centroid_fallback: Optional[str] = None,
    use_gt_centroids: 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(
            max_instances=max_instances,
        ),
        output_stride=output_stride,
        max_stride=max_stride,
    )
    self.max_instances = max_instances
    self.anchor_ind = anchor_ind
    self.centroid_method, self.centroid_fallback = degrade_anchor_if_unresolved(
        centroid_method
        or ("anchor" if anchor_ind is not None else "center_of_mass"),
        centroid_fallback,
        anchor_ind,
    )
    self.use_gt_centroids = use_gt_centroids

postprocess(raw_out, info)

Decode confmaps → centroids; coord-unscale; topk + NaN-pad.

Mirrors the legacy CentroidCrop.forward() shape contract: returns (B, max_instances, 2) centroids and (B, max_instances) values, NaN-padded where no detection.

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

def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
    """Decode confmaps → centroids; coord-unscale; topk + NaN-pad.

    Mirrors the legacy ``CentroidCrop.forward()`` shape contract:
    returns ``(B, max_instances, 2)`` centroids and ``(B, max_instances)``
    values, NaN-padded where no detection.
    """
    # Always the torch decode path: this layer is only built with a
    # ``TorchBackend``; the exported path uses ``ExportedCentroidLayer``
    # (so no double coord ladder, #584).
    confmaps = self._extract_confmaps(raw_out)
    peaks, peak_vals, sample_inds, _channel_inds = find_local_peaks(
        confmaps.detach(),
        threshold=self.postprocess_config.peak_threshold,
        refinement=self.postprocess_config.effective_refinement,
        integral_patch_size=self.postprocess_config.integral_patch_size,
    )

    # Coord ladder: confmap → input pixels → original-image pixels.
    peaks = undo_stride(peaks, info.output_stride)
    peaks = undo_input_scale(peaks, info.input_scale)

    # Batch size from confmaps shape (always available on the torch path).
    B = int(confmaps.shape[0])

    max_instances = (
        getattr(self.postprocess_config, "max_instances", None)
        or self.max_instances
        or self._infer_max_instances(sample_inds)
    )
    if max_instances == 0:
        max_instances = 1  # always emit at least one slot for shape stability

    # Allocate the padded outputs on the same device as the peaks so the
    # scatter below doesn't trip the device check on cuda / mps. Falling
    # back to CPU produces correct results on CPU but silently routes
    # cuda / mps results through CPU (or errors on a downstream scatter).
    device = peaks.device
    padded_peaks = torch.full((B, max_instances, 2), float("nan"), device=device)
    padded_vals = torch.full((B, max_instances), float("nan"), device=device)

    for b in range(B):
        mask = sample_inds == b
        sample_peaks = peaks[mask]  # (n_b, 2)
        sample_vals = peak_vals[mask]  # (n_b,)
        if sample_peaks.numel() == 0:
            continue
        if sample_peaks.shape[0] > max_instances:
            sample_vals, idx = torch.topk(sample_vals, max_instances)
            sample_peaks = sample_peaks[idx]
        n = sample_peaks.shape[0]
        padded_peaks[b, :n] = sample_peaks
        padded_vals[b, :n] = sample_vals

    # Reverse the per-sample sizematcher last (matches legacy ordering).
    padded_peaks = undo_eff_scale(padded_peaks, info.eff_scale)

    outputs = Outputs(
        pred_centroids=padded_peaks,
        pred_centroid_values=padded_vals,
        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(image, instances=None)

Run centroid prediction on image.

Parameters:

Name	Type	Description	Default
image	ImageInput	np.ndarray or torch.Tensor in any of the shapes accepted by :meth:InferenceLayer._to_4d_float_tensor.	required
instances	Optional[Tensor]	(B, max_instances, n_nodes, 2) GT instance keypoints. Required when use_gt_centroids=True; ignored otherwise.	None

Returns:

Type	Description
Outputs	Outputs populated with pred_centroids and pred_centroid_values (and optionally pred_confmaps if the postprocess config asks for it).

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

def predict(
    self,
    image: ImageInput,
    instances: Optional[torch.Tensor] = None,
) -> Outputs:
    """Run centroid prediction on ``image``.

    Args:
        image: ``np.ndarray`` or ``torch.Tensor`` in any of the shapes
            accepted by :meth:`InferenceLayer._to_4d_float_tensor`.
        instances: ``(B, max_instances, n_nodes, 2)`` GT instance
            keypoints. Required when ``use_gt_centroids=True``;
            ignored otherwise.

    Returns:
        ``Outputs`` populated with ``pred_centroids`` and
        ``pred_centroid_values`` (and optionally ``pred_confmaps`` if
        the postprocess config asks for it).
    """
    if self.use_gt_centroids:
        if instances is None:
            raise ValueError(
                "use_gt_centroids=True requires `instances` to be passed "
                "(typically from a LabelsReader's ground-truth field)."
            )
        return self._predict_from_gt(image, instances)
    return super().predict(image)