Skip to content

bottomup

sleap_nn.inference.bottomup

Inference modules for BottomUp models.

Classes:

Name Description
BottomUpInferenceModel

BottomUp Inference model.
BottomUpMultiClassInferenceModel

BottomUp Inference model for multi-class models.

BottomUpInferenceModel

Bases: LightningModule

BottomUp Inference model.

This model encapsulates the bottom-up approach. The images are passed to a peak detector to get the predicted instances and then fed into PAF to combine nodes belonging to the same instance.

Attributes:

Name Type Description
torch_model

A nn.Module that accepts rank-5 images as input and predicts rank-4 confidence maps as output. This should be a model that is trained on MultiInstanceConfMaps.
paf_scorer

A sleap_nn.inference.paf_grouping.PAFScorer instance configured to group instances based on peaks and PAFs produced by the model.
cms_output_stride

Output stride of the model, denoting the scale of the output confidence maps relative to the images (after input scaling). This is used for adjusting the peak coordinates to the image grid.
pafs_output_stride

Output stride of the model, denoting the scale of the output pafs relative to the images (after input scaling). This is used for adjusting the peak coordinates to the image grid.
peak_threshold

Minimum confidence map value to consider a global peak as valid.
refinement

If None, returns the grid-aligned peaks with no refinement. If "integral", peaks will be refined with integral regression. If "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.
integral_patch_size

Size of patches to crop around each rough peak for integral refinement as an integer scalar.
return_confmaps

If True, the confidence maps will be returned together with the predicted peaks. This will result in slower inference times since the data must be copied off of the GPU, but is useful for visualizing the raw output of the model.
return_pafs

If True, the part affinity fields will be returned together with the predicted instances. This will result in slower inference times since the data must be copied off of the GPU, but is useful for visualizing the raw output of the model.
return_paf_graph

If True, the part affinity field graph will be returned together with the predicted instances. The graph is obtained by parsing the part affinity fields with the paf_scorer instance and is an intermediate representation used during instance grouping.
input_scale

Float indicating if the images should be resized before being passed to the model.

Methods:

Name Description
__init__

Initialise the model attributes.
forward

Predict confidence maps and infer peak coordinates.
forward_gpu

Run the GPU portion: network forward pass + peak extraction.
postprocess_cpu

Run CPU post-processing: PAF scoring, matching, and instance grouping.
Source code in sleap_nn/inference/bottomup.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
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
class BottomUpInferenceModel(L.LightningModule):
    """BottomUp Inference model.

    This model encapsulates the bottom-up approach. The images are passed to a peak detector
    to get the predicted instances and then fed into PAF to combine nodes belonging to
    the same instance.

    Attributes:
        torch_model: A `nn.Module` that accepts rank-5 images as input and predicts
            rank-4 confidence maps as output. This should be a model that is trained on
            MultiInstanceConfMaps.
        paf_scorer: A `sleap_nn.inference.paf_grouping.PAFScorer` instance configured to group
            instances based on peaks and PAFs produced by the model.
        cms_output_stride: Output stride of the model, denoting the scale of the output
            confidence maps relative to the images (after input scaling). This is used
            for adjusting the peak coordinates to the image grid.
        pafs_output_stride: Output stride of the model, denoting the scale of the output
            pafs relative to the images (after input scaling). This is used
            for adjusting the peak coordinates to the image grid.
        peak_threshold: Minimum confidence map value to consider a global peak as valid.
        refinement: If `None`, returns the grid-aligned peaks with no refinement. If
            `"integral"`, peaks will be refined with integral regression. If `"local"`,
            peaks will be refined with quarter pixel local gradient offset. This has no
            effect if the model has an offset regression head.
        integral_patch_size: Size of patches to crop around each rough peak for integral
            refinement as an integer scalar.
        return_confmaps: If `True`, the confidence maps will be returned together with
            the predicted peaks. This will result in slower inference times since
            the data must be copied off of the GPU, but is useful for visualizing the
            raw output of the model.
        return_pafs: If `True`, the part affinity fields will be returned together with
            the predicted instances. This will result in slower inference times since
            the data must be copied off of the GPU, but is useful for visualizing the
            raw output of the model.
        return_paf_graph: If `True`, the part affinity field graph will be returned
            together with the predicted instances. The graph is obtained by parsing the
            part affinity fields with the `paf_scorer` instance and is an intermediate
            representation used during instance grouping.
        input_scale: Float indicating if the images should be resized before being
            passed to the model.
    """

    def __init__(
        self,
        torch_model: L.LightningModule,
        paf_scorer: PAFScorer,
        cms_output_stride: Optional[int] = None,
        pafs_output_stride: Optional[int] = None,
        peak_threshold: float = 0.0,
        refinement: Optional[str] = "integral",
        integral_patch_size: int = 5,
        return_confmaps: Optional[bool] = False,
        return_pafs: Optional[bool] = False,
        return_paf_graph: Optional[bool] = False,
        input_scale: float = 1.0,
        max_peaks_per_node: Optional[int] = None,
    ):
        """Initialise the model attributes.

        Args:
            torch_model: A `nn.Module` that accepts images and predicts confidence maps.
            paf_scorer: A `PAFScorer` instance for grouping instances.
            cms_output_stride: Output stride of confidence maps relative to images.
            pafs_output_stride: Output stride of PAFs relative to images.
            peak_threshold: Minimum confidence map value for valid peaks.
            refinement: Peak refinement method: None, "integral", or "local".
            integral_patch_size: Size of patches for integral refinement.
            return_confmaps: If True, return confidence maps in output.
            return_pafs: If True, return PAFs in output.
            return_paf_graph: If True, return intermediate PAF graph in output.
            input_scale: Scale factor applied to input images.
            max_peaks_per_node: Maximum number of peaks allowed per node before
                skipping PAF scoring. If any node has more peaks than this limit,
                empty predictions are returned. This prevents combinatorial explosion
                during early training when confidence maps are noisy. Set to None to
                disable this check (default). Recommended value: 100.
        """
        super().__init__()
        self.torch_model = torch_model
        self.paf_scorer = paf_scorer
        self.peak_threshold = peak_threshold
        self.refinement = refinement
        self.integral_patch_size = integral_patch_size
        self.cms_output_stride = cms_output_stride
        self.pafs_output_stride = pafs_output_stride
        self.return_confmaps = return_confmaps
        self.return_pafs = return_pafs
        self.return_paf_graph = return_paf_graph
        self.input_scale = input_scale
        self.max_peaks_per_node = max_peaks_per_node

    def _generate_cms_peaks(self, cms):
        # TODO: append nans to batch them -> tensor (vectorize the initial paf grouping steps)
        peaks, peak_vals, sample_inds, peak_channel_inds = find_local_peaks(
            cms.detach(),
            threshold=self.peak_threshold,
            refinement=self.refinement,
            integral_patch_size=self.integral_patch_size,
        )
        # Adjust for stride and scale.
        peaks = peaks * self.cms_output_stride  # (n_centroids, 2)

        cms_peaks, cms_peak_vals, cms_peak_channel_inds = [], [], []

        for b in range(self.batch_size):
            cms_peaks.append(peaks[sample_inds == b])
            cms_peak_vals.append(peak_vals[sample_inds == b].to(torch.float32))
            cms_peak_channel_inds.append(peak_channel_inds[sample_inds == b])

        # cms_peaks: [(#nodes, 2), ...]
        return cms_peaks, cms_peak_vals, cms_peak_channel_inds

    def forward_gpu(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """Run the GPU portion: network forward pass + peak extraction.

        This performs the torch model forward pass and confidence map peak finding,
        but stops before PAF scoring (which is CPU-bound). This split enables
        pipelined inference where GPU and CPU work can overlap.

        Args:
            inputs: Dictionary with "image" and "eff_scale" as keys.

        Returns:
            A dictionary containing the GPU outputs needed for CPU post-processing:
            - All keys from inputs (image, frame_idx, video_idx, etc.)
            - "cms": Confidence maps tensor.
            - "pafs": Part affinity fields tensor (batch, h, w, 2*edges).
            - "cms_peaks": List of peak coordinates per sample.
            - "cms_peak_vals": List of peak values per sample.
            - "cms_peak_channel_inds": List of peak channel indices per sample.
            - "skip_paf_scoring": Whether to skip PAF scoring.
            - "n_nodes": Number of nodes in the confidence maps.
        """
        # Network forward pass.
        self.batch_size = inputs["image"].shape[0]
        output = self.torch_model(inputs["image"])
        cms = output["MultiInstanceConfmapsHead"]
        pafs = output["PartAffinityFieldsHead"].permute(
            0, 2, 3, 1
        )  # (batch, h, w, 2*edges)
        cms_peaks, cms_peak_vals, cms_peak_channel_inds = self._generate_cms_peaks(cms)

        # Check if too many peaks per node (prevents combinatorial explosion)
        skip_paf_scoring = False
        if self.max_peaks_per_node is not None:
            n_nodes = cms.shape[1]
            for b in range(self.batch_size):
                for node_idx in range(n_nodes):
                    n_peaks = int((cms_peak_channel_inds[b] == node_idx).sum().item())
                    if n_peaks > self.max_peaks_per_node:
                        logger.warning(
                            f"Skipping PAF scoring: node {node_idx} has {n_peaks} peaks "
                            f"(max_peaks_per_node={self.max_peaks_per_node}). "
                            f"Model may need more training."
                        )
                        skip_paf_scoring = True
                        break
                if skip_paf_scoring:
                    break

        gpu_output = dict(inputs)
        gpu_output["cms"] = cms
        gpu_output["pafs"] = pafs
        gpu_output["cms_peaks"] = cms_peaks
        gpu_output["cms_peak_vals"] = cms_peak_vals
        gpu_output["cms_peak_channel_inds"] = cms_peak_channel_inds
        gpu_output["skip_paf_scoring"] = skip_paf_scoring
        gpu_output["n_nodes"] = cms.shape[1]
        return gpu_output

    def postprocess_cpu(
        self, gpu_output: Dict[str, torch.Tensor]
    ) -> List[Dict[str, torch.Tensor]]:
        """Run CPU post-processing: PAF scoring, matching, and instance grouping.

        Takes the output from forward_gpu() and performs PAF-based instance grouping,
        scale adjustments, and assembles the final output dictionary.

        Args:
            gpu_output: Dictionary returned by forward_gpu().

        Returns:
            A list containing one dictionary of outputs with keys:
            - "pred_instance_peaks": Predicted peaks per instance.
            - "pred_peak_values": Confidence values per peak.
            - "instance_scores": Instance grouping scores.
            Plus optional confmaps/pafs/paf_graph keys.
        """
        cms = gpu_output["cms"]
        pafs = gpu_output["pafs"]
        cms_peaks = gpu_output["cms_peaks"]
        cms_peak_vals = gpu_output["cms_peak_vals"]
        cms_peak_channel_inds = gpu_output["cms_peak_channel_inds"]
        skip_paf_scoring = gpu_output["skip_paf_scoring"]
        n_nodes = gpu_output["n_nodes"]
        batch_size = cms.shape[0]

        if skip_paf_scoring:
            # Return empty predictions for each sample
            device = cms.device
            predicted_instances_adjusted = []
            predicted_peak_scores = []
            predicted_instance_scores = []
            for _ in range(batch_size):
                predicted_instances_adjusted.append(
                    torch.full((0, n_nodes, 2), float("nan"), device=device)
                )
                predicted_peak_scores.append(
                    torch.full((0, n_nodes), float("nan"), device=device)
                )
                predicted_instance_scores.append(torch.tensor([], device=device))
            edge_inds = [
                torch.tensor([], dtype=torch.int32, device=device)
            ] * batch_size
            edge_peak_inds = [
                torch.tensor([], dtype=torch.int32, device=device).reshape(0, 2)
            ] * batch_size
            line_scores = [torch.tensor([], device=device)] * batch_size
        else:
            (
                predicted_instances,
                predicted_peak_scores,
                predicted_instance_scores,
                edge_inds,
                edge_peak_inds,
                line_scores,
            ) = self.paf_scorer.predict(
                pafs=pafs,
                peaks=cms_peaks,
                peak_vals=cms_peak_vals,
                peak_channel_inds=cms_peak_channel_inds,
            )

            predicted_instances = [p / self.input_scale for p in predicted_instances]
            predicted_instances_adjusted = []
            for idx, p in enumerate(predicted_instances):
                predicted_instances_adjusted.append(
                    p / gpu_output["eff_scale"][idx].to(p.device)
                )

        out = {
            "pred_instance_peaks": predicted_instances_adjusted,
            "pred_peak_values": predicted_peak_scores,
            "instance_scores": predicted_instance_scores,
        }

        if self.return_confmaps:
            out["pred_confmaps"] = cms.detach()
        if self.return_pafs:
            out["pred_part_affinity_fields"] = pafs.detach()
        if self.return_paf_graph:
            out["peaks"] = cms_peaks
            out["peak_vals"] = cms_peak_vals
            out["peak_channel_inds"] = cms_peak_channel_inds
            out["edge_inds"] = edge_inds
            out["edge_peak_inds"] = edge_peak_inds
            out["line_scores"] = line_scores

        # Merge with input metadata
        result = dict(gpu_output)
        # Remove intermediate GPU-stage keys
        for key in (
            "cms",
            "pafs",
            "cms_peaks",
            "cms_peak_vals",
            "cms_peak_channel_inds",
            "skip_paf_scoring",
            "n_nodes",
        ):
            result.pop(key, None)
        result.update(out)
        return [result]

    def forward(self, inputs: Dict[str, torch.Tensor]) -> List[Dict[str, torch.Tensor]]:
        """Predict confidence maps and infer peak coordinates.

        This composes forward_gpu() and postprocess_cpu() for the standard
        sequential inference path.

        Args:
            inputs: Dictionary with "image" as one of the keys.

        Returns:
            A list containing one dictionary of outputs with keys:

            `"pred_instance_peaks"`: The predicted peaks for each instance in the batch
                as a `torch.Tensor` of shape `(samples, nodes, 2)`.
            `"pred_peak_vals"`: The value of the confidence maps at the predicted
                peaks for each instance in the batch as a `torch.Tensor` of shape
                `(samples, nodes)`.
        """
        gpu_output = self.forward_gpu(inputs)
        return self.postprocess_cpu(gpu_output)

__init__(torch_model, paf_scorer, cms_output_stride=None, pafs_output_stride=None, peak_threshold=0.0, refinement='integral', integral_patch_size=5, return_confmaps=False, return_pafs=False, return_paf_graph=False, input_scale=1.0, max_peaks_per_node=None)

Initialise the model attributes.

Parameters:

Name Type Description Default
torch_model LightningModule

A nn.Module that accepts images and predicts confidence maps.

 required
paf_scorer PAFScorer

A PAFScorer instance for grouping instances.

 required
cms_output_stride Optional[int]

Output stride of confidence maps relative to images.

 None
pafs_output_stride Optional[int]

Output stride of PAFs relative to images.

 None
peak_threshold float

Minimum confidence map value for valid peaks.

 0.0
refinement Optional[str]

Peak refinement method: None, "integral", or "local".

 'integral'
integral_patch_size int

Size of patches for integral refinement.

 5
return_confmaps Optional[bool]

If True, return confidence maps in output.

 False
return_pafs Optional[bool]

If True, return PAFs in output.

 False
return_paf_graph Optional[bool]

If True, return intermediate PAF graph in output.

 False
input_scale float

Scale factor applied to input images.

 1.0
max_peaks_per_node Optional[int]

Maximum number of peaks allowed per node before skipping PAF scoring. If any node has more peaks than this limit, empty predictions are returned. This prevents combinatorial explosion during early training when confidence maps are noisy. Set to None to disable this check (default). Recommended value: 100.

 None
Source code in sleap_nn/inference/bottomup.py 
 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
def __init__(
    self,
    torch_model: L.LightningModule,
    paf_scorer: PAFScorer,
    cms_output_stride: Optional[int] = None,
    pafs_output_stride: Optional[int] = None,
    peak_threshold: float = 0.0,
    refinement: Optional[str] = "integral",
    integral_patch_size: int = 5,
    return_confmaps: Optional[bool] = False,
    return_pafs: Optional[bool] = False,
    return_paf_graph: Optional[bool] = False,
    input_scale: float = 1.0,
    max_peaks_per_node: Optional[int] = None,
):
    """Initialise the model attributes.

    Args:
        torch_model: A `nn.Module` that accepts images and predicts confidence maps.
        paf_scorer: A `PAFScorer` instance for grouping instances.
        cms_output_stride: Output stride of confidence maps relative to images.
        pafs_output_stride: Output stride of PAFs relative to images.
        peak_threshold: Minimum confidence map value for valid peaks.
        refinement: Peak refinement method: None, "integral", or "local".
        integral_patch_size: Size of patches for integral refinement.
        return_confmaps: If True, return confidence maps in output.
        return_pafs: If True, return PAFs in output.
        return_paf_graph: If True, return intermediate PAF graph in output.
        input_scale: Scale factor applied to input images.
        max_peaks_per_node: Maximum number of peaks allowed per node before
            skipping PAF scoring. If any node has more peaks than this limit,
            empty predictions are returned. This prevents combinatorial explosion
            during early training when confidence maps are noisy. Set to None to
            disable this check (default). Recommended value: 100.
    """
    super().__init__()
    self.torch_model = torch_model
    self.paf_scorer = paf_scorer
    self.peak_threshold = peak_threshold
    self.refinement = refinement
    self.integral_patch_size = integral_patch_size
    self.cms_output_stride = cms_output_stride
    self.pafs_output_stride = pafs_output_stride
    self.return_confmaps = return_confmaps
    self.return_pafs = return_pafs
    self.return_paf_graph = return_paf_graph
    self.input_scale = input_scale
    self.max_peaks_per_node = max_peaks_per_node

forward(inputs)

Predict confidence maps and infer peak coordinates.

This composes forward_gpu() and postprocess_cpu() for the standard sequential inference path.

Parameters:

Name Type Description Default
inputs Dict[str, Tensor]

Dictionary with "image" as one of the keys.

 required

Returns:

Type Description
List[Dict[str, Tensor]]

A list containing one dictionary of outputs with keys:

"pred_instance_peaks": The predicted peaks for each instance in the batch as a torch.Tensor of shape (samples, nodes, 2). "pred_peak_vals": The value of the confidence maps at the predicted peaks for each instance in the batch as a torch.Tensor of shape (samples, nodes).
Source code in sleap_nn/inference/bottomup.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
def forward(self, inputs: Dict[str, torch.Tensor]) -> List[Dict[str, torch.Tensor]]:
    """Predict confidence maps and infer peak coordinates.

    This composes forward_gpu() and postprocess_cpu() for the standard
    sequential inference path.

    Args:
        inputs: Dictionary with "image" as one of the keys.

    Returns:
        A list containing one dictionary of outputs with keys:

        `"pred_instance_peaks"`: The predicted peaks for each instance in the batch
            as a `torch.Tensor` of shape `(samples, nodes, 2)`.
        `"pred_peak_vals"`: The value of the confidence maps at the predicted
            peaks for each instance in the batch as a `torch.Tensor` of shape
            `(samples, nodes)`.
    """
    gpu_output = self.forward_gpu(inputs)
    return self.postprocess_cpu(gpu_output)

forward_gpu(inputs)

Run the GPU portion: network forward pass + peak extraction.

This performs the torch model forward pass and confidence map peak finding, but stops before PAF scoring (which is CPU-bound). This split enables pipelined inference where GPU and CPU work can overlap.

Parameters:

Name Type Description Default
inputs Dict[str, Tensor]

Dictionary with "image" and "eff_scale" as keys.

 required

Returns:

Type Description
Dict[str, Tensor]

A dictionary containing the GPU outputs needed for CPU post-processing: - All keys from inputs (image, frame_idx, video_idx, etc.) - "cms": Confidence maps tensor. - "pafs": Part affinity fields tensor (batch, h, w, 2*edges). - "cms_peaks": List of peak coordinates per sample. - "cms_peak_vals": List of peak values per sample. - "cms_peak_channel_inds": List of peak channel indices per sample. - "skip_paf_scoring": Whether to skip PAF scoring. - "n_nodes": Number of nodes in the confidence maps.
Source code in sleap_nn/inference/bottomup.py 
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
def forward_gpu(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """Run the GPU portion: network forward pass + peak extraction.

    This performs the torch model forward pass and confidence map peak finding,
    but stops before PAF scoring (which is CPU-bound). This split enables
    pipelined inference where GPU and CPU work can overlap.

    Args:
        inputs: Dictionary with "image" and "eff_scale" as keys.

    Returns:
        A dictionary containing the GPU outputs needed for CPU post-processing:
        - All keys from inputs (image, frame_idx, video_idx, etc.)
        - "cms": Confidence maps tensor.
        - "pafs": Part affinity fields tensor (batch, h, w, 2*edges).
        - "cms_peaks": List of peak coordinates per sample.
        - "cms_peak_vals": List of peak values per sample.
        - "cms_peak_channel_inds": List of peak channel indices per sample.
        - "skip_paf_scoring": Whether to skip PAF scoring.
        - "n_nodes": Number of nodes in the confidence maps.
    """
    # Network forward pass.
    self.batch_size = inputs["image"].shape[0]
    output = self.torch_model(inputs["image"])
    cms = output["MultiInstanceConfmapsHead"]
    pafs = output["PartAffinityFieldsHead"].permute(
        0, 2, 3, 1
    )  # (batch, h, w, 2*edges)
    cms_peaks, cms_peak_vals, cms_peak_channel_inds = self._generate_cms_peaks(cms)

    # Check if too many peaks per node (prevents combinatorial explosion)
    skip_paf_scoring = False
    if self.max_peaks_per_node is not None:
        n_nodes = cms.shape[1]
        for b in range(self.batch_size):
            for node_idx in range(n_nodes):
                n_peaks = int((cms_peak_channel_inds[b] == node_idx).sum().item())
                if n_peaks > self.max_peaks_per_node:
                    logger.warning(
                        f"Skipping PAF scoring: node {node_idx} has {n_peaks} peaks "
                        f"(max_peaks_per_node={self.max_peaks_per_node}). "
                        f"Model may need more training."
                    )
                    skip_paf_scoring = True
                    break
            if skip_paf_scoring:
                break

    gpu_output = dict(inputs)
    gpu_output["cms"] = cms
    gpu_output["pafs"] = pafs
    gpu_output["cms_peaks"] = cms_peaks
    gpu_output["cms_peak_vals"] = cms_peak_vals
    gpu_output["cms_peak_channel_inds"] = cms_peak_channel_inds
    gpu_output["skip_paf_scoring"] = skip_paf_scoring
    gpu_output["n_nodes"] = cms.shape[1]
    return gpu_output

postprocess_cpu(gpu_output)

Run CPU post-processing: PAF scoring, matching, and instance grouping.

Takes the output from forward_gpu() and performs PAF-based instance grouping, scale adjustments, and assembles the final output dictionary.

Parameters:

Name Type Description Default
gpu_output Dict[str, Tensor]

Dictionary returned by forward_gpu().

 required

Returns:

Type Description
List[Dict[str, Tensor]]

A list containing one dictionary of outputs with keys: - "pred_instance_peaks": Predicted peaks per instance. - "pred_peak_values": Confidence values per peak. - "instance_scores": Instance grouping scores. Plus optional confmaps/pafs/paf_graph keys.
Source code in sleap_nn/inference/bottomup.py 
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
def postprocess_cpu(
    self, gpu_output: Dict[str, torch.Tensor]
) -> List[Dict[str, torch.Tensor]]:
    """Run CPU post-processing: PAF scoring, matching, and instance grouping.

    Takes the output from forward_gpu() and performs PAF-based instance grouping,
    scale adjustments, and assembles the final output dictionary.

    Args:
        gpu_output: Dictionary returned by forward_gpu().

    Returns:
        A list containing one dictionary of outputs with keys:
        - "pred_instance_peaks": Predicted peaks per instance.
        - "pred_peak_values": Confidence values per peak.
        - "instance_scores": Instance grouping scores.
        Plus optional confmaps/pafs/paf_graph keys.
    """
    cms = gpu_output["cms"]
    pafs = gpu_output["pafs"]
    cms_peaks = gpu_output["cms_peaks"]
    cms_peak_vals = gpu_output["cms_peak_vals"]
    cms_peak_channel_inds = gpu_output["cms_peak_channel_inds"]
    skip_paf_scoring = gpu_output["skip_paf_scoring"]
    n_nodes = gpu_output["n_nodes"]
    batch_size = cms.shape[0]

    if skip_paf_scoring:
        # Return empty predictions for each sample
        device = cms.device
        predicted_instances_adjusted = []
        predicted_peak_scores = []
        predicted_instance_scores = []
        for _ in range(batch_size):
            predicted_instances_adjusted.append(
                torch.full((0, n_nodes, 2), float("nan"), device=device)
            )
            predicted_peak_scores.append(
                torch.full((0, n_nodes), float("nan"), device=device)
            )
            predicted_instance_scores.append(torch.tensor([], device=device))
        edge_inds = [
            torch.tensor([], dtype=torch.int32, device=device)
        ] * batch_size
        edge_peak_inds = [
            torch.tensor([], dtype=torch.int32, device=device).reshape(0, 2)
        ] * batch_size
        line_scores = [torch.tensor([], device=device)] * batch_size
    else:
        (
            predicted_instances,
            predicted_peak_scores,
            predicted_instance_scores,
            edge_inds,
            edge_peak_inds,
            line_scores,
        ) = self.paf_scorer.predict(
            pafs=pafs,
            peaks=cms_peaks,
            peak_vals=cms_peak_vals,
            peak_channel_inds=cms_peak_channel_inds,
        )

        predicted_instances = [p / self.input_scale for p in predicted_instances]
        predicted_instances_adjusted = []
        for idx, p in enumerate(predicted_instances):
            predicted_instances_adjusted.append(
                p / gpu_output["eff_scale"][idx].to(p.device)
            )

    out = {
        "pred_instance_peaks": predicted_instances_adjusted,
        "pred_peak_values": predicted_peak_scores,
        "instance_scores": predicted_instance_scores,
    }

    if self.return_confmaps:
        out["pred_confmaps"] = cms.detach()
    if self.return_pafs:
        out["pred_part_affinity_fields"] = pafs.detach()
    if self.return_paf_graph:
        out["peaks"] = cms_peaks
        out["peak_vals"] = cms_peak_vals
        out["peak_channel_inds"] = cms_peak_channel_inds
        out["edge_inds"] = edge_inds
        out["edge_peak_inds"] = edge_peak_inds
        out["line_scores"] = line_scores

    # Merge with input metadata
    result = dict(gpu_output)
    # Remove intermediate GPU-stage keys
    for key in (
        "cms",
        "pafs",
        "cms_peaks",
        "cms_peak_vals",
        "cms_peak_channel_inds",
        "skip_paf_scoring",
        "n_nodes",
    ):
        result.pop(key, None)
    result.update(out)
    return [result]

BottomUpMultiClassInferenceModel

Bases: LightningModule

BottomUp Inference model for multi-class models.

This model encapsulates the bottom-up approach. The images are passed to a local peak detector to get the predicted instances and then grouped into instances by their identity classifications.

Attributes:

Name Type Description
torch_model

A nn.Module that accepts rank-5 images as input and predicts rank-4 confidence maps as output. This should be a model that is trained on MultiInstanceConfMaps.
cms_output_stride

Output stride of the model, denoting the scale of the output confidence maps relative to the images (after input scaling). This is used for adjusting the peak coordinates to the image grid.
class_maps_output_stride

Output stride of the model, denoting the scale of the output pafs relative to the images (after input scaling). This is used for adjusting the peak coordinates to the image grid.
peak_threshold

Minimum confidence map value to consider a global peak as valid.
refinement

If None, returns the grid-aligned peaks with no refinement. If "integral", peaks will be refined with integral regression. If "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.
integral_patch_size

Size of patches to crop around each rough peak for integral refinement as an integer scalar.
return_confmaps

If True, the confidence maps will be returned together with the predicted peaks. This will result in slower inference times since the data must be copied off of the GPU, but is useful for visualizing the raw output of the model.
return_class_maps

If True, the class maps will be returned together with the predicted instances. This will result in slower inference times since the data must be copied off of the GPU, but is useful for visualizing the raw output of the model.
input_scale

Float indicating if the images should be resized before being passed to the model.

Methods:

Name Description
__init__

Initialise the model attributes.
forward

Predict confidence maps and infer peak coordinates.
Source code in sleap_nn/inference/bottomup.py 
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
class BottomUpMultiClassInferenceModel(L.LightningModule):
    """BottomUp Inference model for multi-class models.

    This model encapsulates the bottom-up approach. The images are passed to a local peak detector
    to get the predicted instances and then grouped into instances by their identity
    classifications.

    Attributes:
        torch_model: A `nn.Module` that accepts rank-5 images as input and predicts
            rank-4 confidence maps as output. This should be a model that is trained on
            MultiInstanceConfMaps.
        cms_output_stride: Output stride of the model, denoting the scale of the output
            confidence maps relative to the images (after input scaling). This is used
            for adjusting the peak coordinates to the image grid.
        class_maps_output_stride: Output stride of the model, denoting the scale of the output
            pafs relative to the images (after input scaling). This is used
            for adjusting the peak coordinates to the image grid.
        peak_threshold: Minimum confidence map value to consider a global peak as valid.
        refinement: If `None`, returns the grid-aligned peaks with no refinement. If
            `"integral"`, peaks will be refined with integral regression. If `"local"`,
            peaks will be refined with quarter pixel local gradient offset. This has no
            effect if the model has an offset regression head.
        integral_patch_size: Size of patches to crop around each rough peak for integral
            refinement as an integer scalar.
        return_confmaps: If `True`, the confidence maps will be returned together with
            the predicted peaks. This will result in slower inference times since
            the data must be copied off of the GPU, but is useful for visualizing the
            raw output of the model.
        return_class_maps: If `True`, the class maps will be returned together with
            the predicted instances. This will result in slower inference times since
            the data must be copied off of the GPU, but is useful for visualizing the
            raw output of the model.
        input_scale: Float indicating if the images should be resized before being
            passed to the model.
    """

    def __init__(
        self,
        torch_model: L.LightningModule,
        cms_output_stride: Optional[int] = None,
        class_maps_output_stride: Optional[int] = None,
        peak_threshold: float = 0.0,
        refinement: Optional[str] = "integral",
        integral_patch_size: int = 5,
        return_confmaps: Optional[bool] = False,
        return_class_maps: Optional[bool] = False,
        input_scale: float = 1.0,
    ):
        """Initialise the model attributes."""
        super().__init__()
        self.torch_model = torch_model
        self.peak_threshold = peak_threshold
        self.refinement = refinement
        self.integral_patch_size = integral_patch_size
        self.cms_output_stride = cms_output_stride
        self.class_maps_output_stride = class_maps_output_stride
        self.return_confmaps = return_confmaps
        self.return_class_maps = return_class_maps
        self.input_scale = input_scale

    def _generate_cms_peaks(self, cms):
        # TODO: append nans to batch them -> tensor (vectorize the initial paf grouping steps)
        peaks, peak_vals, sample_inds, peak_channel_inds = find_local_peaks(
            cms.detach(),
            threshold=self.peak_threshold,
            refinement=self.refinement,
            integral_patch_size=self.integral_patch_size,
        )
        # Adjust for stride and scale.
        peaks = peaks * self.cms_output_stride  # (n_centroids, 2)

        return peaks, peak_vals, sample_inds, peak_channel_inds

    def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """Predict confidence maps and infer peak coordinates.

        Args:
            inputs: Dictionary with "image" as one of the keys.

        Returns:
            A dictionary of outputs with keys:

            `"pred_instance_peaks"`: The predicted peaks for each instance in the batch
                as a `torch.Tensor` of shape `(samples, nodes, 2)`.
            `"pred_peak_vals"`: The value of the confidence maps at the predicted
                peaks for each instance in the batch as a `torch.Tensor` of shape
                `(samples, nodes)`.

        """
        # Network forward pass.
        self.batch_size = inputs["image"].shape[0]
        output = self.torch_model(inputs["image"])
        cms = output["MultiInstanceConfmapsHead"]
        class_maps = output["ClassMapsHead"]  # (batch, n_classes, h, w)
        cms_peaks, cms_peak_vals, cms_peak_sample_inds, cms_peak_channel_inds = (
            self._generate_cms_peaks(cms.detach())
        )

        cms_peaks = cms_peaks / self.class_maps_output_stride
        (
            predicted_instances,
            predicted_peak_scores,
            predicted_instance_scores,
        ) = classify_peaks_from_maps(
            class_maps.detach(),
            cms_peaks,
            cms_peak_vals,
            cms_peak_sample_inds,
            cms_peak_channel_inds,
            n_channels=cms.shape[-3],
        )
        predicted_instances = [
            p * self.class_maps_output_stride for p in predicted_instances
        ]

        # Adjust for input scaling.
        if self.input_scale != 1.0:
            predicted_instances = [p / self.input_scale for p in predicted_instances]

        predicted_instances_adjusted = []
        for idx, p in enumerate(predicted_instances):
            predicted_instances_adjusted.append(
                p / inputs["eff_scale"][idx].to(p.device)
            )
        out = {
            "pred_instance_peaks": predicted_instances_adjusted,
            "pred_peak_values": predicted_peak_scores,
            "instance_scores": predicted_instance_scores,
        }

        if self.return_confmaps:
            out["pred_confmaps"] = cms.detach()
        if self.return_class_maps:
            out["pred_class_maps"] = class_maps.detach()

        inputs.update(out)
        return [inputs]

__init__(torch_model, cms_output_stride=None, class_maps_output_stride=None, peak_threshold=0.0, refinement='integral', integral_patch_size=5, return_confmaps=False, return_class_maps=False, input_scale=1.0)

Initialise the model attributes.

Source code in sleap_nn/inference/bottomup.py 
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
def __init__(
    self,
    torch_model: L.LightningModule,
    cms_output_stride: Optional[int] = None,
    class_maps_output_stride: Optional[int] = None,
    peak_threshold: float = 0.0,
    refinement: Optional[str] = "integral",
    integral_patch_size: int = 5,
    return_confmaps: Optional[bool] = False,
    return_class_maps: Optional[bool] = False,
    input_scale: float = 1.0,
):
    """Initialise the model attributes."""
    super().__init__()
    self.torch_model = torch_model
    self.peak_threshold = peak_threshold
    self.refinement = refinement
    self.integral_patch_size = integral_patch_size
    self.cms_output_stride = cms_output_stride
    self.class_maps_output_stride = class_maps_output_stride
    self.return_confmaps = return_confmaps
    self.return_class_maps = return_class_maps
    self.input_scale = input_scale

forward(inputs)

Predict confidence maps and infer peak coordinates.

Parameters:

Name Type Description Default
inputs Dict[str, Tensor]

Dictionary with "image" as one of the keys.

 required

Returns:

Type Description
Dict[str, Tensor]

A dictionary of outputs with keys:

"pred_instance_peaks": The predicted peaks for each instance in the batch as a torch.Tensor of shape (samples, nodes, 2). "pred_peak_vals": The value of the confidence maps at the predicted peaks for each instance in the batch as a torch.Tensor of shape (samples, nodes).
Source code in sleap_nn/inference/bottomup.py 
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """Predict confidence maps and infer peak coordinates.

    Args:
        inputs: Dictionary with "image" as one of the keys.

    Returns:
        A dictionary of outputs with keys:

        `"pred_instance_peaks"`: The predicted peaks for each instance in the batch
            as a `torch.Tensor` of shape `(samples, nodes, 2)`.
        `"pred_peak_vals"`: The value of the confidence maps at the predicted
            peaks for each instance in the batch as a `torch.Tensor` of shape
            `(samples, nodes)`.

    """
    # Network forward pass.
    self.batch_size = inputs["image"].shape[0]
    output = self.torch_model(inputs["image"])
    cms = output["MultiInstanceConfmapsHead"]
    class_maps = output["ClassMapsHead"]  # (batch, n_classes, h, w)
    cms_peaks, cms_peak_vals, cms_peak_sample_inds, cms_peak_channel_inds = (
        self._generate_cms_peaks(cms.detach())
    )

    cms_peaks = cms_peaks / self.class_maps_output_stride
    (
        predicted_instances,
        predicted_peak_scores,
        predicted_instance_scores,
    ) = classify_peaks_from_maps(
        class_maps.detach(),
        cms_peaks,
        cms_peak_vals,
        cms_peak_sample_inds,
        cms_peak_channel_inds,
        n_channels=cms.shape[-3],
    )
    predicted_instances = [
        p * self.class_maps_output_stride for p in predicted_instances
    ]

    # Adjust for input scaling.
    if self.input_scale != 1.0:
        predicted_instances = [p / self.input_scale for p in predicted_instances]

    predicted_instances_adjusted = []
    for idx, p in enumerate(predicted_instances):
        predicted_instances_adjusted.append(
            p / inputs["eff_scale"][idx].to(p.device)
        )
    out = {
        "pred_instance_peaks": predicted_instances_adjusted,
        "pred_peak_values": predicted_peak_scores,
        "instance_scores": predicted_instance_scores,
    }

    if self.return_confmaps:
        out["pred_confmaps"] = cms.detach()
    if self.return_class_maps:
        out["pred_class_maps"] = class_maps.detach()

    inputs.update(out)
    return [inputs]