Skip to content

bottomup

sleap_nn.inference.layers.bottomup

BottomUpLayer — single-stage multi-instance inference via PAF grouping.

The bottom-up model emits two heads: a multi-instance confidence map (MultiInstanceConfmapsHead) and part-affinity fields (PartAffinityFieldsHead). The layer:

  1. Runs the model.
  2. Finds local peaks on the confmaps.
  3. Scores PAF lines between candidate keypoints.
  4. Groups peaks into instances via the existing :class:PAFScorer.
  5. NaN-pads the variable per-frame instance count to max_instances so Outputs.pred_keypoints retains its canonical (B, I, N, 2) shape.

Steps 1-3 are GPU-friendly tensor ops; step 4 is a CPU-bound scipy.linear_sum_assignment + BFS instance assembly. The two phases are split into :meth:_score_pafs_on_gpu (GPU) and the free function :func:sleap_nn.inference.streaming.group_scored_batch (CPU). The split enables a worker pool for the CPU phase; the inline path simply calls them back-to-back inside :meth:postprocess.

Classes:

Name Description
BottomUpLayer

Bottom-up multi-instance inference layer.

BottomUpLayer

Bases: InferenceLayer

Bottom-up multi-instance inference layer.

Parameters:

Name Type Description Default
backend ModelBackend

Runtime backend wrapping the bottom-up Lightning module (which emits both confmaps and PAFs).

 required
paf_scorer PAFScorer

Pre-configured :class:PAFScorer for instance grouping.

 required
cms_output_stride int

Stride between the confmap output and the scaled-input pixels.

 required
pafs_output_stride int

Stride between the PAF output and the scaled-input pixels (often equal to cms_output_stride).

 required
max_instances Optional[int]

Cap on instances per frame. Variable bottom-up output is NaN-padded to this fixed shape.

 None
max_stride int

Maximum stride the model requires for input divisibility (padding applied bottom-right).

 1
max_peaks_per_node Optional[int]

Skip PAF scoring entirely if any node has more peaks than this — prevents combinatorial explosion on noisy early-training models. None disables.

 None
preprocess_config / postprocess_config

Standard knobs.

 required

Methods:

Name Description
__init__

Compose the layer with the PAF scorer + stride config.
grouping_params

Snapshot the layer's grouping configuration as a value type.
postprocess

GPU peak/PAF scoring + CPU grouping, both in this process.
Source code in sleap_nn/inference/layers/bottomup.py 
 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
class BottomUpLayer(InferenceLayer):
    """Bottom-up multi-instance inference layer.

    Args:
        backend: Runtime backend wrapping the bottom-up Lightning module
            (which emits both confmaps and PAFs).
        paf_scorer: Pre-configured :class:`PAFScorer` for instance grouping.
        cms_output_stride: Stride between the confmap output and the
            scaled-input pixels.
        pafs_output_stride: Stride between the PAF output and the
            scaled-input pixels (often equal to ``cms_output_stride``).
        max_instances: Cap on instances per frame. Variable bottom-up
            output is NaN-padded to this fixed shape.
        max_stride: Maximum stride the model requires for input
            divisibility (padding applied bottom-right).
        max_peaks_per_node: Skip PAF scoring entirely if any node has
            more peaks than this — prevents combinatorial explosion on
            noisy early-training models. ``None`` disables.
        preprocess_config / postprocess_config: Standard knobs.
    """

    def __init__(
        self,
        backend: ModelBackend,
        paf_scorer: PAFScorer,
        cms_output_stride: int,
        pafs_output_stride: int,
        max_instances: Optional[int] = None,
        max_stride: int = 1,
        max_peaks_per_node: Optional[int] = None,
        preprocess_config: Optional[PreprocessConfig] = None,
        postprocess_config: Optional[PostprocessConfig] = None,
    ) -> None:
        """Compose the layer with the PAF scorer + stride config."""
        super().__init__(
            backend=backend,
            preprocess_config=preprocess_config or PreprocessConfig(),
            postprocess_config=postprocess_config or PostprocessConfig(),
            output_stride=cms_output_stride,
            max_stride=max_stride,
        )
        self.paf_scorer = paf_scorer
        self.cms_output_stride = cms_output_stride
        self.pafs_output_stride = pafs_output_stride
        self.max_instances = max_instances
        self.max_peaks_per_node = max_peaks_per_node

    # ──────────────────────────────────────────────────────────────────
    # GPU stage — peaks + PAF line scoring
    # ──────────────────────────────────────────────────────────────────

    def _score_pafs_on_gpu(self, raw_out: dict, info: PreprocInfo) -> ScoredBatch:
        """Run the GPU phase: local-peak find + PAF line scoring.

        Output is a CPU-resident :class:`ScoredBatch` so it can be
        shipped to a worker process verbatim (or fed straight into
        :func:`group_scored_batch` inline).
        """
        cms = raw_out["MultiInstanceConfmapsHead"]
        pafs = raw_out["PartAffinityFieldsHead"].permute(0, 2, 3, 1)  # (B, H, W, 2*E)

        peaks, peak_vals, sample_inds, peak_channel_inds = find_local_peaks(
            cms.detach(),
            threshold=self.postprocess_config.peak_threshold,
            refinement=self.postprocess_config.effective_refinement,
            integral_patch_size=self.postprocess_config.integral_patch_size,
        )
        peaks = peaks * self.cms_output_stride

        B = cms.shape[0]
        n_nodes = cms.shape[1]

        # Per-batch peak lists for the PAF scorer.
        cms_peaks: list[torch.Tensor] = []
        cms_peak_vals: list[torch.Tensor] = []
        cms_peak_channel_inds: list[torch.Tensor] = []
        for b in range(B):
            mask = sample_inds == b
            cms_peaks.append(peaks[mask])
            cms_peak_vals.append(peak_vals[mask].to(torch.float32))
            cms_peak_channel_inds.append(peak_channel_inds[mask])

        # Skip PAF scoring if any node has too many peaks (combinatorial blowup).
        skip_paf = False
        if self.max_peaks_per_node is not None:
            for ch_inds in cms_peak_channel_inds:
                for ni in range(n_nodes):
                    n_peaks = int((ch_inds == ni).sum().item())
                    if n_peaks > self.max_peaks_per_node:
                        skip_paf = True
                        # Surface the skip (legacy logged it); an over-limit node
                        # usually means an under-trained / noisy model (#584).
                        logger.warning(
                            "Skipping PAF scoring for a bottom-up batch: node "
                            f"{ni} has {n_peaks} peaks "
                            f"(max_peaks_per_node={self.max_peaks_per_node}). "
                            "Instances will be all-NaN; the model may need more "
                            "training or a higher peak threshold."
                        )
                        break
                if skip_paf:
                    break

        if skip_paf:
            return ScoredBatch(
                cms_peaks=cms_peaks,
                cms_peak_vals=cms_peak_vals,
                cms_peak_channel_inds=cms_peak_channel_inds,
                edge_inds=[],
                edge_peak_inds=[],
                line_scores=[],
                info=info,
                n_samples=B,
                n_nodes=n_nodes,
                skip_paf=True,
                cms=cms.detach() if self.postprocess_config.return_confmaps else None,
                pafs=pafs.detach() if self.postprocess_config.return_pafs else None,
            ).to_cpu()

        edge_inds, edge_peak_inds, line_scores = self.paf_scorer.score_paf_lines(
            pafs, cms_peaks, cms_peak_channel_inds
        )

        scored = ScoredBatch(
            cms_peaks=cms_peaks,
            cms_peak_vals=cms_peak_vals,
            cms_peak_channel_inds=cms_peak_channel_inds,
            edge_inds=edge_inds,
            edge_peak_inds=edge_peak_inds,
            line_scores=line_scores,
            info=info,
            n_samples=B,
            n_nodes=n_nodes,
            skip_paf=False,
            cms=(
                cms.detach()
                if (
                    self.postprocess_config.return_confmaps
                    or self.postprocess_config.return_paf_graph
                )
                else None
            ),
            pafs=(
                pafs.detach()
                if (
                    self.postprocess_config.return_pafs
                    or self.postprocess_config.return_paf_graph
                )
                else None
            ),
        )
        return scored.to_cpu()

    # ──────────────────────────────────────────────────────────────────
    # Layer-level grouping params (passed to inline + pool paths alike)
    # ──────────────────────────────────────────────────────────────────

    def grouping_params(self) -> GroupingParams:
        """Snapshot the layer's grouping configuration as a value type.

        The result is picklable and can be sent to a worker process.
        """
        # Prefer a predict-time override (set on ``postprocess_config`` by
        # ``Predictor._scoped_postprocess_layer``) over the build-time value
        # carried on ``self.max_instances`` (#582).
        max_instances = getattr(self.postprocess_config, "max_instances", None)
        if max_instances is None:
            max_instances = self.max_instances
        return GroupingParams(
            paf_scorer_kwargs={
                "part_names": list(self.paf_scorer.part_names),
                "edges": [tuple(e) for e in self.paf_scorer.edges],
                "pafs_stride": self.paf_scorer.pafs_stride,
                "max_edge_length_ratio": self.paf_scorer.max_edge_length_ratio,
                "dist_penalty_weight": self.paf_scorer.dist_penalty_weight,
                "n_points": self.paf_scorer.n_points,
                "min_instance_peaks": self.paf_scorer.min_instance_peaks,
                "min_line_scores": self.paf_scorer.min_line_scores,
            },
            max_instances=max_instances,
            return_confmaps=self.postprocess_config.return_confmaps,
            return_pafs=self.postprocess_config.return_pafs,
            return_paf_graph=self.postprocess_config.return_paf_graph,
        )

    # ──────────────────────────────────────────────────────────────────
    # Postprocess — inline composition (parity path)
    # ──────────────────────────────────────────────────────────────────

    def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
        """GPU peak/PAF scoring + CPU grouping, both in this process."""
        scored = self._score_pafs_on_gpu(raw_out, info)
        return group_scored_batch(scored, self.grouping_params())

__init__(backend, paf_scorer, cms_output_stride, pafs_output_stride, max_instances=None, max_stride=1, max_peaks_per_node=None, preprocess_config=None, postprocess_config=None)

Compose the layer with the PAF scorer + stride config.

Source code in sleap_nn/inference/layers/bottomup.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def __init__(
    self,
    backend: ModelBackend,
    paf_scorer: PAFScorer,
    cms_output_stride: int,
    pafs_output_stride: int,
    max_instances: Optional[int] = None,
    max_stride: int = 1,
    max_peaks_per_node: Optional[int] = None,
    preprocess_config: Optional[PreprocessConfig] = None,
    postprocess_config: Optional[PostprocessConfig] = None,
) -> None:
    """Compose the layer with the PAF scorer + stride config."""
    super().__init__(
        backend=backend,
        preprocess_config=preprocess_config or PreprocessConfig(),
        postprocess_config=postprocess_config or PostprocessConfig(),
        output_stride=cms_output_stride,
        max_stride=max_stride,
    )
    self.paf_scorer = paf_scorer
    self.cms_output_stride = cms_output_stride
    self.pafs_output_stride = pafs_output_stride
    self.max_instances = max_instances
    self.max_peaks_per_node = max_peaks_per_node

grouping_params()

Snapshot the layer's grouping configuration as a value type.

The result is picklable and can be sent to a worker process.

Source code in sleap_nn/inference/layers/bottomup.py 
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
def grouping_params(self) -> GroupingParams:
    """Snapshot the layer's grouping configuration as a value type.

    The result is picklable and can be sent to a worker process.
    """
    # Prefer a predict-time override (set on ``postprocess_config`` by
    # ``Predictor._scoped_postprocess_layer``) over the build-time value
    # carried on ``self.max_instances`` (#582).
    max_instances = getattr(self.postprocess_config, "max_instances", None)
    if max_instances is None:
        max_instances = self.max_instances
    return GroupingParams(
        paf_scorer_kwargs={
            "part_names": list(self.paf_scorer.part_names),
            "edges": [tuple(e) for e in self.paf_scorer.edges],
            "pafs_stride": self.paf_scorer.pafs_stride,
            "max_edge_length_ratio": self.paf_scorer.max_edge_length_ratio,
            "dist_penalty_weight": self.paf_scorer.dist_penalty_weight,
            "n_points": self.paf_scorer.n_points,
            "min_instance_peaks": self.paf_scorer.min_instance_peaks,
            "min_line_scores": self.paf_scorer.min_line_scores,
        },
        max_instances=max_instances,
        return_confmaps=self.postprocess_config.return_confmaps,
        return_pafs=self.postprocess_config.return_pafs,
        return_paf_graph=self.postprocess_config.return_paf_graph,
    )

postprocess(raw_out, info)

GPU peak/PAF scoring + CPU grouping, both in this process.

Source code in sleap_nn/inference/layers/bottomup.py 
233
234
235
236
def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
    """GPU peak/PAF scoring + CPU grouping, both in this process."""
    scored = self._score_pafs_on_gpu(raw_out, info)
    return group_scored_batch(scored, self.grouping_params())