Skip to content

bottomup_multiclass

sleap_nn.export.wrappers.bottomup_multiclass

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

Classes:

Name Description
BottomUpMultiClassONNXWrapper

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

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