peaks

sleap_nn.inference.ops.peaks

Peak finding ops for confidence maps.

Functions provided:

• :func:morphological_dilation — pure-PyTorch 3×3 dilation (max over the 8 neighbors).
• :func:morphological_dilation_split — the same 8-neighbor max, split into the neighbors before and after each pixel in raster order.
• :func:integral_regression — sub-pixel refinement by integrating expected coordinates over a local patch.
• :func:find_global_peaks_rough / :func:find_global_peaks — single peak per (sample, channel) for centroid / single-instance confmaps.
• :func:find_local_peaks_rough / :func:find_local_peaks — multi-peak finding for bottom-up confmaps.
• :func:quarter_offset_refinement — sub-pixel refinement of global peaks (refinement="local").

PR 1 is a pure relocation — every function below preserves the exact signature and behavior of its old home in sleap_nn/inference/peak_finding.py.

Functions:

Name	Description
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).
morphological_dilation_split	Compute the 8-neighbor max split into earlier and later neighbors.
quarter_offset_refinement	Refine global peaks with the quarter-offset rule (refinement="local").

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 (or "none") returns grid-aligned peaks; "integral" runs sub-pixel integral refinement on a small patch around each peak; "local" moves each peak a quarter pixel towards its higher neighbor along each axis (:func:quarter_offset_refinement).	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.

Raises:

Type	Description
ValueError	If refinement is not one of the names above.

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`` (or ``"none"``) returns grid-aligned peaks;
            ``"integral"`` runs sub-pixel integral refinement on a small patch
            around each peak; ``"local"`` moves each peak a quarter pixel towards
            its higher neighbor along each axis (:func:`quarter_offset_refinement`).
        integral_patch_size: Side length of the refinement patch.

    Returns:
        ``(peak_points, peak_vals)`` as in :func:`find_global_peaks_rough`.

    Raises:
        ValueError: If ``refinement`` is not one of the names above.
    """
    refinement = _check_refinement(refinement)
    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 == "local":
        return quarter_offset_refinement(cms, 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 / "none" (no refinement), "integral" or "local" (quarter-pixel offset towards the higher neighbor, see :func:quarter_offset_refinement); other names raise ValueError.

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`` / ``"none"`` (no refinement), ``"integral"`` or
    ``"local"`` (quarter-pixel offset towards the higher neighbor, see
    :func:`quarter_offset_refinement`); other names raise ``ValueError``.
    """
    refinement = _check_refinement(refinement)
    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 == "local":
        refined_peaks = _quarter_offset_local_peaks(
            cms, rough_peaks, peak_sample_inds, peak_channel_inds
        )
        return refined_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))

morphological_dilation_split(image)

Compute the 8-neighbor max split into earlier and later neighbors.

The 8 neighbors of each pixel are split by raster order (row, then column): the 4 that come before it (NW, N, NE, W) and the 4 that come after it (E, SW, S, SE). The elementwise max of the two outputs equals :func:morphological_dilation.

Parameters:

Name	Type	Description	Default
image	Tensor	Input tensor of shape (..., H, W).	required

Returns:

Type	Description
Tuple[Tensor, Tensor]	(before, after), each the same shape as image: the max over the 4 earlier and over the 4 later neighbors (out-of-image neighbors are -inf).

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

def morphological_dilation_split(
    image: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute the 8-neighbor max split into earlier and later neighbors.

    The 8 neighbors of each pixel are split by raster order (row, then column):
    the 4 that come before it (NW, N, NE, W) and the 4 that come after it
    (E, SW, S, SE). The elementwise max of the two outputs equals
    :func:`morphological_dilation`.

    Args:
        image: Input tensor of shape ``(..., H, W)``.

    Returns:
        ``(before, after)``, each the same shape as ``image``: the max over the
        4 earlier and over the 4 later neighbors (out-of-image neighbors are
        ``-inf``).
    """
    padded = F.pad(image, (1, 1, 1, 1), mode="constant", value=float("-inf"))
    return _neighbor_max_split(padded)

quarter_offset_refinement(cms, rough_peaks)

Refine global peaks with the quarter-offset rule (refinement="local").

Each coordinate moves 0.25 confmap pixels towards its higher neighbor (x += 0.25 * sign(cm[y, x+1] - cm[y, x-1]), same for y), as in SLEAP's "local" refinement and the Hourglass/SimpleBaseline decoders.

Parameters:

Name	Type	Description	Default
cms	Tensor	(samples, channels, height, width) confidence maps.	required
rough_peaks	Tensor	(samples, channels, 2) grid peaks from :func:find_global_peaks_rough (NaN for missing).	required

Returns:

Type	Description
Tensor	(samples, channels, 2) refined (x, y) peaks (NaN preserved).

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

def quarter_offset_refinement(
    cms: torch.Tensor, rough_peaks: torch.Tensor
) -> torch.Tensor:
    """Refine global peaks with the quarter-offset rule (``refinement="local"``).

    Each coordinate moves 0.25 confmap pixels towards its higher neighbor
    (``x += 0.25 * sign(cm[y, x+1] - cm[y, x-1])``, same for ``y``), as in
    SLEAP's ``"local"`` refinement and the Hourglass/SimpleBaseline decoders.

    Args:
        cms: ``(samples, channels, height, width)`` confidence maps.
        rough_peaks: ``(samples, channels, 2)`` grid peaks from
            :func:`find_global_peaks_rough` (NaN for missing).

    Returns:
        ``(samples, channels, 2)`` refined ``(x, y)`` peaks (NaN preserved).
    """
    p = _gather_neighborhood(cms, rough_peaks, radius=1).to(torch.float32)
    dx = torch.sign(p[..., 1, 2] - p[..., 1, 0])
    dy = torch.sign(p[..., 2, 1] - p[..., 0, 1])
    return rough_peaks + 0.25 * torch.stack([dx, dy], dim=-1)