Skip to content

wrappers

sleap_nn.export.wrappers

ONNX/TensorRT export wrappers.

Modules:

Name Description
base

Base classes and shared helpers for export wrappers.
bottomup

Bottom-up ONNX wrapper.
bottomup_multiclass

ONNX wrapper for bottom-up multiclass (supervised ID) models.
centered_instance

Centered-instance ONNX wrapper.
centroid

Centroid ONNX wrapper.
embedding

ONNX export wrapper for the embedding model type.
single_instance

Single-instance ONNX wrapper.
topdown

Top-down ONNX wrapper.
topdown_multiclass

ONNX wrapper for top-down multiclass (supervised ID) models.

Classes:

Name Description
BaseExportWrapper

Base class for ONNX-exportable wrappers.
BottomUpMultiClassONNXWrapper

ONNX-exportable wrapper for bottom-up multiclass (supervised ID) models.
BottomUpONNXWrapper

ONNX-exportable wrapper for bottom-up inference up to PAF scoring.
CenteredInstanceONNXWrapper

ONNX-exportable wrapper for centered-instance models.
CentroidONNXWrapper

ONNX-exportable wrapper for centroid models.
EmbeddingONNXWrapper

Wrap an embedding model for ONNX export: crop -> appearance vector.
SingleInstanceONNXWrapper

ONNX-exportable wrapper for single-instance models.
TopDownMultiClassCombinedONNXWrapper

ONNX-exportable wrapper for combined centroid + multiclass instance models.
TopDownMultiClassONNXWrapper

ONNX-exportable wrapper for top-down multiclass (supervised ID) models.
TopDownONNXWrapper

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

BaseExportWrapper

Bases: Module

Base class for ONNX-exportable wrappers.

Methods:

Name Description
__init__

Initialize wrapper with the underlying model.
Source code in sleap_nn/export/wrappers/base.py 
 12
 13
 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
class BaseExportWrapper(nn.Module):
    """Base class for ONNX-exportable wrappers."""

    def __init__(self, model: nn.Module):
        """Initialize wrapper with the underlying model.

        Args:
            model: The PyTorch model to wrap for export.
        """
        super().__init__()
        self.model = model

    @staticmethod
    def _normalize_uint8(image: torch.Tensor) -> torch.Tensor:
        """Normalize unnormalized uint8 (or [0, 255] float) images to [0, 1]."""
        if image.dtype != torch.float32:
            image = image.float()
        return image / 255.0

    @staticmethod
    def _extract_tensor(output, key_hints: Iterable[str]) -> torch.Tensor:
        if isinstance(output, dict):
            for key in output:
                for hint in key_hints:
                    if hint.lower() in key.lower():
                        return output[key]
            return next(iter(output.values()))
        return output

    @staticmethod
    def _neighbor_max(x: torch.Tensor) -> torch.Tensor:
        """Compute max of 8 neighbors excluding center pixel.

        Uses -inf padding to match PyTorch dilation semantics (confmap heads
        are identity-activated, so negative values are possible).

        All ops (F.pad, slicing, torch.max) export cleanly to ONNX.
        """
        p = F.pad(x, [1, 1, 1, 1], mode="constant", value=float("-inf"))
        # 8 shifted views (excluding center)
        tl = p[:, :, :-2, :-2]  # top-left
        tc = p[:, :, :-2, 1:-1]  # top-center
        tr = p[:, :, :-2, 2:]  # top-right
        ml = p[:, :, 1:-1, :-2]  # middle-left
        mr = p[:, :, 1:-1, 2:]  # middle-right
        bl = p[:, :, 2:, :-2]  # bottom-left
        bc = p[:, :, 2:, 1:-1]  # bottom-center
        br = p[:, :, 2:, 2:]  # bottom-right
        return torch.max(
            torch.max(
                torch.max(
                    torch.max(torch.max(torch.max(torch.max(tl, tc), tr), ml), mr), bl
                ),
                bc,
            ),
            br,
        )

    @staticmethod
    def _find_topk_peaks(
        confmaps: torch.Tensor, k: int, peak_threshold: float = 0.2
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Top-K peak finding with center-excluded neighbor-max NMS.

        Matches PyTorch path semantics: strict inequality (center > all
        neighbors) and configurable confidence threshold.
        """
        batch_size, _, height, width = confmaps.shape
        neighbor_max = BaseExportWrapper._neighbor_max(confmaps)
        is_peak = (confmaps > neighbor_max) & (confmaps > peak_threshold)

        confmaps_flat = confmaps.reshape(batch_size, height * width)
        is_peak_flat = is_peak.reshape(batch_size, height * width)
        masked = torch.where(
            is_peak_flat, confmaps_flat, torch.full_like(confmaps_flat, -1e9)
        )
        values, indices = torch.topk(masked, k=k, dim=1)

        y = indices // width
        x = indices % width
        peaks = torch.stack([x.float(), y.float()], dim=-1)
        valid = values > peak_threshold
        return peaks, values, valid

    @staticmethod
    def _find_topk_peaks_per_node(
        confmaps: torch.Tensor, k: int, peak_threshold: float = 0.2
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Top-K peak finding per channel with center-excluded neighbor-max NMS.

        Matches PyTorch path semantics: strict inequality (center > all
        neighbors) and configurable confidence threshold.
        """
        batch_size, n_nodes, height, width = confmaps.shape
        neighbor_max = BaseExportWrapper._neighbor_max(confmaps)
        is_peak = (confmaps > neighbor_max) & (confmaps > peak_threshold)

        confmaps_flat = confmaps.reshape(batch_size, n_nodes, height * width)
        is_peak_flat = is_peak.reshape(batch_size, n_nodes, height * width)
        masked = torch.where(
            is_peak_flat, confmaps_flat, torch.full_like(confmaps_flat, -1e9)
        )
        values, indices = torch.topk(masked, k=k, dim=2)

        y = indices // width
        x = indices % width
        peaks = torch.stack([x.float(), y.float()], dim=-1)
        valid = values > peak_threshold
        return peaks, values, valid

    @staticmethod
    def _find_global_peaks(
        confmaps: torch.Tensor, peak_threshold: float = 0.2
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Find global maxima per channel with threshold.

        Peaks with confidence below threshold are set to NaN coordinates and
        zero confidence, matching ``find_global_peaks_rough`` in the PyTorch
        path.
        """
        batch_size, channels, height, width = confmaps.shape
        flat = confmaps.reshape(batch_size, channels, height * width)
        values, indices = flat.max(dim=-1)
        y = indices // width
        x = indices % width
        peaks = torch.stack([x.float(), y.float()], dim=-1)

        below = values < peak_threshold
        peaks = peaks.masked_fill(below.unsqueeze(-1), float("nan"))
        values = values.masked_fill(below, 0.0)

        return peaks, values

__init__(model)

Initialize wrapper with the underlying model.

Parameters:

Name Type Description Default
model Module

The PyTorch model to wrap for export.

 required
Source code in sleap_nn/export/wrappers/base.py 
15
16
17
18
19
20
21
22
def __init__(self, model: nn.Module):
    """Initialize wrapper with the underlying model.

    Args:
        model: The PyTorch model to wrap for export.
    """
    super().__init__()
    self.model = model

BottomUpMultiClassONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for bottom-up multiclass (supervised ID) models.

This wrapper handles models that output both confidence maps for keypoint detection and class maps for identity classification. Unlike PAF-based bottom-up models, multiclass models use class maps to assign identity to each detected peak, then group peaks by identity.

The wrapper performs: 1. Peak detection in confidence maps (GPU) 2. Class probability sampling at peak locations (GPU) 3. Returns fixed-size tensors for CPU-side grouping

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

Attributes:

Name Type Description
model

The underlying PyTorch model.
n_nodes

Number of keypoint nodes in the skeleton.
n_classes

Number of identity classes.
max_peaks_per_node

Maximum number of peaks to detect per node.
cms_output_stride

Output stride of the confidence map head.
class_maps_output_stride

Output stride of the class maps head.
input_scale

Scale factor applied to input images before inference.

Methods:

Name Description
__init__

Initialize the wrapper.
forward

Run bottom-up multiclass inference.
Source code in sleap_nn/export/wrappers/bottomup_multiclass.py 
 12
 13
 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
189
190
191
192
193
194
195
196
197
198
class BottomUpMultiClassONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for bottom-up multiclass (supervised ID) models.

    This wrapper handles models that output both confidence maps for keypoint
    detection and class maps for identity classification. Unlike PAF-based
    bottom-up models, multiclass models use class maps to assign identity to
    each detected peak, then group peaks by identity.

    The wrapper performs:
    1. Peak detection in confidence maps (GPU)
    2. Class probability sampling at peak locations (GPU)
    3. Returns fixed-size tensors for CPU-side grouping

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

    Attributes:
        model: The underlying PyTorch model.
        n_nodes: Number of keypoint nodes in the skeleton.
        n_classes: Number of identity classes.
        max_peaks_per_node: Maximum number of peaks to detect per node.
        cms_output_stride: Output stride of the confidence map head.
        class_maps_output_stride: Output stride of the class maps head.
        input_scale: Scale factor applied to input images before inference.
    """

    def __init__(
        self,
        model: nn.Module,
        n_nodes: int,
        n_classes: int = 2,
        max_peaks_per_node: int = 20,
        cms_output_stride: int = 4,
        class_maps_output_stride: int = 8,
        input_scale: float = 1.0,
        peak_threshold: float = 0.2,
    ):
        """Initialize the wrapper.

        Args:
            model: The underlying PyTorch model.
            n_nodes: Number of keypoint nodes.
            n_classes: Number of identity classes (e.g., 2 for male/female).
            max_peaks_per_node: Maximum peaks per node to detect.
            cms_output_stride: Output stride of confidence maps.
            class_maps_output_stride: Output stride of class maps.
            input_scale: Scale factor for input images.
            peak_threshold: Minimum confidence for a peak to be considered valid.
        """
        super().__init__(model)
        self.n_nodes = n_nodes
        self.n_classes = n_classes
        self.max_peaks_per_node = max_peaks_per_node
        self.cms_output_stride = cms_output_stride
        self.class_maps_output_stride = class_maps_output_stride
        self.input_scale = input_scale
        self.peak_threshold = peak_threshold

    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run bottom-up multiclass inference.

        Args:
            image: Input image tensor of shape (batch, channels, height, width).
                   Expected to be uint8 in [0, 255].

        Returns:
            Dictionary with keys:
                - "peaks": Detected peak coordinates (batch, n_nodes, max_peaks, 2).
                    Coordinates are in input image space (x, y).
                - "peak_vals": Peak confidence values (batch, n_nodes, max_peaks).
                - "peak_mask": Boolean mask for valid peaks (batch, n_nodes, max_peaks).
                - "class_probs": Class probabilities at each peak location
                    (batch, n_nodes, max_peaks, n_classes).

            Postprocessing on CPU uses `classify_peaks_from_maps()` to group
            peaks by identity using Hungarian matching.
        """
        # Normalize uint8 [0, 255] to float32 [0, 1]
        image = self._normalize_uint8(image)

        # Apply input scaling if needed
        if self.input_scale != 1.0:
            height = int(image.shape[-2] * self.input_scale)
            width = int(image.shape[-1] * self.input_scale)
            image = F.interpolate(
                image, size=(height, width), mode="bilinear", align_corners=False
            )

        batch_size = image.shape[0]

        # Forward pass
        out = self.model(image)

        # Extract outputs
        # Note: Use "classmaps" as a single hint to avoid "map" matching "confmaps"
        confmaps = self._extract_tensor(out, ["confmap", "multiinstance"])
        class_maps = self._extract_tensor(out, ["classmaps", "classmapshead"])

        # Find top-k peaks per node
        peaks, peak_vals, peak_mask = self._find_topk_peaks_per_node(
            confmaps, self.max_peaks_per_node, self.peak_threshold
        )

        # Scale peaks to input image space
        peaks = peaks * self.cms_output_stride

        # Sample class maps at peak locations
        class_probs = self._sample_class_maps_at_peaks(class_maps, peaks, peak_mask)

        # Scale peaks for output (accounting for input scale)
        if self.input_scale != 1.0:
            peaks = peaks / self.input_scale

        return {
            "peaks": peaks,
            "peak_vals": peak_vals,
            "peak_mask": peak_mask,
            "class_probs": class_probs,
        }

    def _sample_class_maps_at_peaks(
        self,
        class_maps: torch.Tensor,
        peaks: torch.Tensor,
        peak_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Sample class map values at peak locations.

        Args:
            class_maps: Class maps of shape (batch, n_classes, height, width).
            peaks: Peak coordinates in cms_output_stride space,
                shape (batch, n_nodes, max_peaks, 2) in (x, y) order.
            peak_mask: Boolean mask for valid peaks (batch, n_nodes, max_peaks).

        Returns:
            Class probabilities at each peak location,
            shape (batch, n_nodes, max_peaks, n_classes).
        """
        batch_size, n_classes, cm_height, cm_width = class_maps.shape
        _, n_nodes, max_peaks, _ = peaks.shape
        device = peaks.device

        # Initialize output tensor
        class_probs = torch.zeros(
            (batch_size, n_nodes, max_peaks, n_classes),
            device=device,
            dtype=class_maps.dtype,
        )

        # Convert peak coordinates to class map space
        # peaks are in full image space (after cms_output_stride scaling)
        peaks_cm = peaks / self.class_maps_output_stride

        # Clamp coordinates to valid range
        peaks_cm_x = peaks_cm[..., 0].clamp(0, cm_width - 1)
        peaks_cm_y = peaks_cm[..., 1].clamp(0, cm_height - 1)

        # Use grid_sample for bilinear interpolation
        # Normalize coordinates to [-1, 1] for grid_sample
        grid_x = (peaks_cm_x / (cm_width - 1)) * 2 - 1
        grid_y = (peaks_cm_y / (cm_height - 1)) * 2 - 1

        # Reshape for grid_sample: (batch, n_nodes * max_peaks, 1, 2)
        grid = torch.stack([grid_x, grid_y], dim=-1)
        grid_flat = grid.reshape(batch_size, n_nodes * max_peaks, 1, 2)

        # Sample class maps: (batch, n_classes, n_nodes * max_peaks, 1)
        sampled = F.grid_sample(
            class_maps,
            grid_flat,
            mode="bilinear",
            padding_mode="zeros",
            align_corners=True,
        )

        # Reshape to (batch, n_nodes, max_peaks, n_classes)
        sampled = sampled.squeeze(-1)  # (batch, n_classes, n_nodes * max_peaks)
        sampled = sampled.permute(0, 2, 1)  # (batch, n_nodes * max_peaks, n_classes)
        sampled = sampled.reshape(batch_size, n_nodes, max_peaks, n_classes)

        # Apply softmax to get probabilities (optional - depends on training)
        # For now, return raw values as the grouping function expects logits
        class_probs = sampled

        # Mask invalid peaks
        class_probs = class_probs * peak_mask.unsqueeze(-1).float()

        return class_probs

__init__(model, n_nodes, n_classes=2, max_peaks_per_node=20, cms_output_stride=4, class_maps_output_stride=8, input_scale=1.0, peak_threshold=0.2)

Initialize the wrapper.

Parameters:

Name Type Description Default
model Module

The underlying PyTorch model.

 required
n_nodes int

Number of keypoint nodes.

 required
n_classes int

Number of identity classes (e.g., 2 for male/female).

 2
max_peaks_per_node int

Maximum peaks per node to detect.

 20
cms_output_stride int

Output stride of confidence maps.

 4
class_maps_output_stride int

Output stride of class maps.

 8
input_scale float

Scale factor for input images.

 1.0
peak_threshold float

Minimum confidence for a peak to be considered valid.

 0.2
Source code in sleap_nn/export/wrappers/bottomup_multiclass.py 
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,
    model: nn.Module,
    n_nodes: int,
    n_classes: int = 2,
    max_peaks_per_node: int = 20,
    cms_output_stride: int = 4,
    class_maps_output_stride: int = 8,
    input_scale: float = 1.0,
    peak_threshold: float = 0.2,
):
    """Initialize the wrapper.

    Args:
        model: The underlying PyTorch model.
        n_nodes: Number of keypoint nodes.
        n_classes: Number of identity classes (e.g., 2 for male/female).
        max_peaks_per_node: Maximum peaks per node to detect.
        cms_output_stride: Output stride of confidence maps.
        class_maps_output_stride: Output stride of class maps.
        input_scale: Scale factor for input images.
        peak_threshold: Minimum confidence for a peak to be considered valid.
    """
    super().__init__(model)
    self.n_nodes = n_nodes
    self.n_classes = n_classes
    self.max_peaks_per_node = max_peaks_per_node
    self.cms_output_stride = cms_output_stride
    self.class_maps_output_stride = class_maps_output_stride
    self.input_scale = input_scale
    self.peak_threshold = peak_threshold

forward(image)

Run bottom-up multiclass inference.

Parameters:

Name Type Description Default
image Tensor

Input image tensor of shape (batch, channels, height, width). Expected to be uint8 in [0, 255].

 required

Returns:

Type Description
Dict[str, Tensor]

Dictionary with keys: - "peaks": Detected peak coordinates (batch, n_nodes, max_peaks, 2). Coordinates are in input image space (x, y). - "peak_vals": Peak confidence values (batch, n_nodes, max_peaks). - "peak_mask": Boolean mask for valid peaks (batch, n_nodes, max_peaks). - "class_probs": Class probabilities at each peak location (batch, n_nodes, max_peaks, n_classes).

Postprocessing on CPU uses classify_peaks_from_maps() to group peaks by identity using Hungarian matching.
Source code in sleap_nn/export/wrappers/bottomup_multiclass.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
def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run bottom-up multiclass inference.

    Args:
        image: Input image tensor of shape (batch, channels, height, width).
               Expected to be uint8 in [0, 255].

    Returns:
        Dictionary with keys:
            - "peaks": Detected peak coordinates (batch, n_nodes, max_peaks, 2).
                Coordinates are in input image space (x, y).
            - "peak_vals": Peak confidence values (batch, n_nodes, max_peaks).
            - "peak_mask": Boolean mask for valid peaks (batch, n_nodes, max_peaks).
            - "class_probs": Class probabilities at each peak location
                (batch, n_nodes, max_peaks, n_classes).

        Postprocessing on CPU uses `classify_peaks_from_maps()` to group
        peaks by identity using Hungarian matching.
    """
    # Normalize uint8 [0, 255] to float32 [0, 1]
    image = self._normalize_uint8(image)

    # Apply input scaling if needed
    if self.input_scale != 1.0:
        height = int(image.shape[-2] * self.input_scale)
        width = int(image.shape[-1] * self.input_scale)
        image = F.interpolate(
            image, size=(height, width), mode="bilinear", align_corners=False
        )

    batch_size = image.shape[0]

    # Forward pass
    out = self.model(image)

    # Extract outputs
    # Note: Use "classmaps" as a single hint to avoid "map" matching "confmaps"
    confmaps = self._extract_tensor(out, ["confmap", "multiinstance"])
    class_maps = self._extract_tensor(out, ["classmaps", "classmapshead"])

    # Find top-k peaks per node
    peaks, peak_vals, peak_mask = self._find_topk_peaks_per_node(
        confmaps, self.max_peaks_per_node, self.peak_threshold
    )

    # Scale peaks to input image space
    peaks = peaks * self.cms_output_stride

    # Sample class maps at peak locations
    class_probs = self._sample_class_maps_at_peaks(class_maps, peaks, peak_mask)

    # Scale peaks for output (accounting for input scale)
    if self.input_scale != 1.0:
        peaks = peaks / self.input_scale

    return {
        "peaks": peaks,
        "peak_vals": peak_vals,
        "peak_mask": peak_mask,
        "class_probs": class_probs,
    }

BottomUpONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for bottom-up inference up to PAF scoring.

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

Methods:

Name Description
__init__

Initialize bottom-up ONNX wrapper.
forward

Run bottom-up inference and return fixed-size outputs.
Source code in sleap_nn/export/wrappers/bottomup.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
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
class BottomUpONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for bottom-up inference up to PAF scoring.

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

    def __init__(
        self,
        model: nn.Module,
        skeleton_edges: list,
        n_nodes: int,
        max_peaks_per_node: int = 20,
        n_line_points: int = 10,
        cms_output_stride: int = 4,
        pafs_output_stride: int = 8,
        max_edge_length_ratio: float = 0.25,
        dist_penalty_weight: float = 1.0,
        input_scale: float = 1.0,
        peak_threshold: float = 0.2,
    ) -> None:
        """Initialize bottom-up ONNX wrapper.

        Args:
            model: Bottom-up model producing confidence maps and PAFs.
            skeleton_edges: List of (src, dst) edge tuples defining skeleton.
            n_nodes: Number of nodes in the skeleton.
            max_peaks_per_node: Maximum peaks to detect per node type.
            n_line_points: Points to sample along PAF edges.
            cms_output_stride: Confidence map output stride.
            pafs_output_stride: PAF output stride.
            max_edge_length_ratio: Maximum edge length as ratio of image size.
            dist_penalty_weight: Weight for distance penalty in scoring.
            input_scale: Input scaling factor.
            peak_threshold: Minimum confidence for a peak to be considered valid.
        """
        super().__init__(model)
        self.n_nodes = n_nodes
        self.n_edges = len(skeleton_edges)
        self.max_peaks_per_node = max_peaks_per_node
        self.n_line_points = n_line_points
        self.cms_output_stride = cms_output_stride
        self.pafs_output_stride = pafs_output_stride
        self.max_edge_length_ratio = max_edge_length_ratio
        self.dist_penalty_weight = dist_penalty_weight
        self.input_scale = input_scale
        self.peak_threshold = peak_threshold

        edge_src = torch.tensor([e[0] for e in skeleton_edges], dtype=torch.long)
        edge_dst = torch.tensor([e[1] for e in skeleton_edges], dtype=torch.long)
        self.register_buffer("edge_src", edge_src)
        self.register_buffer("edge_dst", edge_dst)

        line_samples = torch.linspace(0, 1, n_line_points, dtype=torch.float32)
        self.register_buffer("line_samples", line_samples)

    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run bottom-up inference and return fixed-size outputs.

        Note: confmaps and pafs are NOT returned to avoid D2H transfer bottleneck.
        Peak detection and PAF scoring are performed on GPU within this wrapper.
        """
        image = self._normalize_uint8(image)
        if self.input_scale != 1.0:
            height = int(image.shape[-2] * self.input_scale)
            width = int(image.shape[-1] * self.input_scale)
            image = F.interpolate(
                image, size=(height, width), mode="bilinear", align_corners=False
            )

        batch_size, _, height, width = image.shape

        out = self.model(image)
        if isinstance(out, dict):
            confmaps = self._extract_tensor(out, ["confmap", "multiinstance"])
            pafs = self._extract_tensor(out, ["paf", "affinity"])
        else:
            confmaps, pafs = out[:2]

        peaks, peak_vals, peak_mask = self._find_topk_peaks_per_node(
            confmaps, self.max_peaks_per_node, self.peak_threshold
        )

        peaks = peaks * self.cms_output_stride

        # Compute max_edge_length to match PyTorch implementation:
        # max_edge_length = ratio * max(paf_dims) * pafs_stride
        # PAFs shape is (batch, 2*edges, H, W)
        _, n_paf_channels, paf_height, paf_width = pafs.shape
        max_paf_dim = max(n_paf_channels, paf_height, paf_width)
        max_edge_length = torch.tensor(
            self.max_edge_length_ratio * max_paf_dim * self.pafs_output_stride,
            dtype=peaks.dtype,
            device=peaks.device,
        )

        line_scores, candidate_mask = self._score_all_candidates(
            pafs, peaks, peak_mask, max_edge_length
        )

        # Only return final outputs needed for CPU-side grouping.
        # Do NOT return confmaps/pafs - they are large (~29 MB/batch) and
        # cause D2H transfer bottleneck. Peak detection and PAF scoring
        # are already done on GPU above.
        return {
            "peaks": peaks,
            "peak_vals": peak_vals,
            "peak_mask": peak_mask,
            "line_scores": line_scores,
            "candidate_mask": candidate_mask,
        }

    def _score_all_candidates(
        self,
        pafs: torch.Tensor,
        peaks: torch.Tensor,
        peak_mask: torch.Tensor,
        max_edge_length: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Score all K*K candidate connections for each edge."""
        batch_size = peaks.shape[0]
        k = self.max_peaks_per_node
        n_edges = self.n_edges

        _, _, paf_height, paf_width = pafs.shape

        src_peaks = peaks[:, self.edge_src, :, :]
        dst_peaks = peaks[:, self.edge_dst, :, :]

        src_mask = peak_mask[:, self.edge_src, :]
        dst_mask = peak_mask[:, self.edge_dst, :]

        src_peaks_exp = src_peaks.unsqueeze(3).expand(-1, -1, -1, k, -1)
        dst_peaks_exp = dst_peaks.unsqueeze(2).expand(-1, -1, k, -1, -1)

        src_mask_exp = src_mask.unsqueeze(3).expand(-1, -1, -1, k)
        dst_mask_exp = dst_mask.unsqueeze(2).expand(-1, -1, k, -1)
        candidate_mask = src_mask_exp & dst_mask_exp

        src_peaks_flat = src_peaks_exp.reshape(batch_size, n_edges, k * k, 2)
        dst_peaks_flat = dst_peaks_exp.reshape(batch_size, n_edges, k * k, 2)
        candidate_mask_flat = candidate_mask.reshape(batch_size, n_edges, k * k)

        spatial_vecs = dst_peaks_flat - src_peaks_flat
        spatial_lengths = torch.norm(spatial_vecs, dim=-1, keepdim=True).clamp(min=1e-6)
        spatial_vecs_norm = spatial_vecs / spatial_lengths

        line_samples = self.line_samples.view(1, 1, 1, -1, 1)
        src_exp = src_peaks_flat.unsqueeze(3)
        dst_exp = dst_peaks_flat.unsqueeze(3)
        line_points = src_exp + line_samples * (dst_exp - src_exp)

        line_points_paf = line_points / self.pafs_output_stride
        line_x = line_points_paf[..., 0].clamp(0, paf_width - 1)
        line_y = line_points_paf[..., 1].clamp(0, paf_height - 1)

        line_scores = self._sample_and_score_lines(
            pafs,
            line_x,
            line_y,
            spatial_vecs_norm,
            spatial_lengths.squeeze(-1),
            max_edge_length,
        )

        line_scores = line_scores.masked_fill(~candidate_mask_flat, -2.0)
        return line_scores, candidate_mask_flat

    def _sample_and_score_lines(
        self,
        pafs: torch.Tensor,
        line_x: torch.Tensor,
        line_y: torch.Tensor,
        spatial_vecs_norm: torch.Tensor,
        spatial_lengths: torch.Tensor,
        max_edge_length: torch.Tensor,
    ) -> torch.Tensor:
        """Sample PAF values along lines and compute scores."""
        batch_size, n_edges, k2, n_points = line_x.shape
        _, _, paf_height, paf_width = pafs.shape

        all_scores = []
        for edge_idx in range(n_edges):
            paf_x = pafs[:, 2 * edge_idx, :, :]
            paf_y = pafs[:, 2 * edge_idx + 1, :, :]

            lx = line_x[:, edge_idx, :, :]
            ly = line_y[:, edge_idx, :, :]

            lx_norm = (lx / (paf_width - 1)) * 2 - 1
            ly_norm = (ly / (paf_height - 1)) * 2 - 1

            grid = torch.stack([lx_norm, ly_norm], dim=-1)

            paf_x_samples = F.grid_sample(
                paf_x.unsqueeze(1),
                grid,
                mode="bilinear",
                padding_mode="zeros",
                align_corners=True,
            ).squeeze(1)

            paf_y_samples = F.grid_sample(
                paf_y.unsqueeze(1),
                grid,
                mode="bilinear",
                padding_mode="zeros",
                align_corners=True,
            ).squeeze(1)

            paf_samples = torch.stack([paf_x_samples, paf_y_samples], dim=-1)
            disp_vec = spatial_vecs_norm[:, edge_idx, :, :]

            dot_products = (paf_samples * disp_vec.unsqueeze(2)).sum(dim=-1)
            mean_scores = dot_products.mean(dim=-1)

            edge_lengths = spatial_lengths[:, edge_idx, :]
            dist_penalty = self._compute_distance_penalty(edge_lengths, max_edge_length)

            all_scores.append(mean_scores + dist_penalty)

        return torch.stack(all_scores, dim=1)

    def _compute_distance_penalty(
        self, distances: torch.Tensor, max_edge_length: torch.Tensor
    ) -> torch.Tensor:
        """Compute distance penalty for edge candidates.

        Matches the PyTorch implementation in sleap_nn.inference.paf_grouping.
        Penalty is 0 when distance <= max_edge_length, and negative when longer.
        """
        # Match PyTorch: penalty = clamp((max_edge_length / distance) - 1, max=0) * weight
        penalty = torch.clamp((max_edge_length / distances) - 1, max=0)
        return penalty * self.dist_penalty_weight

__init__(model, skeleton_edges, n_nodes, max_peaks_per_node=20, n_line_points=10, cms_output_stride=4, pafs_output_stride=8, max_edge_length_ratio=0.25, dist_penalty_weight=1.0, input_scale=1.0, peak_threshold=0.2)

Initialize bottom-up ONNX wrapper.

Parameters:

Name Type Description Default
model Module

Bottom-up model producing confidence maps and PAFs.

 required
skeleton_edges list

List of (src, dst) edge tuples defining skeleton.

 required
n_nodes int

Number of nodes in the skeleton.

 required
max_peaks_per_node int

Maximum peaks to detect per node type.

 20
n_line_points int

Points to sample along PAF edges.

 10
cms_output_stride int

Confidence map output stride.

 4
pafs_output_stride int

PAF output stride.

 8
max_edge_length_ratio float

Maximum edge length as ratio of image size.

 0.25
dist_penalty_weight float

Weight for distance penalty in scoring.

 1.0
input_scale float

Input scaling factor.

 1.0
peak_threshold float

Minimum confidence for a peak to be considered valid.

 0.2
Source code in sleap_nn/export/wrappers/bottomup.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,
    model: nn.Module,
    skeleton_edges: list,
    n_nodes: int,
    max_peaks_per_node: int = 20,
    n_line_points: int = 10,
    cms_output_stride: int = 4,
    pafs_output_stride: int = 8,
    max_edge_length_ratio: float = 0.25,
    dist_penalty_weight: float = 1.0,
    input_scale: float = 1.0,
    peak_threshold: float = 0.2,
) -> None:
    """Initialize bottom-up ONNX wrapper.

    Args:
        model: Bottom-up model producing confidence maps and PAFs.
        skeleton_edges: List of (src, dst) edge tuples defining skeleton.
        n_nodes: Number of nodes in the skeleton.
        max_peaks_per_node: Maximum peaks to detect per node type.
        n_line_points: Points to sample along PAF edges.
        cms_output_stride: Confidence map output stride.
        pafs_output_stride: PAF output stride.
        max_edge_length_ratio: Maximum edge length as ratio of image size.
        dist_penalty_weight: Weight for distance penalty in scoring.
        input_scale: Input scaling factor.
        peak_threshold: Minimum confidence for a peak to be considered valid.
    """
    super().__init__(model)
    self.n_nodes = n_nodes
    self.n_edges = len(skeleton_edges)
    self.max_peaks_per_node = max_peaks_per_node
    self.n_line_points = n_line_points
    self.cms_output_stride = cms_output_stride
    self.pafs_output_stride = pafs_output_stride
    self.max_edge_length_ratio = max_edge_length_ratio
    self.dist_penalty_weight = dist_penalty_weight
    self.input_scale = input_scale
    self.peak_threshold = peak_threshold

    edge_src = torch.tensor([e[0] for e in skeleton_edges], dtype=torch.long)
    edge_dst = torch.tensor([e[1] for e in skeleton_edges], dtype=torch.long)
    self.register_buffer("edge_src", edge_src)
    self.register_buffer("edge_dst", edge_dst)

    line_samples = torch.linspace(0, 1, n_line_points, dtype=torch.float32)
    self.register_buffer("line_samples", line_samples)

forward(image)

Run bottom-up inference and return fixed-size outputs.

Note: confmaps and pafs are NOT returned to avoid D2H transfer bottleneck. Peak detection and PAF scoring are performed on GPU within this wrapper.

Source code in sleap_nn/export/wrappers/bottomup.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
def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run bottom-up inference and return fixed-size outputs.

    Note: confmaps and pafs are NOT returned to avoid D2H transfer bottleneck.
    Peak detection and PAF scoring are performed on GPU within this wrapper.
    """
    image = self._normalize_uint8(image)
    if self.input_scale != 1.0:
        height = int(image.shape[-2] * self.input_scale)
        width = int(image.shape[-1] * self.input_scale)
        image = F.interpolate(
            image, size=(height, width), mode="bilinear", align_corners=False
        )

    batch_size, _, height, width = image.shape

    out = self.model(image)
    if isinstance(out, dict):
        confmaps = self._extract_tensor(out, ["confmap", "multiinstance"])
        pafs = self._extract_tensor(out, ["paf", "affinity"])
    else:
        confmaps, pafs = out[:2]

    peaks, peak_vals, peak_mask = self._find_topk_peaks_per_node(
        confmaps, self.max_peaks_per_node, self.peak_threshold
    )

    peaks = peaks * self.cms_output_stride

    # Compute max_edge_length to match PyTorch implementation:
    # max_edge_length = ratio * max(paf_dims) * pafs_stride
    # PAFs shape is (batch, 2*edges, H, W)
    _, n_paf_channels, paf_height, paf_width = pafs.shape
    max_paf_dim = max(n_paf_channels, paf_height, paf_width)
    max_edge_length = torch.tensor(
        self.max_edge_length_ratio * max_paf_dim * self.pafs_output_stride,
        dtype=peaks.dtype,
        device=peaks.device,
    )

    line_scores, candidate_mask = self._score_all_candidates(
        pafs, peaks, peak_mask, max_edge_length
    )

    # Only return final outputs needed for CPU-side grouping.
    # Do NOT return confmaps/pafs - they are large (~29 MB/batch) and
    # cause D2H transfer bottleneck. Peak detection and PAF scoring
    # are already done on GPU above.
    return {
        "peaks": peaks,
        "peak_vals": peak_vals,
        "peak_mask": peak_mask,
        "line_scores": line_scores,
        "candidate_mask": candidate_mask,
    }

CenteredInstanceONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for centered-instance models.

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

Methods:

Name Description
__init__

Initialize centered instance ONNX wrapper.
forward

Run centered-instance inference on crops.
Source code in sleap_nn/export/wrappers/centered_instance.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
class CenteredInstanceONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for centered-instance models.

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

    def __init__(
        self,
        model: nn.Module,
        output_stride: int = 4,
        input_scale: float = 1.0,
        peak_threshold: float = 0.2,
    ):
        """Initialize centered instance ONNX wrapper.

        Args:
            model: Centered instance model for pose estimation.
            output_stride: Output stride for confidence maps.
            input_scale: Input scaling factor.
            peak_threshold: Minimum confidence for a peak to be considered valid.
        """
        super().__init__(model)
        self.output_stride = output_stride
        self.input_scale = input_scale
        self.peak_threshold = peak_threshold

    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run centered-instance inference on crops."""
        image = self._normalize_uint8(image)
        if self.input_scale != 1.0:
            height = int(image.shape[-2] * self.input_scale)
            width = int(image.shape[-1] * self.input_scale)
            image = F.interpolate(
                image, size=(height, width), mode="bilinear", align_corners=False
            )

        confmaps = self._extract_tensor(
            self.model(image), ["centered", "instance", "confmap"]
        )
        peaks, values = self._find_global_peaks(confmaps, self.peak_threshold)
        peaks = peaks * (self.output_stride / self.input_scale)

        return {
            "peaks": peaks,
            "peak_vals": values,
        }

__init__(model, output_stride=4, input_scale=1.0, peak_threshold=0.2)

Initialize centered instance ONNX wrapper.

Parameters:

Name Type Description Default
model Module

Centered instance model for pose estimation.

 required
output_stride int

Output stride for confidence maps.

 4
input_scale float

Input scaling factor.

 1.0
peak_threshold float

Minimum confidence for a peak to be considered valid.

 0.2
Source code in sleap_nn/export/wrappers/centered_instance.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def __init__(
    self,
    model: nn.Module,
    output_stride: int = 4,
    input_scale: float = 1.0,
    peak_threshold: float = 0.2,
):
    """Initialize centered instance ONNX wrapper.

    Args:
        model: Centered instance model for pose estimation.
        output_stride: Output stride for confidence maps.
        input_scale: Input scaling factor.
        peak_threshold: Minimum confidence for a peak to be considered valid.
    """
    super().__init__(model)
    self.output_stride = output_stride
    self.input_scale = input_scale
    self.peak_threshold = peak_threshold

forward(image)

Run centered-instance inference on crops.

Source code in sleap_nn/export/wrappers/centered_instance.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run centered-instance inference on crops."""
    image = self._normalize_uint8(image)
    if self.input_scale != 1.0:
        height = int(image.shape[-2] * self.input_scale)
        width = int(image.shape[-1] * self.input_scale)
        image = F.interpolate(
            image, size=(height, width), mode="bilinear", align_corners=False
        )

    confmaps = self._extract_tensor(
        self.model(image), ["centered", "instance", "confmap"]
    )
    peaks, values = self._find_global_peaks(confmaps, self.peak_threshold)
    peaks = peaks * (self.output_stride / self.input_scale)

    return {
        "peaks": peaks,
        "peak_vals": values,
    }

CentroidONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for centroid models.

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

Methods:

Name Description
__init__

Initialize centroid ONNX wrapper.
forward

Run centroid inference and return fixed-size outputs.
Source code in sleap_nn/export/wrappers/centroid.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
class CentroidONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for centroid models.

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

    def __init__(
        self,
        model: nn.Module,
        max_instances: int = 20,
        output_stride: int = 2,
        input_scale: float = 1.0,
        peak_threshold: float = 0.2,
    ):
        """Initialize centroid ONNX wrapper.

        Args:
            model: Centroid detection model.
            max_instances: Maximum number of instances to detect.
            output_stride: Output stride for confidence maps.
            input_scale: Input scaling factor.
            peak_threshold: Minimum confidence for a peak to be considered valid.
        """
        super().__init__(model)
        self.max_instances = max_instances
        self.output_stride = output_stride
        self.input_scale = input_scale
        self.peak_threshold = peak_threshold

    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run centroid inference and return fixed-size outputs."""
        image = self._normalize_uint8(image)
        if self.input_scale != 1.0:
            height = int(image.shape[-2] * self.input_scale)
            width = int(image.shape[-1] * self.input_scale)
            image = F.interpolate(
                image, size=(height, width), mode="bilinear", align_corners=False
            )

        confmaps = self._extract_tensor(self.model(image), ["centroid", "confmap"])
        peaks, values, valid = self._find_topk_peaks(
            confmaps, self.max_instances, self.peak_threshold
        )
        peaks = peaks * (self.output_stride / self.input_scale)

        return {
            "centroids": peaks,
            "centroid_vals": values,
            "instance_valid": valid,
        }

__init__(model, max_instances=20, output_stride=2, input_scale=1.0, peak_threshold=0.2)

Initialize centroid ONNX wrapper.

Parameters:

Name Type Description Default
model Module

Centroid detection model.

 required
max_instances int

Maximum number of instances to detect.

 20
output_stride int

Output stride for confidence maps.

 2
input_scale float

Input scaling factor.

 1.0
peak_threshold float

Minimum confidence for a peak to be considered valid.

 0.2
Source code in sleap_nn/export/wrappers/centroid.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def __init__(
    self,
    model: nn.Module,
    max_instances: int = 20,
    output_stride: int = 2,
    input_scale: float = 1.0,
    peak_threshold: float = 0.2,
):
    """Initialize centroid ONNX wrapper.

    Args:
        model: Centroid detection model.
        max_instances: Maximum number of instances to detect.
        output_stride: Output stride for confidence maps.
        input_scale: Input scaling factor.
        peak_threshold: Minimum confidence for a peak to be considered valid.
    """
    super().__init__(model)
    self.max_instances = max_instances
    self.output_stride = output_stride
    self.input_scale = input_scale
    self.peak_threshold = peak_threshold

forward(image)

Run centroid inference and return fixed-size outputs.

Source code in sleap_nn/export/wrappers/centroid.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run centroid inference and return fixed-size outputs."""
    image = self._normalize_uint8(image)
    if self.input_scale != 1.0:
        height = int(image.shape[-2] * self.input_scale)
        width = int(image.shape[-1] * self.input_scale)
        image = F.interpolate(
            image, size=(height, width), mode="bilinear", align_corners=False
        )

    confmaps = self._extract_tensor(self.model(image), ["centroid", "confmap"])
    peaks, values, valid = self._find_topk_peaks(
        confmaps, self.max_instances, self.peak_threshold
    )
    peaks = peaks * (self.output_stride / self.input_scale)

    return {
        "centroids": peaks,
        "centroid_vals": values,
        "instance_valid": valid,
    }

EmbeddingONNXWrapper

Bases: BaseExportWrapper

Wrap an embedding model for ONNX export: crop -> appearance vector.

The simplest wrapper (single input/output, no peak finding): WHOLE-crop per-crop standardize the input, run the encoder + head, optionally L2-normalize. Output: {"embedding": (B, D)}.

Parity: this exactly reproduces native inference for a burn_in=False embedder (whose _standardize also normalizes over the whole crop). A burn_in=True embedder standardizes over the FOREGROUND (mask) only and fills the background, which this single-input graph cannot replicate — its exported embeddings therefore DIVERGE from native masked inference. The export CLI records burn_in/background_fill in the metadata and warns on a burn_in=True export; use the native sleap-nn predict ... --save_embeddings slp path for exact parity with such a model.

Methods:

Name Description
__init__

Initialize.
forward

image: (B, C, H, W) [0, 255] -> {"embedding": (B, D)}.
Source code in sleap_nn/export/wrappers/embedding.py 
11
12
13
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
class EmbeddingONNXWrapper(BaseExportWrapper):
    """Wrap an embedding model for ONNX export: crop -> appearance vector.

    The simplest wrapper (single input/output, no peak finding): WHOLE-crop per-crop
    standardize the input, run the encoder + head, optionally L2-normalize. Output:
    ``{"embedding": (B, D)}``.

    Parity: this exactly reproduces native inference for a ``burn_in=False`` embedder
    (whose ``_standardize`` also normalizes over the whole crop). A ``burn_in=True``
    embedder standardizes over the FOREGROUND (mask) only and fills the background, which
    this single-input graph cannot replicate — its exported embeddings therefore DIVERGE
    from native masked inference. The export CLI records ``burn_in``/``background_fill``
    in the metadata and warns on a ``burn_in=True`` export; use the native
    ``sleap-nn predict ... --save_embeddings slp`` path for exact parity with such
    a model.
    """

    def __init__(self, model, normalize: bool = True, eps: float = 1e-5):
        """Initialize.

        Args:
            model: The underlying ``Model`` (encoder + EmbeddingHead).
            normalize: L2-normalize the output embedding.
            eps: Standardization epsilon.
        """
        super().__init__(model)
        self.normalize = normalize
        self.eps = eps

    def forward(self, image: torch.Tensor):
        """image: (B, C, H, W) [0, 255] -> {"embedding": (B, D)}.

        Replicates ``EmbeddingLightningModule._standardize``'s MASKLESS path exactly
        (i.e. the ``burn_in=False`` native path): reduce over the spatial dims only so
        each channel is standardized independently. For grayscale (C=1) this is the plain
        per-crop standardize; for RGB (C=3) it is a true per-channel zero-mean/unit-std,
        matching the PyTorch inference path. Computing the count from a ones tensor
        (rather than a baked H*W constant) keeps the graph valid under dynamic spatial
        axes. NOTE: a ``burn_in=True`` model's masked (foreground-only) standardize is NOT
        reproduced here — see the class docstring.
        """
        x = image.float()
        ones = torch.ones_like(x[:, :1])
        cnt = ones.sum((2, 3), keepdim=True).clamp(min=1)
        mu = (x * ones).sum((2, 3), keepdim=True) / cnt
        var = ((x - mu) ** 2 * ones).sum((2, 3), keepdim=True) / cnt
        x = (x - mu) / (var.sqrt() + self.eps)
        feat = self._extract_tensor(self.model(x), ["embedding", "vector"])
        if feat.dim() > 2:
            feat = feat.flatten(1)
        if self.normalize:
            feat = F.normalize(feat, p=2, dim=-1)
        return {"embedding": feat}

__init__(model, normalize=True, eps=1e-05)

Initialize.

Parameters:

Name Type Description Default
model

The underlying Model (encoder + EmbeddingHead).

 required
normalize bool

L2-normalize the output embedding.

 True
eps float

Standardization epsilon.

 1e-05
Source code in sleap_nn/export/wrappers/embedding.py 
28
29
30
31
32
33
34
35
36
37
38
def __init__(self, model, normalize: bool = True, eps: float = 1e-5):
    """Initialize.

    Args:
        model: The underlying ``Model`` (encoder + EmbeddingHead).
        normalize: L2-normalize the output embedding.
        eps: Standardization epsilon.
    """
    super().__init__(model)
    self.normalize = normalize
    self.eps = eps

forward(image)

image: (B, C, H, W) [0, 255] -> {"embedding": (B, D)}.

Replicates EmbeddingLightningModule._standardize's MASKLESS path exactly (i.e. the burn_in=False native path): reduce over the spatial dims only so each channel is standardized independently. For grayscale (C=1) this is the plain per-crop standardize; for RGB (C=3) it is a true per-channel zero-mean/unit-std, matching the PyTorch inference path. Computing the count from a ones tensor (rather than a baked H*W constant) keeps the graph valid under dynamic spatial axes. NOTE: a burn_in=True model's masked (foreground-only) standardize is NOT reproduced here — see the class docstring.

Source code in sleap_nn/export/wrappers/embedding.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def forward(self, image: torch.Tensor):
    """image: (B, C, H, W) [0, 255] -> {"embedding": (B, D)}.

    Replicates ``EmbeddingLightningModule._standardize``'s MASKLESS path exactly
    (i.e. the ``burn_in=False`` native path): reduce over the spatial dims only so
    each channel is standardized independently. For grayscale (C=1) this is the plain
    per-crop standardize; for RGB (C=3) it is a true per-channel zero-mean/unit-std,
    matching the PyTorch inference path. Computing the count from a ones tensor
    (rather than a baked H*W constant) keeps the graph valid under dynamic spatial
    axes. NOTE: a ``burn_in=True`` model's masked (foreground-only) standardize is NOT
    reproduced here — see the class docstring.
    """
    x = image.float()
    ones = torch.ones_like(x[:, :1])
    cnt = ones.sum((2, 3), keepdim=True).clamp(min=1)
    mu = (x * ones).sum((2, 3), keepdim=True) / cnt
    var = ((x - mu) ** 2 * ones).sum((2, 3), keepdim=True) / cnt
    x = (x - mu) / (var.sqrt() + self.eps)
    feat = self._extract_tensor(self.model(x), ["embedding", "vector"])
    if feat.dim() > 2:
        feat = feat.flatten(1)
    if self.normalize:
        feat = F.normalize(feat, p=2, dim=-1)
    return {"embedding": feat}

SingleInstanceONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for single-instance models.

This wrapper handles full-frame inference assuming a single instance per frame. For each body part (channel), it finds the global maximum in the confidence map.

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

Attributes:

Name Type Description
model

The trained backbone model that outputs confidence maps.
output_stride

Output stride of the model (e.g., 4 means confmaps are ¼ the input resolution).
input_scale

Factor to scale input images before inference.

Methods:

Name Description
__init__

Initialize the single-instance wrapper.
forward

Run single-instance inference.
Source code in sleap_nn/export/wrappers/single_instance.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
class SingleInstanceONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for single-instance models.

    This wrapper handles full-frame inference assuming a single instance per frame.
    For each body part (channel), it finds the global maximum in the confidence map.

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

    Attributes:
        model: The trained backbone model that outputs confidence maps.
        output_stride: Output stride of the model (e.g., 4 means confmaps are 1/4 the
            input resolution).
        input_scale: Factor to scale input images before inference.
    """

    def __init__(
        self,
        model: nn.Module,
        output_stride: int = 4,
        input_scale: float = 1.0,
        peak_threshold: float = 0.2,
    ):
        """Initialize the single-instance wrapper.

        Args:
            model: The trained backbone model.
            output_stride: Output stride of the model. Default: 4.
            input_scale: Factor to scale input images. Default: 1.0.
            peak_threshold: Minimum confidence for a peak to be considered valid.
        """
        super().__init__(model)
        self.output_stride = output_stride
        self.input_scale = input_scale
        self.peak_threshold = peak_threshold

    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run single-instance inference.

        Args:
            image: Input image tensor of shape (batch, channels, height, width).
                Expected as uint8 [0, 255] values.

        Returns:
            Dictionary with:
                peaks: Peak coordinates of shape (batch, n_nodes, 2) in (x, y) format.
                peak_vals: Peak confidence values of shape (batch, n_nodes).
        """
        # Normalize uint8 [0, 255] to float32 [0, 1]
        image = self._normalize_uint8(image)

        # Apply input scaling if needed. ``antialias=True`` matches the PyTorch inference
        # path, which resizes via ``torchvision.transforms.functional.resize``
        # (antialiased by default). Without it a downscaling resize (input_scale < 1)
        # produces visibly different pixels, which shifts confidence-map peaks and —
        # once scaled back by ``output_stride / input_scale`` — the exported keypoints.
        if self.input_scale != 1.0:
            height = int(image.shape[-2] * self.input_scale)
            width = int(image.shape[-1] * self.input_scale)
            image = F.interpolate(
                image,
                size=(height, width),
                mode="bilinear",
                align_corners=False,
                antialias=True,
            )

        # Run model to get confidence maps: (batch, n_nodes, height, width)
        confmaps = self._extract_tensor(
            self.model(image), ["single", "instance", "confmap"]
        )

        # Find global peak for each channel: (batch, n_nodes, 2), (batch, n_nodes)
        peaks, values = self._find_global_peaks(confmaps, self.peak_threshold)

        # Scale peaks from confmap coordinates to image coordinates
        peaks = peaks * (self.output_stride / self.input_scale)

        return {
            "peaks": peaks,
            "peak_vals": values,
        }

__init__(model, output_stride=4, input_scale=1.0, peak_threshold=0.2)

Initialize the single-instance wrapper.

Parameters:

Name Type Description Default
model Module

The trained backbone model.

 required
output_stride int

Output stride of the model. Default: 4.

 4
input_scale float

Factor to scale input images. Default: 1.0.

 1.0
peak_threshold float

Minimum confidence for a peak to be considered valid.

 0.2
Source code in sleap_nn/export/wrappers/single_instance.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    model: nn.Module,
    output_stride: int = 4,
    input_scale: float = 1.0,
    peak_threshold: float = 0.2,
):
    """Initialize the single-instance wrapper.

    Args:
        model: The trained backbone model.
        output_stride: Output stride of the model. Default: 4.
        input_scale: Factor to scale input images. Default: 1.0.
        peak_threshold: Minimum confidence for a peak to be considered valid.
    """
    super().__init__(model)
    self.output_stride = output_stride
    self.input_scale = input_scale
    self.peak_threshold = peak_threshold

forward(image)

Run single-instance inference.

Parameters:

Name Type Description Default
image Tensor

Input image tensor of shape (batch, channels, height, width). Expected as uint8 [0, 255] values.

 required

Returns:

Type Description
Dict[str, Tensor]

Dictionary with: peaks: Peak coordinates of shape (batch, n_nodes, 2) in (x, y) format. peak_vals: Peak confidence values of shape (batch, n_nodes).
Source code in sleap_nn/export/wrappers/single_instance.py 
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
def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run single-instance inference.

    Args:
        image: Input image tensor of shape (batch, channels, height, width).
            Expected as uint8 [0, 255] values.

    Returns:
        Dictionary with:
            peaks: Peak coordinates of shape (batch, n_nodes, 2) in (x, y) format.
            peak_vals: Peak confidence values of shape (batch, n_nodes).
    """
    # Normalize uint8 [0, 255] to float32 [0, 1]
    image = self._normalize_uint8(image)

    # Apply input scaling if needed. ``antialias=True`` matches the PyTorch inference
    # path, which resizes via ``torchvision.transforms.functional.resize``
    # (antialiased by default). Without it a downscaling resize (input_scale < 1)
    # produces visibly different pixels, which shifts confidence-map peaks and —
    # once scaled back by ``output_stride / input_scale`` — the exported keypoints.
    if self.input_scale != 1.0:
        height = int(image.shape[-2] * self.input_scale)
        width = int(image.shape[-1] * self.input_scale)
        image = F.interpolate(
            image,
            size=(height, width),
            mode="bilinear",
            align_corners=False,
            antialias=True,
        )

    # Run model to get confidence maps: (batch, n_nodes, height, width)
    confmaps = self._extract_tensor(
        self.model(image), ["single", "instance", "confmap"]
    )

    # Find global peak for each channel: (batch, n_nodes, 2), (batch, n_nodes)
    peaks, values = self._find_global_peaks(confmaps, self.peak_threshold)

    # Scale peaks from confmap coordinates to image coordinates
    peaks = peaks * (self.output_stride / self.input_scale)

    return {
        "peaks": peaks,
        "peak_vals": values,
    }

TopDownMultiClassCombinedONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for combined centroid + multiclass instance models.

This wrapper combines a centroid detection model with a centered instance multiclass model. It performs: 1. Centroid detection on full images 2. Cropping around each centroid using vectorized grid_sample 3. Instance keypoint detection + identity classification on each crop

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

Methods:

Name Description
__init__

Initialize the combined wrapper.
forward

Run combined top-down multiclass inference.
Source code in sleap_nn/export/wrappers/topdown_multiclass.py 
 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
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
class TopDownMultiClassCombinedONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for combined centroid + multiclass instance models.

    This wrapper combines a centroid detection model with a centered instance
    multiclass model. It performs:
    1. Centroid detection on full images
    2. Cropping around each centroid using vectorized grid_sample
    3. Instance keypoint detection + identity classification on each crop

    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 = (192, 192),
        centroid_output_stride: int = 4,
        instance_output_stride: int = 2,
        centroid_input_scale: float = 1.0,
        instance_input_scale: float = 1.0,
        n_nodes: int = 13,
        n_classes: int = 2,
        centroid_peak_threshold: float = 0.2,
        instance_peak_threshold: float = 0.2,
    ):
        """Initialize the combined wrapper.

        Args:
            centroid_model: Model for centroid detection.
            instance_model: Model for instance keypoints + class prediction.
            max_instances: Maximum number of instances to detect.
            crop_size: Size of crops around centroids (height, width).
            centroid_output_stride: Output stride of centroid model.
            instance_output_stride: Output stride of instance model.
            centroid_input_scale: Input scale for centroid model.
            instance_input_scale: Input scale for instance model.
            n_nodes: Number of keypoint nodes per instance.
            n_classes: Number of identity classes.
            centroid_peak_threshold: Minimum confidence for centroid peaks.
            instance_peak_threshold: Minimum confidence for instance peaks.
        """
        super().__init__(centroid_model)  # Primary model is centroid
        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.n_classes = n_classes
        self.centroid_peak_threshold = centroid_peak_threshold
        self.instance_peak_threshold = instance_peak_threshold

        # Pre-compute base grid for crop extraction (same as TopDownONNXWrapper)
        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 combined top-down multiclass inference.

        Args:
            image: Input image tensor of shape (batch, channels, height, width).
                   Expected to be uint8 in [0, 255].

        Returns:
            Dictionary with keys:
                - "centroids": Detected centroids (batch, max_instances, 2).
                - "centroid_vals": Centroid confidence values (batch, max_instances).
                - "peaks": Instance peaks (batch, max_instances, n_nodes, 2).
                - "peak_vals": Peak values (batch, max_instances, n_nodes).
                - "class_logits": Class logits per instance (batch, max_instances, n_classes).
                - "instance_valid": Validity mask (batch, max_instances).
        """
        # Normalize input
        image = self._normalize_uint8(image)
        batch_size, channels, height, width = image.shape

        # Apply centroid input scaling
        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 detection
        centroid_out = self.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
        )

        # Extract crops using vectorized grid_sample (same as TopDownONNXWrapper)
        crops = self._extract_crops(image, centroids)
        crops_flat = crops.reshape(
            batch_size * self.max_instances,
            channels,
            self.crop_size[0],
            self.crop_size[1],
        )

        # Apply instance input scaling if needed
        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 model forward (batch all crops)
        instance_out = self.instance_model(crops_flat)
        instance_cms = self._extract_tensor(
            instance_out, ["centered", "instance", "confmap"]
        )
        instance_class = self._extract_tensor(instance_out, ["class", "vector"])

        # Find peaks in all crops
        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
        )

        # Reshape to batch x instances x nodes x 2
        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)

        # Reshape class logits
        class_logits = instance_class.reshape(
            batch_size, self.max_instances, self.n_classes
        )

        # Transform peaks from crop coordinates to full image coordinates
        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

        # Zero out invalid instances
        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)
        class_logits = class_logits.masked_fill(invalid_mask.unsqueeze(-1), 0.0)

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

    def _extract_crops(
        self,
        image: torch.Tensor,
        centroids: torch.Tensor,
    ) -> torch.Tensor:
        """Extract crops around centroids using grid_sample.

        This is the same vectorized implementation as TopDownONNXWrapper.
        """
        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=4, instance_output_stride=2, centroid_input_scale=1.0, instance_input_scale=1.0, n_nodes=13, n_classes=2, centroid_peak_threshold=0.2, instance_peak_threshold=0.2)

Initialize the combined wrapper.

Parameters:

Name Type Description Default
centroid_model Module

Model for centroid detection.

 required
instance_model Module

Model for instance keypoints + class prediction.

 required
max_instances int

Maximum number of instances to detect.

 20
crop_size tuple

Size of crops around centroids (height, width).

 (192, 192)
centroid_output_stride int

Output stride of centroid model.

 4
instance_output_stride int

Output stride of instance model.

 2
centroid_input_scale float

Input scale for centroid model.

 1.0
instance_input_scale float

Input scale for instance model.

 1.0
n_nodes int

Number of keypoint nodes per instance.

 13
n_classes int

Number of identity classes.

 2
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_multiclass.py 
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
def __init__(
    self,
    centroid_model: nn.Module,
    instance_model: nn.Module,
    max_instances: int = 20,
    crop_size: tuple = (192, 192),
    centroid_output_stride: int = 4,
    instance_output_stride: int = 2,
    centroid_input_scale: float = 1.0,
    instance_input_scale: float = 1.0,
    n_nodes: int = 13,
    n_classes: int = 2,
    centroid_peak_threshold: float = 0.2,
    instance_peak_threshold: float = 0.2,
):
    """Initialize the combined wrapper.

    Args:
        centroid_model: Model for centroid detection.
        instance_model: Model for instance keypoints + class prediction.
        max_instances: Maximum number of instances to detect.
        crop_size: Size of crops around centroids (height, width).
        centroid_output_stride: Output stride of centroid model.
        instance_output_stride: Output stride of instance model.
        centroid_input_scale: Input scale for centroid model.
        instance_input_scale: Input scale for instance model.
        n_nodes: Number of keypoint nodes per instance.
        n_classes: Number of identity classes.
        centroid_peak_threshold: Minimum confidence for centroid peaks.
        instance_peak_threshold: Minimum confidence for instance peaks.
    """
    super().__init__(centroid_model)  # Primary model is centroid
    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.n_classes = n_classes
    self.centroid_peak_threshold = centroid_peak_threshold
    self.instance_peak_threshold = instance_peak_threshold

    # Pre-compute base grid for crop extraction (same as TopDownONNXWrapper)
    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 combined top-down multiclass inference.

Parameters:

Name Type Description Default
image Tensor

Input image tensor of shape (batch, channels, height, width). Expected to be uint8 in [0, 255].

 required

Returns:

Type Description
Dict[str, Tensor]

Dictionary with keys: - "centroids": Detected centroids (batch, max_instances, 2). - "centroid_vals": Centroid confidence values (batch, max_instances). - "peaks": Instance peaks (batch, max_instances, n_nodes, 2). - "peak_vals": Peak values (batch, max_instances, n_nodes). - "class_logits": Class logits per instance (batch, max_instances, n_classes). - "instance_valid": Validity mask (batch, max_instances).
Source code in sleap_nn/export/wrappers/topdown_multiclass.py 
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
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run combined top-down multiclass inference.

    Args:
        image: Input image tensor of shape (batch, channels, height, width).
               Expected to be uint8 in [0, 255].

    Returns:
        Dictionary with keys:
            - "centroids": Detected centroids (batch, max_instances, 2).
            - "centroid_vals": Centroid confidence values (batch, max_instances).
            - "peaks": Instance peaks (batch, max_instances, n_nodes, 2).
            - "peak_vals": Peak values (batch, max_instances, n_nodes).
            - "class_logits": Class logits per instance (batch, max_instances, n_classes).
            - "instance_valid": Validity mask (batch, max_instances).
    """
    # Normalize input
    image = self._normalize_uint8(image)
    batch_size, channels, height, width = image.shape

    # Apply centroid input scaling
    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 detection
    centroid_out = self.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
    )

    # Extract crops using vectorized grid_sample (same as TopDownONNXWrapper)
    crops = self._extract_crops(image, centroids)
    crops_flat = crops.reshape(
        batch_size * self.max_instances,
        channels,
        self.crop_size[0],
        self.crop_size[1],
    )

    # Apply instance input scaling if needed
    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 model forward (batch all crops)
    instance_out = self.instance_model(crops_flat)
    instance_cms = self._extract_tensor(
        instance_out, ["centered", "instance", "confmap"]
    )
    instance_class = self._extract_tensor(instance_out, ["class", "vector"])

    # Find peaks in all crops
    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
    )

    # Reshape to batch x instances x nodes x 2
    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)

    # Reshape class logits
    class_logits = instance_class.reshape(
        batch_size, self.max_instances, self.n_classes
    )

    # Transform peaks from crop coordinates to full image coordinates
    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

    # Zero out invalid instances
    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)
    class_logits = class_logits.masked_fill(invalid_mask.unsqueeze(-1), 0.0)

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

TopDownMultiClassONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for top-down multiclass (supervised ID) models.

This wrapper handles models that output both confidence maps for keypoint detection and class logits for identity classification. It runs on instance crops (centered around detected centroids).

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

Attributes:

Name Type Description
model

The underlying PyTorch model (centered instance + class vectors heads).
output_stride

Output stride of the confmap head.
input_scale

Scale factor applied to input images before inference.
n_classes

Number of identity classes.

Methods:

Name Description
__init__

Initialize the wrapper.
forward

Run top-down multiclass inference on crops.
Source code in sleap_nn/export/wrappers/topdown_multiclass.py 
12
13
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
class TopDownMultiClassONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for top-down multiclass (supervised ID) models.

    This wrapper handles models that output both confidence maps for keypoint
    detection and class logits for identity classification. It runs on instance
    crops (centered around detected centroids).

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

    Attributes:
        model: The underlying PyTorch model (centered instance + class vectors heads).
        output_stride: Output stride of the confmap head.
        input_scale: Scale factor applied to input images before inference.
        n_classes: Number of identity classes.
    """

    def __init__(
        self,
        model: nn.Module,
        output_stride: int = 2,
        input_scale: float = 1.0,
        n_classes: int = 2,
        peak_threshold: float = 0.2,
    ):
        """Initialize the wrapper.

        Args:
            model: The underlying PyTorch model.
            output_stride: Output stride of the confidence maps.
            input_scale: Scale factor for input images.
            n_classes: Number of identity classes (e.g., 2 for male/female).
            peak_threshold: Minimum confidence for a peak to be considered valid.
        """
        super().__init__(model)
        self.output_stride = output_stride
        self.input_scale = input_scale
        self.n_classes = n_classes
        self.peak_threshold = peak_threshold

    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run top-down multiclass inference on crops.

        Args:
            image: Input image tensor of shape (batch, channels, height, width).
                   Expected to be uint8 in [0, 255].

        Returns:
            Dictionary with keys:
                - "peaks": Predicted peak coordinates (batch, n_nodes, 2) in (x, y).
                - "peak_vals": Peak confidence values (batch, n_nodes).
                - "class_logits": Raw class logits (batch, n_classes).

            The class assignment is done on CPU using Hungarian matching
            via `get_class_inds_from_vectors()`.
        """
        # Normalize uint8 [0, 255] to float32 [0, 1]
        image = self._normalize_uint8(image)

        # Apply input scaling if needed
        if self.input_scale != 1.0:
            height = int(image.shape[-2] * self.input_scale)
            width = int(image.shape[-1] * self.input_scale)
            image = F.interpolate(
                image, size=(height, width), mode="bilinear", align_corners=False
            )

        # Forward pass
        out = self.model(image)

        # Extract outputs
        confmaps = self._extract_tensor(out, ["centered", "instance", "confmap"])
        class_logits = self._extract_tensor(out, ["class", "vector"])

        # Find global peaks (one per node)
        peaks, peak_vals = self._find_global_peaks(confmaps, self.peak_threshold)

        # Scale peaks back to input coordinates
        peaks = peaks * (self.output_stride / self.input_scale)

        return {
            "peaks": peaks,
            "peak_vals": peak_vals,
            "class_logits": class_logits,
        }

__init__(model, output_stride=2, input_scale=1.0, n_classes=2, peak_threshold=0.2)

Initialize the wrapper.

Parameters:

Name Type Description Default
model Module

The underlying PyTorch model.

 required
output_stride int

Output stride of the confidence maps.

 2
input_scale float

Scale factor for input images.

 1.0
n_classes int

Number of identity classes (e.g., 2 for male/female).

 2
peak_threshold float

Minimum confidence for a peak to be considered valid.

 0.2
Source code in sleap_nn/export/wrappers/topdown_multiclass.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def __init__(
    self,
    model: nn.Module,
    output_stride: int = 2,
    input_scale: float = 1.0,
    n_classes: int = 2,
    peak_threshold: float = 0.2,
):
    """Initialize the wrapper.

    Args:
        model: The underlying PyTorch model.
        output_stride: Output stride of the confidence maps.
        input_scale: Scale factor for input images.
        n_classes: Number of identity classes (e.g., 2 for male/female).
        peak_threshold: Minimum confidence for a peak to be considered valid.
    """
    super().__init__(model)
    self.output_stride = output_stride
    self.input_scale = input_scale
    self.n_classes = n_classes
    self.peak_threshold = peak_threshold

forward(image)

Run top-down multiclass inference on crops.

Parameters:

Name Type Description Default
image Tensor

Input image tensor of shape (batch, channels, height, width). Expected to be uint8 in [0, 255].

 required

Returns:

Type Description
Dict[str, Tensor]

Dictionary with keys: - "peaks": Predicted peak coordinates (batch, n_nodes, 2) in (x, y). - "peak_vals": Peak confidence values (batch, n_nodes). - "class_logits": Raw class logits (batch, n_classes).

The class assignment is done on CPU using Hungarian matching via get_class_inds_from_vectors().
Source code in sleap_nn/export/wrappers/topdown_multiclass.py 
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
def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run top-down multiclass inference on crops.

    Args:
        image: Input image tensor of shape (batch, channels, height, width).
               Expected to be uint8 in [0, 255].

    Returns:
        Dictionary with keys:
            - "peaks": Predicted peak coordinates (batch, n_nodes, 2) in (x, y).
            - "peak_vals": Peak confidence values (batch, n_nodes).
            - "class_logits": Raw class logits (batch, n_classes).

        The class assignment is done on CPU using Hungarian matching
        via `get_class_inds_from_vectors()`.
    """
    # Normalize uint8 [0, 255] to float32 [0, 1]
    image = self._normalize_uint8(image)

    # Apply input scaling if needed
    if self.input_scale != 1.0:
        height = int(image.shape[-2] * self.input_scale)
        width = int(image.shape[-1] * self.input_scale)
        image = F.interpolate(
            image, size=(height, width), mode="bilinear", align_corners=False
        )

    # Forward pass
    out = self.model(image)

    # Extract outputs
    confmaps = self._extract_tensor(out, ["centered", "instance", "confmap"])
    class_logits = self._extract_tensor(out, ["class", "vector"])

    # Find global peaks (one per node)
    peaks, peak_vals = self._find_global_peaks(confmaps, self.peak_threshold)

    # Scale peaks back to input coordinates
    peaks = peaks * (self.output_stride / self.input_scale)

    return {
        "peaks": peaks,
        "peak_vals": peak_vals,
        "class_logits": class_logits,
    }

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,
    }