peak_finding

sleap_nn.inference.peak_finding

Backward-compatibility re-export shim for peak finding.

The implementations live in :mod:sleap_nn.inference.ops.peaks and :mod:sleap_nn.inference.ops.crops after PR 1 of #508. This module preserves the old import path for existing callers; it is scheduled for deletion in #519 alongside the rest of the legacy inference layout.

Functions:

Name	Description
crop_bboxes	Crop bounding boxes from a batch of images.
find_global_peaks	Find global peaks with optional refinement.
find_global_peaks_rough	Find the global maximum for each sample and channel.
find_local_peaks	Find local peaks with optional refinement.
find_local_peaks_rough	Find local maxima via non-maximum suppression.
integral_regression	Compute regression by integrating over the confidence maps on a grid.
morphological_dilation	Compute the per-pixel max over the 8-neighborhood (excluding center).

crop_bboxes(images, bboxes, sample_inds, crop_size=None)

Crop bounding boxes from a batch of images.

Parameters:

Name	Type	Description	Default
images	Tensor	(samples, channels, height, width) tensor.	required
bboxes	Tensor	(n_bboxes, 4, 2) float32 corner tensor in top-left → top-right → bottom-right → bottom-left order. Build these with :func:make_centered_bboxes.	required
sample_inds	Tensor	(n_bboxes,) int tensor selecting which sample each bbox crops from.	required
crop_size	Optional[Tuple[int, int]]	(height, width) of the crops when the caller knows it (e.g. a top-down layer's crop_size); skips deriving it from the first box (and the device sync that costs). None derives it.	None

Returns:

Type	Description
Tensor	(n_bboxes, channels, crop_height, crop_width) of the same dtype as images. The crop size is inferred from the first bbox.

Notes

Bbox top-lefts are floored to the integer grid before extraction, matching the prior .to(torch.long) behavior. Out-of-image sample positions are zero-padded.

See Also

:func:make_centered_bboxes.

Source code in sleap_nn/inference/ops/crops.py

def crop_bboxes(
    images: torch.Tensor,
    bboxes: torch.Tensor,
    sample_inds: torch.Tensor,
    crop_size: Optional[Tuple[int, int]] = None,
) -> torch.Tensor:
    """Crop bounding boxes from a batch of images.

    Args:
        images: ``(samples, channels, height, width)`` tensor.
        bboxes: ``(n_bboxes, 4, 2)`` ``float32`` corner tensor in
            top-left → top-right → bottom-right → bottom-left order. Build
            these with :func:`make_centered_bboxes`.
        sample_inds: ``(n_bboxes,)`` int tensor selecting which sample each
            bbox crops from.
        crop_size: ``(height, width)`` of the crops when the caller knows it
            (e.g. a top-down layer's ``crop_size``); skips deriving it from the
            first box (and the device sync that costs). ``None`` derives it.

    Returns:
        ``(n_bboxes, channels, crop_height, crop_width)`` of the same dtype
        as ``images``. The crop size is inferred from the first bbox.

    Notes:
        Bbox top-lefts are floored to the integer grid before extraction,
        matching the prior ``.to(torch.long)`` behavior. Out-of-image
        sample positions are zero-padded.

    See Also:
        :func:`make_centered_bboxes`.
    """
    n_crops = bboxes.shape[0]
    if n_crops == 0:
        return torch.empty(
            0, images.shape[1], 0, 0, device=images.device, dtype=images.dtype
        )

    if crop_size is not None:
        height, width = int(crop_size[0]), int(crop_size[1])
    else:
        # Crop size from the first bbox. ``.item()`` makes this a Python int —
        # required to allocate fixed-shape index tensors. PR 7 (#515)
        # parameterizes ONNX wrappers with the constexpr crop size to bypass
        # this call. Round, don't truncate: `make_centered_bboxes` corners are
        # float32, and for ~5% of centroid positions their span comes out just
        # under `size - 1` (190.99998 for a 192 px crop), which truncation turned
        # into a crop 1 px short -- for every crop of the batch (#778).
        height = int(torch.round(abs(bboxes[0, 3, 1] - bboxes[0, 0, 1])).item()) + 1
        width = int(torch.round(abs(bboxes[0, 1, 0] - bboxes[0, 0, 0])).item()) + 1

    device = images.device
    bboxes_on_device = bboxes.to(device)

    # Reproduce the LEGACY crop top-left exactly. Legacy padded each side by
    # ``dim // 2`` and indexed the unfold patch grid at ``trunc(top_left + half)``;
    # in original-image coords that is ``trunc(top_left + half) - half``. Plain
    # ``.long()`` (truncate-toward-zero) on the raw top-left diverges by one pixel
    # when the bbox top-left is negative AND fractional (an instance overhanging
    # the top/left image edge), shifting the crop fed to the centered-instance
    # model. ``trunc(x + half) - half`` is identical to ``.long()`` for the
    # integer-aligned case and reproduces legacy for all REACHABLE centroids
    # (model peaks/centroids are >= 0). NOTE: for off-frame centroids <= -1 px,
    # legacy additionally clamps the padded patch index to 0 (a further 1 px
    # shift) that this does not replicate; that far-out-of-bounds clamp
    # divergence is unreachable via the model path and tracked in the #530
    # follow-ups (#584).
    half_xy = torch.tensor(
        [width // 2, height // 2], device=device, dtype=bboxes_on_device.dtype
    )
    crop_topleft = (bboxes_on_device[:, 0, :] + half_xy).to(
        torch.long
    ) - half_xy.long()  # (n, 2) -- (x, y)

    # Pad the source image with zeros so out-of-bounds samples become 0.
    # The pad amount is at least ``max(width, height)`` so any in-image
    # sub-pixel can shift up to one full crop without escaping the padded
    # region (used as a static, ONNX-friendly upper bound).
    pad_h, pad_w = height, width
    images_padded = F.pad(
        images, (pad_w, pad_w, pad_h, pad_h), mode="constant", value=0
    )
    padded_h, padded_w = images_padded.shape[-2], images_padded.shape[-1]

    # Build per-crop sample indices over the padded image.
    yy, xx = torch.meshgrid(
        torch.arange(height, dtype=torch.long, device=device),
        torch.arange(width, dtype=torch.long, device=device),
        indexing="ij",
    )
    # offsets shape: (height, width)
    # crop top-left in padded coords = original top-left + (pad_w, pad_h)
    abs_x = crop_topleft[:, 0:1, None] + xx + pad_w  # (n, h, w)
    abs_y = crop_topleft[:, 1:2, None] + yy + pad_h  # (n, h, w)

    # Clamp to padded bounds (defends against extremely-out-of-bounds bboxes).
    abs_x = abs_x.clamp(0, padded_w - 1)
    abs_y = abs_y.clamp(0, padded_h - 1)

    # Gather. Result shape: (n_crops, channels, height, width).
    if not isinstance(sample_inds, torch.Tensor):
        sample_inds = torch.tensor(sample_inds, device=device)
    sample_inds_long = sample_inds.to(device=device, dtype=torch.long)
    sample_idx = sample_inds_long[:, None, None]  # (n, 1, 1)
    crops = images_padded[sample_idx, :, abs_y, abs_x]  # (n, h, w, c)
    # Move channels back to dim 1: (n, c, h, w)
    return crops.permute(0, 3, 1, 2).contiguous()

find_global_peaks(cms, threshold=0.2, refinement=None, integral_patch_size=5)

Find global peaks with optional refinement.

Parameters:

Name	Type	Description	Default
cms	Tensor	(samples, channels, height, width) confidence maps.	required
threshold	float	Peaks below this are NaN-padded.	0.2
refinement	Optional[str]	None returns grid-aligned peaks; "integral" runs sub-pixel integral refinement on a small patch around each peak.	None
integral_patch_size	int	Side length of the refinement patch.	5

Returns:

Type	Description
Tuple[Tensor, Tensor]	(peak_points, peak_vals) as in :func:find_global_peaks_rough.

Source code in sleap_nn/inference/ops/peaks.py

def find_global_peaks(
    cms: torch.Tensor,
    threshold: float = 0.2,
    refinement: Optional[str] = None,
    integral_patch_size: int = 5,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Find global peaks with optional refinement.

    Args:
        cms: ``(samples, channels, height, width)`` confidence maps.
        threshold: Peaks below this are NaN-padded.
        refinement: ``None`` returns grid-aligned peaks; ``"integral"`` runs
            sub-pixel integral refinement on a small patch around each peak.
        integral_patch_size: Side length of the refinement patch.

    Returns:
        ``(peak_points, peak_vals)`` as in :func:`find_global_peaks_rough`.
    """
    rough_peaks, peak_vals = find_global_peaks_rough(cms, threshold=threshold)

    if refinement is None or torch.isnan(rough_peaks).all():
        return rough_peaks, peak_vals
    if refinement != "integral":
        return rough_peaks, peak_vals

    crop_size = integral_patch_size

    samples = cms.size(0)
    channels = cms.size(1)
    rough_peaks = rough_peaks.view(samples * channels, 2)

    valid_idx = torch.where(~torch.isnan(rough_peaks[:, 0]))[0]
    valid_peaks = rough_peaks[valid_idx]

    bboxes = make_centered_bboxes(
        valid_peaks, box_height=crop_size, box_width=crop_size
    )

    cms = torch.reshape(cms, [samples * channels, 1, cms.size(2), cms.size(3)])
    cm_crops = crop_bboxes(cms, bboxes, valid_idx)

    gv = torch.arange(crop_size, dtype=torch.float32) - ((crop_size - 1) / 2)
    dx_hat, dy_hat = integral_regression(cm_crops, xv=gv, yv=gv)
    offsets = torch.cat([dx_hat, dy_hat], dim=1)

    refined_peaks = rough_peaks.clone()
    refined_peaks[valid_idx] += offsets

    return refined_peaks.reshape(samples, channels, 2), peak_vals

find_global_peaks_rough(cms, threshold=0.1)

Find the global maximum for each sample and channel.

Parameters:

Name	Type	Description	Default
cms	Tensor	(samples, channels, height, width).	required
threshold	float	Peaks below this are replaced with NaN.	0.1

Returns:

Type	Description
Tuple[Tensor, Tensor]	(peak_points, peak_vals) where peak_points is (samples, channels, 2) in (x, y) order and peak_vals is (samples, channels).

Source code in sleap_nn/inference/ops/peaks.py

def find_global_peaks_rough(
    cms: torch.Tensor, threshold: float = 0.1
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Find the global maximum for each sample and channel.

    Args:
        cms: ``(samples, channels, height, width)``.
        threshold: Peaks below this are replaced with NaN.

    Returns:
        ``(peak_points, peak_vals)`` where ``peak_points`` is
        ``(samples, channels, 2)`` in ``(x, y)`` order and ``peak_vals`` is
        ``(samples, channels)``.
    """
    max_values, _max_indices_y = torch.max(cms, dim=2, keepdim=True)
    max_values, max_indices_x = torch.max(max_values, dim=3, keepdim=True)
    # Drop dims one at a time so the ONNX exporter can lower each Squeeze
    # node independently (a single ``dim=(2, 3)`` argument is not supported).
    max_indices_x = max_indices_x.squeeze(3).squeeze(2)

    amax_values, _amax_indices_x = torch.max(cms, dim=3, keepdim=True)
    amax_values, amax_indices_y = torch.max(amax_values, dim=2, keepdim=True)
    amax_indices_y = amax_indices_y.squeeze(3).squeeze(2)

    peak_points = torch.cat(
        [max_indices_x.unsqueeze(-1), amax_indices_y.unsqueeze(-1)], dim=-1
    ).to(torch.float32)
    max_values = max_values.squeeze(-1).squeeze(-1)

    # Below-threshold positions get NaN coords + zero value. We use
    # ``torch.where`` rather than boolean-mask in-place assignment so this
    # function exports to ONNX cleanly (PR 5 of #508).
    below_threshold_mask = max_values < threshold
    peak_points = torch.where(
        below_threshold_mask.unsqueeze(-1).expand_as(peak_points),
        torch.full_like(peak_points, float("nan")),
        peak_points,
    )
    max_values = torch.where(
        below_threshold_mask, torch.zeros_like(max_values), max_values
    )
    return peak_points, max_values

find_local_peaks(cms, threshold=0.2, refinement=None, integral_patch_size=5)

Find local peaks with optional refinement.

Same return shape as :func:find_local_peaks_rough. refinement accepts None (no refinement) or "integral".

Source code in sleap_nn/inference/ops/peaks.py

def find_local_peaks(
    cms: torch.Tensor,
    threshold: float = 0.2,
    refinement: Optional[str] = None,
    integral_patch_size: int = 5,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Find local peaks with optional refinement.

    Same return shape as :func:`find_local_peaks_rough`. ``refinement``
    accepts ``None`` (no refinement) or ``"integral"``.
    """
    rough_peaks, peak_vals, peak_sample_inds, peak_channel_inds = (
        find_local_peaks_rough(cms, threshold=threshold)
    )

    if rough_peaks.size(0) == 0 or refinement is None:
        return rough_peaks, peak_vals, peak_sample_inds, peak_channel_inds
    if refinement != "integral":
        return rough_peaks, peak_vals, peak_sample_inds, peak_channel_inds

    crop_size = integral_patch_size

    # A flat top's rough peak can be fractional; refine around the pixel it
    # floors to, since the crops are taken on the pixel grid.
    anchors = torch.floor(rough_peaks)
    bboxes = make_centered_bboxes(anchors, box_height=crop_size, box_width=crop_size)

    samples = cms.size(0)
    channels = cms.size(1)
    cms = torch.reshape(cms, [samples * channels, 1, cms.size(2), cms.size(3)])
    box_sample_inds = (peak_sample_inds * channels) + peak_channel_inds

    cm_crops = crop_bboxes(cms, bboxes, sample_inds=box_sample_inds)

    gv = torch.arange(crop_size, dtype=torch.float32) - ((crop_size - 1) / 2)
    dx_hat, dy_hat = integral_regression(cm_crops, xv=gv, yv=gv)
    offsets = torch.cat([dx_hat, dy_hat], dim=1)

    refined_peaks = anchors + offsets
    return refined_peaks, peak_vals, peak_sample_inds, peak_channel_inds

find_local_peaks_rough(cms, threshold=0.2)

Find local maxima via non-maximum suppression.

A pixel is a peak if it is above threshold and greater than all 8 of its neighbors. A flat top of equal pixels (common with fp16 confidence maps) is one peak if it fits in a 2x2 block (2 pixels, an L of 3, or all 4) and is higher than every pixel around it; its point is the mean of the pixels' coordinates, so it can be fractional. A wider flat top, or one next to a higher pixel, gives no peak. See #781.

Parameters:

Name	Type	Description	Default
cms	Tensor	(samples, channels, height, width).	required
threshold	float	Peaks below this are dropped.	0.2

Returns:

Type	Description
Tuple[Tensor, Tensor, Tensor, Tensor]	(peak_points, peak_vals, peak_sample_inds, peak_channel_inds): peak_points is (n_peaks, 2) in (x, y) order; peak_vals is (n_peaks,); the index tensors are (n_peaks,) int32.

Source code in sleap_nn/inference/ops/peaks.py

def find_local_peaks_rough(
    cms: torch.Tensor, threshold: float = 0.2
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Find local maxima via non-maximum suppression.

    A pixel is a peak if it is above ``threshold`` and greater than all 8 of
    its neighbors. A flat top of equal pixels (common with fp16 confidence
    maps) is one peak if it fits in a 2x2 block (2 pixels, an L of 3, or all
    4) and is higher than every pixel around it; its point is the mean of the
    pixels' coordinates, so it can be fractional. A wider flat top, or one
    next to a higher pixel, gives no peak. See #781.

    Args:
        cms: ``(samples, channels, height, width)``.
        threshold: Peaks below this are dropped.

    Returns:
        ``(peak_points, peak_vals, peak_sample_inds, peak_channel_inds)``:
        ``peak_points`` is ``(n_peaks, 2)`` in ``(x, y)`` order;
        ``peak_vals`` is ``(n_peaks,)``;
        the index tensors are ``(n_peaks,)`` ``int32``.
    """
    padded = F.pad(cms, (2, 2, 2, 2), mode="constant", value=float("-inf"))
    before_img, after_img = _neighbor_max_split(padded[..., 1:-1, 1:-1])
    # >= all 8 neighbors and > the 4 before it: the first pixel of a flat top.
    argmax_and_thresh_img = (cms > before_img) & (cms >= after_img) & (cms > threshold)

    peak_subs = torch.stack(
        torch.where(argmax_and_thresh_img.permute(0, 2, 3, 1)), axis=-1
    )
    s, y, x, c = peak_subs.unbind(1)
    peak_vals = cms[s, c, y, x]
    peak_points = peak_subs[:, [2, 1]].to(torch.float32)
    tied = peak_vals == after_img[s, c, y, x]
    # Ties are rare in fp32, so the flat-top check only runs when there is one.
    # Traced or compiled graphs always run it (it keeps every untied peak).
    tracing = torch.jit.is_tracing() or torch.compiler.is_compiling()
    if tracing or bool(tied.any()):
        keep, offsets = _flat_tops(padded, peak_subs, peak_vals)
        keep = torch.nonzero(keep)[:, 0]
        peak_subs, peak_vals = peak_subs[keep], peak_vals[keep]
        peak_points = peak_points[keep] + offsets[keep]
    peak_sample_inds = peak_subs[:, 0].to(torch.int32)
    peak_channel_inds = peak_subs[:, 3].to(torch.int32)
    return peak_points, peak_vals, peak_sample_inds, peak_channel_inds

integral_regression(cms, xv, yv)

Compute regression by integrating over the confidence maps on a grid.

Parameters:

Name	Type	Description	Default
cms	Tensor	Confidence maps with shape (samples, channels, height, width).	required
xv	Tensor	float32 x-grid vector of coordinates to sample.	required
yv	Tensor	float32 y-grid vector of coordinates to sample.	required

Returns:

Type	Description
Tuple[Tensor, Tensor]	(x_hat, y_hat) regressed coordinates per channel, each of shape (samples, channels).

Source code in sleap_nn/inference/ops/peaks.py

def integral_regression(
    cms: torch.Tensor, xv: torch.Tensor, yv: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute regression by integrating over the confidence maps on a grid.

    Args:
        cms: Confidence maps with shape ``(samples, channels, height, width)``.
        xv: ``float32`` x-grid vector of coordinates to sample.
        yv: ``float32`` y-grid vector of coordinates to sample.

    Returns:
        ``(x_hat, y_hat)`` regressed coordinates per channel, each of shape
        ``(samples, channels)``.
    """
    z = torch.sum(cms, dim=[2, 3]).to(cms.device)
    xv = xv.to(cms.device)
    yv = yv.to(cms.device)

    x_hat = torch.sum(xv.view(1, 1, 1, -1) * cms, dim=[2, 3]) / z
    y_hat = torch.sum(yv.view(1, 1, -1, 1) * cms, dim=[2, 3]) / z
    return x_hat, y_hat

morphological_dilation(image, kernel)

Compute the per-pixel max over the 8-neighborhood (excluding center).

The kernel argument is preserved for API compatibility but is currently ignored — the 8-neighbor pattern is hardcoded as shifted slices of a padded image combined with elementwise torch.maximum, which exports to ONNX (PR 5 of #508 rewrote the original Tensor.unfold formulation that the legacy ONNX exporter rejected).

Parameters:

Name	Type	Description	Default
image	Tensor	Input tensor of shape (B, 1, H, W).	required
kernel	Tensor	Legacy 3×3 NMS kernel; unused. Kept so existing callers continue to work without modification.	required

Returns:

Type	Description
Tensor	Same shape as image; each output pixel is the max over its eight neighbors in the input (out-of-image neighbors are -inf, i.e. pad-with-minimum).

Source code in sleap_nn/inference/ops/peaks.py

def morphological_dilation(image: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
    """Compute the per-pixel max over the 8-neighborhood (excluding center).

    The ``kernel`` argument is preserved for API compatibility but is
    currently ignored — the 8-neighbor pattern is hardcoded as shifted slices
    of a padded image combined with elementwise ``torch.maximum``, which
    exports to ONNX (PR 5 of #508 rewrote the original ``Tensor.unfold``
    formulation that the legacy ONNX exporter rejected).

    Args:
        image: Input tensor of shape ``(B, 1, H, W)``.
        kernel: Legacy 3×3 NMS kernel; unused. Kept so existing callers
            continue to work without modification.

    Returns:
        Same shape as ``image``; each output pixel is the max over its
        eight neighbors in the input (out-of-image neighbors are ``-inf``,
        i.e. pad-with-minimum).
    """
    del kernel  # see docstring
    return torch.maximum(*morphological_dilation_split(image))