Skip to content

instance_centroids

sleap_nn.data.instance_centroids

Handle calculation of instance centroids.

Functions:

Name Description
find_points_bbox_midpoint

Find the midpoint of the bounding box of a set of points.
find_points_mean

Find the mean position of a set of points, ignoring NaNs.
generate_centroids

Return centroids, falling back to the mean of visible nodes.

find_points_bbox_midpoint(points)

Find the midpoint of the bounding box of a set of points.

Retained as a utility for callers that explicitly want bbox-midpoint behavior. The canonical anchor fallback used by :func:generate_centroids is :func:find_points_mean (mean of visible nodes) — see that function for the project-wide convention.

Parameters:

Name Type Description Default
points Tensor

A torch.Tensor of dtype torch.float32 and of shape (..., n_points, 2), i.e., rank >= 2.

 required

Returns:

Type Description
Tensor

The midpoints between the bounds of each set of points. The output will be of shape (..., 2), reducing the rank of the input by 1. NaNs will be ignored in the calculation.
Notes

The midpoint is calculated as: xy_mid = xy_min + ((xy_max - xy_min) / 2) = ((2 * xy_min) / 2) + ((xy_max - xy_min) / 2) = (2 * xy_min + xy_max - xy_min) / 2 = (xy_min + xy_max) / 2

Source code in sleap_nn/data/instance_centroids.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def find_points_bbox_midpoint(points: torch.Tensor) -> torch.Tensor:
    """Find the midpoint of the bounding box of a set of points.

    Retained as a utility for callers that explicitly want bbox-midpoint behavior.
    The canonical anchor fallback used by :func:`generate_centroids` is
    :func:`find_points_mean` (mean of visible nodes) — see that function for the
    project-wide convention.

    Args:
        points: A torch.Tensor of dtype torch.float32 and of shape (..., n_points, 2),
            i.e., rank >= 2.

    Returns:
        The midpoints between the bounds of each set of points. The output will be of
        shape (..., 2), reducing the rank of the input by 1. NaNs will be ignored in the
        calculation.

    Notes:
        The midpoint is calculated as:
            xy_mid = xy_min + ((xy_max - xy_min) / 2)
                   = ((2 * xy_min) / 2) + ((xy_max - xy_min) / 2)
                   = (2 * xy_min + xy_max - xy_min) / 2
                   = (xy_min + xy_max) / 2
    """
    pts_min = torch.min(
        torch.where(torch.isnan(points), torch.inf, points), dim=-2
    ).values
    pts_max = torch.max(
        torch.where(torch.isnan(points), -torch.inf, points), dim=-2
    ).values

    return (pts_max + pts_min) * 0.5

find_points_mean(points)

Find the mean position of a set of points, ignoring NaNs.

Parameters:

Name Type Description Default
points Tensor

A torch.Tensor of dtype torch.float32 and of shape (..., n_points, 2), i.e., rank >= 2.

 required

Returns:

Type Description
Tensor

The NaN-ignoring mean across the n_points axis. Output shape (..., 2) (rank reduced by 1). When every point in a slot is NaN, the corresponding output row is NaN.
Source code in sleap_nn/data/instance_centroids.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def find_points_mean(points: torch.Tensor) -> torch.Tensor:
    """Find the mean position of a set of points, ignoring NaNs.

    Args:
        points: A torch.Tensor of dtype torch.float32 and of shape (..., n_points, 2),
            i.e., rank >= 2.

    Returns:
        The NaN-ignoring mean across the ``n_points`` axis. Output shape ``(..., 2)``
        (rank reduced by 1). When *every* point in a slot is NaN, the corresponding
        output row is NaN.
    """
    mask = ~torch.isnan(points)
    # Count non-NaN values PER AXIS so each coordinate's mean divides by the
    # number of non-NaN values on that axis. Counting per-point (mask.any)
    # biased the mean toward 0 when only one coordinate of a point was NaN
    # (#584). For the common all-or-nothing-per-point case the two are equal.
    counts = mask.sum(dim=-2).clamp(min=1).to(points.dtype)  # (..., 2)
    safe = torch.where(mask, points, torch.zeros_like(points))
    means = safe.sum(dim=-2) / counts
    all_nan = (~mask.any(dim=-1)).all(dim=-1, keepdim=True)
    return torch.where(all_nan, torch.full_like(means, float("nan")), means)

generate_centroids(points, anchor_ind=None)

Return centroids, falling back to the mean of visible nodes.

Parameters:

Name Type Description Default
points Tensor

A torch.Tensor of dtype torch.float32 and of shape (..., n_nodes, 2), i.e., rank >= 2.

 required
anchor_ind Optional[int]

The index of the node to use as the anchor for the centroid. If not provided, or if the anchor node is NaN (not visible) for a given instance, the centroid falls back to the NaN-ignoring mean of all visible nodes for that instance.

 None

Returns:

Type Description
Tensor

The centroids of the instances. The output will be of shape (..., 2), reducing the rank of the input by 1. NaNs will be ignored in the calculation.
Note

The missing/occluded-anchor fallback is the mean of visible nodes (find_points_mean). Pre-#530 this was the bounding-box midpoint (find_points_bbox_midpoint). The two modes are tracked for a future revisit in https://github.com/talmolab/sleap-nn/issues/586 — keep this consistent with the centroid target generated during training.

Source code in sleap_nn/data/instance_centroids.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def generate_centroids(
    points: torch.Tensor, anchor_ind: Optional[int] = None
) -> torch.Tensor:
    """Return centroids, falling back to the mean of visible nodes.

    Args:
        points: A torch.Tensor of dtype torch.float32 and of shape (..., n_nodes, 2),
            i.e., rank >= 2.
        anchor_ind: The index of the node to use as the anchor for the centroid.
            If not provided, or if the anchor node is NaN (not visible) for a given
            instance, the centroid falls back to the NaN-ignoring mean of all visible
            nodes for that instance.

    Returns:
        The centroids of the instances. The output will be of shape (..., 2),
        reducing the rank of the input by 1. NaNs will be ignored in the calculation.

    Note:
        The missing/occluded-anchor fallback is the mean of visible nodes
        (``find_points_mean``). Pre-#530 this was the bounding-box midpoint
        (``find_points_bbox_midpoint``). The two modes are tracked for a future
        revisit in https://github.com/talmolab/sleap-nn/issues/586 — keep this
        consistent with the centroid target generated during training.
    """
    if anchor_ind is not None:
        centroids = points[..., anchor_ind, :].clone()
    else:
        centroids = torch.full_like(points[..., 0, :], torch.nan).clone()

    missing_anchors = torch.isnan(centroids).any(dim=-1)
    if missing_anchors.any():
        centroids[missing_anchors] = find_points_mean(points[missing_anchors])

    return centroids  # (..., n_instances, 2)