Skip to content

topdown

sleap_nn.export.wrappers.topdown

Top-down ONNX wrapper.

Classes:

Name Description
TopDownONNXWrapper

ONNX-exportable wrapper for top-down (centroid + centered-instance) inference.

TopDownONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for top-down (centroid + centered-instance) inference.

Expects input images as uint8 tensors in [0, 255].

Methods:

Name Description
__init__

Initialize top-down ONNX wrapper.
forward

Run top-down inference and return fixed-size outputs.
Source code in sleap_nn/export/wrappers/topdown.py 
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 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
 63
 64
 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
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
class TopDownONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for top-down (centroid + centered-instance) inference.

    Expects input images as uint8 tensors in [0, 255].
    """

    def __init__(
        self,
        centroid_model: nn.Module,
        instance_model: nn.Module,
        max_instances: int = 20,
        crop_size: Tuple[int, int] = (192, 192),
        centroid_output_stride: int = 2,
        instance_output_stride: int = 4,
        centroid_input_scale: float = 1.0,
        instance_input_scale: float = 1.0,
        n_nodes: int = 1,
        centroid_peak_threshold: float = 0.2,
        instance_peak_threshold: float = 0.2,
    ) -> None:
        """Initialize top-down ONNX wrapper.

        Args:
            centroid_model: Centroid detection model.
            instance_model: Instance pose estimation model.
            max_instances: Maximum number of instances to detect.
            crop_size: Size of instance crops (height, width).
            centroid_output_stride: Centroid model output stride.
            instance_output_stride: Instance model output stride.
            centroid_input_scale: Centroid input scaling factor.
            instance_input_scale: Instance input scaling factor.
            n_nodes: Number of skeleton nodes.
            centroid_peak_threshold: Minimum confidence for centroid peaks.
            instance_peak_threshold: Minimum confidence for instance peaks.
        """
        super().__init__(centroid_model)
        self.centroid_model = centroid_model
        self.instance_model = instance_model
        self.max_instances = max_instances
        self.crop_size = crop_size
        self.centroid_output_stride = centroid_output_stride
        self.instance_output_stride = instance_output_stride
        self.centroid_input_scale = centroid_input_scale
        self.instance_input_scale = instance_input_scale
        self.n_nodes = n_nodes
        self.centroid_peak_threshold = centroid_peak_threshold
        self.instance_peak_threshold = instance_peak_threshold

        crop_h, crop_w = crop_size
        y_crop = torch.linspace(-1, 1, crop_h, dtype=torch.float32)
        x_crop = torch.linspace(-1, 1, crop_w, dtype=torch.float32)
        grid_y, grid_x = torch.meshgrid(y_crop, x_crop, indexing="ij")
        base_grid = torch.stack([grid_x, grid_y], dim=-1)
        self.register_buffer("base_grid", base_grid, persistent=False)

    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run top-down inference and return fixed-size outputs."""
        image = self._normalize_uint8(image)
        batch_size, channels, height, width = image.shape

        scaled_image = image
        if self.centroid_input_scale != 1.0:
            scaled_h = int(height * self.centroid_input_scale)
            scaled_w = int(width * self.centroid_input_scale)
            scaled_image = F.interpolate(
                scaled_image,
                size=(scaled_h, scaled_w),
                mode="bilinear",
                align_corners=False,
            )

        centroid_out = self.centroid_model(scaled_image)
        centroid_cms = self._extract_tensor(centroid_out, ["centroid", "confmap"])

        centroids, centroid_vals, instance_valid = self._find_topk_peaks(
            centroid_cms, self.max_instances, self.centroid_peak_threshold
        )
        centroids = centroids * (
            self.centroid_output_stride / self.centroid_input_scale
        )

        crops = self._extract_crops(image, centroids)
        crops_flat = crops.reshape(
            batch_size * self.max_instances,
            channels,
            self.crop_size[0],
            self.crop_size[1],
        )

        if self.instance_input_scale != 1.0:
            scaled_h = int(self.crop_size[0] * self.instance_input_scale)
            scaled_w = int(self.crop_size[1] * self.instance_input_scale)
            crops_flat = F.interpolate(
                crops_flat,
                size=(scaled_h, scaled_w),
                mode="bilinear",
                align_corners=False,
            )

        instance_out = self.instance_model(crops_flat)
        instance_cms = self._extract_tensor(
            instance_out, ["centered", "instance", "confmap"]
        )

        crop_peaks, crop_peak_vals = self._find_global_peaks(
            instance_cms, self.instance_peak_threshold
        )
        crop_peaks = crop_peaks * (
            self.instance_output_stride / self.instance_input_scale
        )

        crop_peaks = crop_peaks.reshape(batch_size, self.max_instances, self.n_nodes, 2)
        peak_vals = crop_peak_vals.reshape(batch_size, self.max_instances, self.n_nodes)

        crop_offset = centroids.unsqueeze(2) - image.new_tensor(
            [self.crop_size[1] / 2.0, self.crop_size[0] / 2.0]
        )
        peaks = crop_peaks + crop_offset

        invalid_mask = ~instance_valid
        centroids = centroids.masked_fill(invalid_mask.unsqueeze(-1), 0.0)
        centroid_vals = centroid_vals.masked_fill(invalid_mask, 0.0)
        peaks = peaks.masked_fill(invalid_mask.unsqueeze(-1).unsqueeze(-1), 0.0)
        peak_vals = peak_vals.masked_fill(invalid_mask.unsqueeze(-1), 0.0)

        return {
            "centroids": centroids,
            "centroid_vals": centroid_vals,
            "peaks": peaks,
            "peak_vals": peak_vals,
            "instance_valid": instance_valid,
        }

    def _extract_crops(
        self,
        image: torch.Tensor,
        centroids: torch.Tensor,
    ) -> torch.Tensor:
        """Extract crops around centroids using grid_sample."""
        batch_size, channels, height, width = image.shape
        crop_h, crop_w = self.crop_size
        n_instances = centroids.shape[1]

        scale_x = crop_w / width
        scale_y = crop_h / height
        scale = image.new_tensor([scale_x, scale_y])
        base_grid = self.base_grid.to(device=image.device, dtype=image.dtype)
        scaled_grid = base_grid * scale

        scaled_grid = scaled_grid.unsqueeze(0).unsqueeze(0)
        scaled_grid = scaled_grid.expand(batch_size, n_instances, -1, -1, -1)

        norm_centroids = torch.zeros_like(centroids)
        norm_centroids[..., 0] = (centroids[..., 0] / (width - 1)) * 2 - 1
        norm_centroids[..., 1] = (centroids[..., 1] / (height - 1)) * 2 - 1
        offset = norm_centroids.unsqueeze(2).unsqueeze(2)

        sample_grid = scaled_grid + offset

        image_expanded = image.unsqueeze(1).expand(-1, n_instances, -1, -1, -1)
        image_flat = image_expanded.reshape(
            batch_size * n_instances, channels, height, width
        )
        grid_flat = sample_grid.reshape(batch_size * n_instances, crop_h, crop_w, 2)

        crops_flat = F.grid_sample(
            image_flat,
            grid_flat,
            mode="bilinear",
            padding_mode="zeros",
            align_corners=True,
        )

        crops = crops_flat.reshape(batch_size, n_instances, channels, crop_h, crop_w)
        return crops

__init__(centroid_model, instance_model, max_instances=20, crop_size=(192, 192), centroid_output_stride=2, instance_output_stride=4, centroid_input_scale=1.0, instance_input_scale=1.0, n_nodes=1, centroid_peak_threshold=0.2, instance_peak_threshold=0.2)

Initialize top-down ONNX wrapper.

Parameters:

Name Type Description Default
centroid_model Module

Centroid detection model.

 required
instance_model Module

Instance pose estimation model.

 required
max_instances int

Maximum number of instances to detect.

 20
crop_size Tuple[int, int]

Size of instance crops (height, width).

 (192, 192)
centroid_output_stride int

Centroid model output stride.

 2
instance_output_stride int

Instance model output stride.

 4
centroid_input_scale float

Centroid input scaling factor.

 1.0
instance_input_scale float

Instance input scaling factor.

 1.0
n_nodes int

Number of skeleton nodes.

 1
centroid_peak_threshold float

Minimum confidence for centroid peaks.

 0.2
instance_peak_threshold float

Minimum confidence for instance peaks.

 0.2
Source code in sleap_nn/export/wrappers/topdown.py 
20
21
22
23
24
25
26
27
28
29
30
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
63
64
65
66
67
def __init__(
    self,
    centroid_model: nn.Module,
    instance_model: nn.Module,
    max_instances: int = 20,
    crop_size: Tuple[int, int] = (192, 192),
    centroid_output_stride: int = 2,
    instance_output_stride: int = 4,
    centroid_input_scale: float = 1.0,
    instance_input_scale: float = 1.0,
    n_nodes: int = 1,
    centroid_peak_threshold: float = 0.2,
    instance_peak_threshold: float = 0.2,
) -> None:
    """Initialize top-down ONNX wrapper.

    Args:
        centroid_model: Centroid detection model.
        instance_model: Instance pose estimation model.
        max_instances: Maximum number of instances to detect.
        crop_size: Size of instance crops (height, width).
        centroid_output_stride: Centroid model output stride.
        instance_output_stride: Instance model output stride.
        centroid_input_scale: Centroid input scaling factor.
        instance_input_scale: Instance input scaling factor.
        n_nodes: Number of skeleton nodes.
        centroid_peak_threshold: Minimum confidence for centroid peaks.
        instance_peak_threshold: Minimum confidence for instance peaks.
    """
    super().__init__(centroid_model)
    self.centroid_model = centroid_model
    self.instance_model = instance_model
    self.max_instances = max_instances
    self.crop_size = crop_size
    self.centroid_output_stride = centroid_output_stride
    self.instance_output_stride = instance_output_stride
    self.centroid_input_scale = centroid_input_scale
    self.instance_input_scale = instance_input_scale
    self.n_nodes = n_nodes
    self.centroid_peak_threshold = centroid_peak_threshold
    self.instance_peak_threshold = instance_peak_threshold

    crop_h, crop_w = crop_size
    y_crop = torch.linspace(-1, 1, crop_h, dtype=torch.float32)
    x_crop = torch.linspace(-1, 1, crop_w, dtype=torch.float32)
    grid_y, grid_x = torch.meshgrid(y_crop, x_crop, indexing="ij")
    base_grid = torch.stack([grid_x, grid_y], dim=-1)
    self.register_buffer("base_grid", base_grid, persistent=False)

forward(image)

Run top-down inference and return fixed-size outputs.

Source code in sleap_nn/export/wrappers/topdown.py 
 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
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run top-down inference and return fixed-size outputs."""
    image = self._normalize_uint8(image)
    batch_size, channels, height, width = image.shape

    scaled_image = image
    if self.centroid_input_scale != 1.0:
        scaled_h = int(height * self.centroid_input_scale)
        scaled_w = int(width * self.centroid_input_scale)
        scaled_image = F.interpolate(
            scaled_image,
            size=(scaled_h, scaled_w),
            mode="bilinear",
            align_corners=False,
        )

    centroid_out = self.centroid_model(scaled_image)
    centroid_cms = self._extract_tensor(centroid_out, ["centroid", "confmap"])

    centroids, centroid_vals, instance_valid = self._find_topk_peaks(
        centroid_cms, self.max_instances, self.centroid_peak_threshold
    )
    centroids = centroids * (
        self.centroid_output_stride / self.centroid_input_scale
    )

    crops = self._extract_crops(image, centroids)
    crops_flat = crops.reshape(
        batch_size * self.max_instances,
        channels,
        self.crop_size[0],
        self.crop_size[1],
    )

    if self.instance_input_scale != 1.0:
        scaled_h = int(self.crop_size[0] * self.instance_input_scale)
        scaled_w = int(self.crop_size[1] * self.instance_input_scale)
        crops_flat = F.interpolate(
            crops_flat,
            size=(scaled_h, scaled_w),
            mode="bilinear",
            align_corners=False,
        )

    instance_out = self.instance_model(crops_flat)
    instance_cms = self._extract_tensor(
        instance_out, ["centered", "instance", "confmap"]
    )

    crop_peaks, crop_peak_vals = self._find_global_peaks(
        instance_cms, self.instance_peak_threshold
    )
    crop_peaks = crop_peaks * (
        self.instance_output_stride / self.instance_input_scale
    )

    crop_peaks = crop_peaks.reshape(batch_size, self.max_instances, self.n_nodes, 2)
    peak_vals = crop_peak_vals.reshape(batch_size, self.max_instances, self.n_nodes)

    crop_offset = centroids.unsqueeze(2) - image.new_tensor(
        [self.crop_size[1] / 2.0, self.crop_size[0] / 2.0]
    )
    peaks = crop_peaks + crop_offset

    invalid_mask = ~instance_valid
    centroids = centroids.masked_fill(invalid_mask.unsqueeze(-1), 0.0)
    centroid_vals = centroid_vals.masked_fill(invalid_mask, 0.0)
    peaks = peaks.masked_fill(invalid_mask.unsqueeze(-1).unsqueeze(-1), 0.0)
    peak_vals = peak_vals.masked_fill(invalid_mask.unsqueeze(-1), 0.0)

    return {
        "centroids": centroids,
        "centroid_vals": centroid_vals,
        "peaks": peaks,
        "peak_vals": peak_vals,
        "instance_valid": instance_valid,
    }