sleap.nn.peak_finding#

This module contains TensorFlow-based peak finding methods.

In general, the inputs to the functions provided here operate on confidence maps (sometimes referred to as heatmaps), which are image-based representations of the locations of landmark coordinates.

In these representations, landmark locations are encoded as probability that it is present each pixel. This is often represented by an unnormalized 2D Gaussian PDF centered at the true location and evaluated over the entire image grid.

Peak finding entails finding either the global or local maxima of these confidence maps.

sleap.nn.peak_finding.crop_bboxes(images: Tensor, bboxes: Tensor, sample_inds: Tensor) → Tensor[source]#

Crop bounding boxes from a batch of images.

This method serves as a convenience method for specifying the arguments of tf.image.crop_and_resize.

Parameters:

images – Tensor of shape (samples, height, width, channels) of a batch of images.
bboxes – Tensor of shape (n_bboxes, 4) and dtype tf.float32, where the last axis corresponds to unnormalized (y1, x1, y2, x2) coordinates of the bounding boxes. This can be generated from centroids using make_centered_bboxes.
sample_inds – Tensor of shape (n_bboxes,) specifying which samples each bounding box should be cropped from.

Returns:

A tensor of shape (n_bboxes, crop_height, crop_width, channels) of the same dtype as the input image. The crop size is inferred from the bounding box coordinates.

Notes

This function expects bounding boxes with coordinates at the centers of the pixels in the box limits. Technically, the box will span (x1 - 0.5, x2 + 0.5) and (y1 - 0.5, y2 + 0.5).

For example, a 3x3 patch centered at (1, 1) would be specified by (y1, x1, y2, x2) = (0, 0, 2, 2). This would be exactly equivalent to indexing the image with image[0:3, 0:3].

See also: make_centered_bboxes

sleap.nn.peak_finding.find_global_peaks(cms: Tensor, threshold: float = 0.2, refinement: str | None = None, integral_patch_size: int = 5) → Tuple[Tensor, Tensor]#

Find global peaks with optional refinement.

Parameters:

cms – Confidence maps. Tensor of shape (samples, height, width, channels).
threshold – Minimum confidence threshold. Peaks with values below this will ignored.
refinement – If None, returns the grid-aligned peaks with no refinement. If "integral", peaks will be refined with integral regression. If "local", peaks will be refined with quarter pixel local gradient offset.
integral_patch_size – Size of patches to crop around each rough peak as an integer scalar.

Returns:

A tuple of (peak_points, peak_vals).

peak_points: float32 tensor of shape (samples, channels, 2), where the last axis indicates peak locations in xy order.

peak_vals: float32 tensor of shape (samples, channels) containing the values at the peak points.

sleap.nn.peak_finding.find_global_peaks_integral(cms: Tensor, crop_size: int = 5, threshold: float = 0.2) → Tuple[Tensor, Tensor][source]#

Find local peaks with integral refinement.

Integral regression refinement will be computed by taking the weighted average of the local neighborhood around each rough peak.

Parameters:

cms – Confidence maps. Tensor of shape (samples, height, width, channels).
crop_size – Size of patches to crop around each rough peak as an integer scalar.
threshold – Minimum confidence threshold. Peaks with values below this will be set to NaNs.

Returns:

A tuple of (peak_points, peak_vals).

peak_points: float32 tensor of shape (samples, channels, 2), where the last axis indicates peak locations in xy order.

peak_vals: float32 tensor of shape (samples, channels) containing the values at the peak points.

sleap.nn.peak_finding.find_global_peaks_rough(cms: Tensor, threshold: float = 0.1) → Tuple[Tensor, Tensor][source]#

Find the global maximum for each sample and channel.

Parameters:

cms – Tensor of shape (samples, height, width, channels).
threshold – Scalar float specifying the minimum confidence value for peaks. Peaks with values below this threshold will be replaced with NaNs.

Returns:

A tuple of (peak_points, peak_vals).

peak_points: float32 tensor of shape (samples, channels, 2), where the last axis indicates peak locations in xy order.

peak_vals: float32 tensor of shape (samples, channels) containing the values at the peak points.

sleap.nn.peak_finding.find_global_peaks_with_offsets(cms: Tensor, offsets: Tensor, threshold: float = 0.2) → Tuple[Tensor, Tensor]#

Find global peaks and refine with learned offset maps.

Parameters:

cms – Confidence maps tensor of shape (samples, height, width, channels).
offsets – Offset maps tensor of shape (samples, height, width, 2 * channels).
threshold – Scalar float specifying the minimum confidence value for peaks. Peaks with values below this threshold will be replaced with NaNs.

Returns:

A tuple of (peak_points, peak_vals).

peak_points: float32 tensor of shape (samples, channels, 2), where the last axis indicates peak locations in xy order.

peak_vals: float32 tensor of shape (samples, channels) containing the values at the peak points.

sleap.nn.peak_finding.find_local_peaks(cms: Tensor, threshold: float = 0.2, refinement: str | None = None, integral_patch_size: int = 5) → Tuple[Tensor, Tensor]#

Find local peaks with optional refinement.

Parameters:

cms – Confidence maps. Tensor of shape (samples, height, width, channels).
threshold – Minimum confidence threshold. Peaks with values below this will ignored.
refinement – If None, returns the grid-aligned peaks with no refinement. If "integral", peaks will be refined with integral regression. If "local", peaks will be refined with quarter pixel local gradient offset.
integral_patch_size – Size of patches to crop around each rough peak as an integer scalar.

Returns:

A tuple of (peak_points, peak_vals, peak_sample_inds, peak_channel_inds).

peak_points: float32 tensor of shape (n_peaks, 2), where the last axis indicates peak locations in xy order.

peak_vals: float32 tensor of shape (n_peaks,) containing the values at the peak points.

peak_sample_inds: int32 tensor of shape (n_peaks,) containing the indices of the sample each peak belongs to.

peak_channel_inds: int32 tensor of shape (n_peaks,) containing the indices of the channel each peak belongs to.

sleap.nn.peak_finding.find_local_peaks_integral(cms: Tensor, crop_size: int = 5, threshold: float = 0.2) → Tuple[Tensor, Tensor, Tensor, Tensor][source]#

Find local peaks with integral refinement.

Parameters:

cms – Confidence maps. Tensor of shape (samples, height, width, channels).
crop_size
threshold – Minimum confidence threshold. Peaks with values below this will ignored.

Returns:

A tuple of (peak_points, peak_vals, peak_sample_inds, peak_channel_inds).

peak_points: float32 tensor of shape (n_peaks, 2), where the last axis indicates peak locations in xy order.

peak_vals: float32 tensor of shape (n_peaks,) containing the values at the peak points.

peak_sample_inds: int32 tensor of shape (n_peaks,) containing the indices of the sample each peak belongs to.

peak_channel_inds: int32 tensor of shape (n_peaks,) containing the indices of the channel each peak belongs to.

sleap.nn.peak_finding.find_local_peaks_rough(cms: Tensor, threshold: float = 0.2) → Tuple[Tensor, Tensor, Tensor, Tensor][source]#

Find local maxima via non-maximum suppresion.

Parameters:

cms – Tensor of shape (samples, height, width, channels).
threshold – Scalar float specifying the minimum confidence value for peaks. Peaks with values below this threshold will not be returned.

Returns:

A tuple of (peak_points, peak_vals, peak_sample_inds, peak_channel_inds). peak_points: float32 tensor of shape (n_peaks, 2), where the last axis indicates peak locations in xy order.

peak_vals: float32 tensor of shape (n_peaks,) containing the values at the peak points.

peak_sample_inds: int32 tensor of shape (n_peaks,) containing the indices of the sample each peak belongs to.

peak_channel_inds: int32 tensor of shape (n_peaks,) containing the indices of the channel each peak belongs to.

sleap.nn.peak_finding.find_local_peaks_with_offsets(cms: Tensor, offsets: Tensor, threshold: float = 0.2) → Tuple[Tensor, Tensor, Tensor, Tensor]#

Find local peaks and refine with learned offset maps.

Parameters:

cms – Confidence maps. Tensor of shape (samples, height, width, channels).
offsets – Offset maps tensor of shape (samples, height, width, 2 * channels).
threshold – Minimum confidence threshold. Peaks with values below this will ignored.

Returns:

A tuple of (peak_points, peak_vals, peak_sample_inds, peak_channel_inds).

peak_points: float32 tensor of shape (n_peaks, 2), where the last axis indicates peak locations in xy order.

peak_vals: float32 tensor of shape (n_peaks,) containing the values at the peak points.

peak_sample_inds: int32 tensor of shape (n_peaks,) containing the indices of the sample each peak belongs to.

peak_channel_inds: int32 tensor of shape (n_peaks,) containing the indices of the channel each peak belongs to.

sleap.nn.peak_finding.find_offsets_local_direction(centered_patches: Tensor, delta: float = 0.25) → Tensor[source]#

Computes subpixel offsets from the direction of the pixels around the peak.

This function finds the delta-offset from the center pixel of peak-centered patches by finding the direction of the gradient around each center.

Parameters:

centered_patches – A rank-4 tensor of shape (samples, 3, 3, 1) corresponding to the centered crops around the grid-anchored peaks. For multi-channel images, stack the channels along the samples axis before calling this function.
delta – Scalar float that will scaled by the gradient direction.

Returns:

offsets, a float32 tensor of shape (samples, 2) where the columns correspond to the offsets relative to the center pixel for the x and y directions respectively, i.e., for the i-th sample:

dx_i, dy_i = offsets[i]

Notes

For symmetric patches, the offset will be 0.

This function can be used to refine peak coordinates by:

Cropping 3 x 3 patches around each peak.
Stacking patches along the samples axis.
Computing the local gradient around each centered patch.
Applying subpixel offsets to each peak.

This is a commonly used algorithm for subpixel peak refinement, described for pose estimation applications in [1].

Example:

>>> find_offsets_local_direction(np.array(
...     [[0., 1., 0.],
...      [1., 3., 2.],
...      [0., 1., 0.]]).reshape(1, 3, 3, 1), 0.25)
<tf.Tensor: shape=(1, 2), dtype=float64, numpy=array([[0.25, 0.  ]])>

References

sleap.nn.peak_finding.integral_regression(cms: Tensor, xv: Tensor, yv: Tensor) → Tuple[Tensor, Tensor][source]#

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

Parameters:

cms – Confidence maps with shape (samples, height, width, channels).
xv – X grid vector tf.float32 of grid coordinates to sample.
yv – Y grid vector tf.float32 of grid coordinates to sample.

Returns:

A tuple of (x_hat, y_hat) with the regressed x- and y-coordinates for each channel of the confidence maps.

x_hat and y_hat are of shape (samples, channels)

sleap.nn.peak_finding.make_gaussian_kernel(size: int, sigma: float) → Tensor[source]#

Generates a square unnormalized 2D symmetric Gaussian kernel.

Parameters:

size – Length of kernel. This should be an odd integer.
sigma – Standard deviation of the Gaussian specified as a scalar float.

Returns:

kernel, a float32 tensor of shape (size, size) with values corresponding to the unnormalized probability density of a 2D Gaussian distribution with symmetric covariance along the x and y directions.

Note

The maximum value of this kernel will be 1.0. To normalize it, divide each element by (2 * np.pi * sigma ** 2).

sleap.nn.peak_finding.smooth_imgs(imgs: Tensor, kernel_size: int = 5, sigma: float = 1.0) → Tensor[source]#

Smooths the input image by convolving it with a Gaussian kernel.

Parameters:

imgs – A rank-4 tensor of shape (samples, height, width, channels).
kernel_size – An odd-valued scalar integer specifying the width and height of the Gaussian kernel.
sigma – Standard deviation of the Gaussian specified as a scalar float.

Returns:

A tensor of the same shape as imgs after convolving with a Gaussian sample and channelwise.

sleap.nn.peak_finding

Contents

sleap.nn.peak_finding#