segmentation_maps

sleap_nn.data.segmentation_maps

Generate ground truth tensors for instance segmentation from SegmentationMask objects.

Functions:

Name	Description
generate_center_heatmap	Generate Gaussian heatmap at each instance mask centroid.
generate_center_offsets	Generate per-pixel offset vectors pointing to each pixel's instance center.
generate_foreground_mask	Generate binary foreground mask as union of all instance masks.

generate_center_heatmap(masks, img_hw, output_stride=2, sigma=4.0, centers=None)

Generate Gaussian heatmap at each instance mask centroid.

Parameters:

Name	Type	Description	Default
masks	List[ndarray]	List of 2D boolean arrays (H, W), one per instance.	required
img_hw	Tuple[int, int]	Original image size as (height, width).	required
output_stride	int	Stride for downsampling the output.	2
sigma	float	Standard deviation of the Gaussian in pixels (at original resolution).	4.0
centers	Optional[List[Tuple[float, float]]]	Pre-computed list of (x, y) centroid coordinates. If None, centroids will be computed from masks via _compute_mask_centroids.	None

Returns:

Type	Description
Tensor	Tensor of shape (1, 1, H/s, W/s) with float32 values.

Source code in sleap_nn/data/segmentation_maps.py

def generate_center_heatmap(
    masks: List[np.ndarray],
    img_hw: Tuple[int, int],
    output_stride: int = 2,
    sigma: float = 4.0,
    centers: Optional[List[Tuple[float, float]]] = None,
) -> torch.Tensor:
    """Generate Gaussian heatmap at each instance mask centroid.

    Args:
        masks: List of 2D boolean arrays (H, W), one per instance.
        img_hw: Original image size as (height, width).
        output_stride: Stride for downsampling the output.
        sigma: Standard deviation of the Gaussian in pixels (at original resolution).
        centers: Pre-computed list of (x, y) centroid coordinates. If None, centroids
            will be computed from masks via ``_compute_mask_centroids``.

    Returns:
        Tensor of shape (1, 1, H/s, W/s) with float32 values.
    """
    height, width = img_hw
    out_h = height // output_stride
    out_w = width // output_stride
    xv = torch.arange(out_w, dtype=torch.float32) * output_stride + output_stride / 2.0
    yv = torch.arange(out_h, dtype=torch.float32) * output_stride + output_stride / 2.0

    if len(masks) == 0:
        return torch.zeros((1, 1, out_h, out_w), dtype=torch.float32)

    # Compute centroids of each mask if not provided
    if centers is None:
        centers = _compute_mask_centroids(masks)  # (N, 2) in (x, y) pixel coords

    # Build heatmap as max of Gaussians
    heatmap = torch.zeros((1, 1, out_h, out_w), dtype=torch.float32)
    xv_grid = xv.reshape(1, 1, 1, -1)  # (1, 1, 1, W/s)
    yv_grid = yv.reshape(1, 1, -1, 1)  # (1, 1, H/s, 1)
    scaled_sigma = sigma * output_stride

    for cx, cy in centers:
        g = torch.exp(
            -((xv_grid - cx) ** 2 + (yv_grid - cy) ** 2) / (2 * scaled_sigma**2)
        )
        heatmap = torch.maximum(heatmap, g)

    return heatmap  # (1, 1, H/s, W/s)

generate_center_offsets(masks, img_hw, output_stride=2, centers=None)

Generate per-pixel offset vectors pointing to each pixel's instance center.

Parameters:

Name	Type	Description	Default
masks	List[ndarray]	List of 2D boolean arrays (H, W), one per instance.	required
img_hw	Tuple[int, int]	Original image size as (height, width).	required
output_stride	int	Stride for downsampling the output.	2
centers	Optional[List[Tuple[float, float]]]	Pre-computed list of (x, y) centroid coordinates. If None, centroids will be computed from masks via _compute_mask_centroids.	None

Returns:

Type	Description
Tuple[Tensor, Tensor]	Tuple of: offsets: Tensor of shape (1, 2, H/s, W/s) with (dx, dy) offset vectors. Only defined on foreground pixels; background pixels are 0. weight_mask: Tensor of shape (1, 1, H/s, W/s) binary mask indicating where offset loss should be computed (foreground pixels).

Source code in sleap_nn/data/segmentation_maps.py

def generate_center_offsets(
    masks: List[np.ndarray],
    img_hw: Tuple[int, int],
    output_stride: int = 2,
    centers: Optional[List[Tuple[float, float]]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Generate per-pixel offset vectors pointing to each pixel's instance center.

    Args:
        masks: List of 2D boolean arrays (H, W), one per instance.
        img_hw: Original image size as (height, width).
        output_stride: Stride for downsampling the output.
        centers: Pre-computed list of (x, y) centroid coordinates. If None, centroids
            will be computed from masks via ``_compute_mask_centroids``.

    Returns:
        Tuple of:
            offsets: Tensor of shape (1, 2, H/s, W/s) with (dx, dy) offset vectors.
                Only defined on foreground pixels; background pixels are 0.
            weight_mask: Tensor of shape (1, 1, H/s, W/s) binary mask indicating
                where offset loss should be computed (foreground pixels).
    """
    height, width = img_hw
    out_h = height // output_stride
    out_w = width // output_stride

    offsets = torch.zeros((1, 2, out_h, out_w), dtype=torch.float32)
    weight_mask = torch.zeros((1, 1, out_h, out_w), dtype=torch.float32)

    if len(masks) == 0:
        return offsets, weight_mask

    if centers is None:
        centers = _compute_mask_centroids(masks)

    # Sort by area descending so smaller instances overwrite larger ones in overlaps
    areas = [m.sum() for m in masks]
    sorted_indices = sorted(range(len(masks)), key=lambda i: areas[i], reverse=True)

    # Create coordinate grids at output stride resolution
    # Grid values are in original pixel coordinates
    yy = torch.arange(out_h, dtype=torch.float32) * output_stride + output_stride / 2.0
    xx = torch.arange(out_w, dtype=torch.float32) * output_stride + output_stride / 2.0
    grid_x, grid_y = torch.meshgrid(xx, yy, indexing="xy")  # both (out_h, out_w)

    for idx in sorted_indices:
        m = masks[idx]
        # Downsample mask
        m_tensor = (
            torch.from_numpy(m[:height, :width].astype(np.float32))
            .unsqueeze(0)
            .unsqueeze(0)
        )
        if output_stride > 1:
            m_ds = F.interpolate(m_tensor, size=(out_h, out_w), mode="area")
        else:
            m_ds = m_tensor
        m_binary = m_ds[0, 0] > 0.5  # (out_h, out_w)

        cx, cy = centers[idx]

        # Offset = center - pixel_coord (so pixel + offset = center)
        dx = cx - grid_x  # (out_h, out_w)
        dy = cy - grid_y

        # Only set offsets for this instance's foreground pixels
        offsets[0, 0][m_binary] = dx[m_binary]
        offsets[0, 1][m_binary] = dy[m_binary]
        weight_mask[0, 0][m_binary] = 1.0

    return offsets, weight_mask

generate_foreground_mask(masks, img_hw, output_stride=2, maxpool=False)

Generate binary foreground mask as union of all instance masks.

Parameters:

Name	Type	Description	Default
masks	List[ndarray]	List of 2D boolean arrays (H, W), one per instance.	required
img_hw	Tuple[int, int]	Original image size as (height, width).	required
output_stride	int	Stride for downsampling the output mask.	2
maxpool	bool	When True, a stride cell is foreground if ANY of its source pixels is foreground (max-pool semantics: area-downsample then keep > 0). When False (default), the cell must have >50% foreground coverage (area-average > 0.5) — byte-for-byte the previous behavior. maxpool preserves thin structures (e.g. plant roots) that would otherwise erode below the 50% threshold when output_stride > 1. Inert at output_stride=1 (no downsample).	False

Returns:

Type	Description
Tensor	Tensor of shape (1, 1, H/s, W/s) with float32 values in [0, 1].

Source code in sleap_nn/data/segmentation_maps.py

def generate_foreground_mask(
    masks: List[np.ndarray],
    img_hw: Tuple[int, int],
    output_stride: int = 2,
    maxpool: bool = False,
) -> torch.Tensor:
    """Generate binary foreground mask as union of all instance masks.

    Args:
        masks: List of 2D boolean arrays (H, W), one per instance.
        img_hw: Original image size as (height, width).
        output_stride: Stride for downsampling the output mask.
        maxpool: When ``True``, a stride cell is foreground if ANY of its source
            pixels is foreground (max-pool semantics: area-downsample then keep
            ``> 0``). When ``False`` (default), the cell must have >50% foreground
            coverage (area-average > 0.5) — byte-for-byte the previous behavior.
            ``maxpool`` preserves thin structures (e.g. plant roots) that would
            otherwise erode below the 50% threshold when ``output_stride`` > 1.
            Inert at ``output_stride=1`` (no downsample).

    Returns:
        Tensor of shape (1, 1, H/s, W/s) with float32 values in [0, 1].
    """
    height, width = img_hw
    out_h = height // output_stride
    out_w = width // output_stride

    if len(masks) == 0:
        return torch.zeros((1, 1, out_h, out_w), dtype=torch.float32)

    # Union of all masks at original resolution
    union = np.zeros((height, width), dtype=bool)
    for m in masks:
        # Handle masks that may be different sizes than image
        mh, mw = m.shape
        h_end = min(mh, height)
        w_end = min(mw, width)
        union[:h_end, :w_end] |= m[:h_end, :w_end]

    # Convert to tensor and downsample via area interpolation
    fg = torch.from_numpy(union.astype(np.float32)).unsqueeze(0).unsqueeze(0)
    if output_stride > 1:
        fg = F.interpolate(fg, size=(out_h, out_w), mode="area")
    # Binarize the (soft, from area interpolation) coverage. ``maxpool`` keeps any
    # nonzero coverage (thin-structure-preserving); the default keeps >50%.
    fg = (fg > (0.0 if maxpool else 0.5)).float()

    return fg  # (1, 1, H/s, W/s)