diff --git a/smauglab/transforms/synthseg/functional.py b/smauglab/transforms/synthseg/functional.py index 93d9cc9..954f8f7 100644 --- a/smauglab/transforms/synthseg/functional.py +++ b/smauglab/transforms/synthseg/functional.py @@ -670,15 +670,28 @@ def distinct_values(volume: torch.Tensor) -> torch.Tensor: def _label_counts_per_sample(label_map: torch.Tensor) -> torch.Tensor: - """`(B, max_label + 1)` voxel counts per label, per sample.""" + """`(B, max_label + 1)` voxel counts per label, per sample. Negative values are not counted. + + A negative value is not a label. nnU-Net writes -1 for voxels outside the annotated + region, and `clamp_min(0)` counted every one of them as background -- while + `em_subdivide_labels`, the only caller, then selects that label's voxels with + `label_map == p`. The count exceeded the number of voxels selected, so the + `randint(0, count)` subsample indexed past the end of the gathered values. On CUDA + that is an asynchronous `index out of bounds` device-side assert, reported at + whichever kernel happens to run next, and it only fires when an oversized draw is + actually taken -- so it presented as an occasional, unattributable crash. + """ batch = label_map.shape[0] - flat = label_map.reshape(batch, -1).clamp_min(0) - n_bins = int(flat.max()) + 1 + flat = label_map.reshape(batch, -1) + keep = flat >= 0 + # The clamp is on the offset only, so a negative label cannot reach another sample's + # bins; `keep` is what removes it from the histogram. + n_bins = int(flat.max().clamp_min(0)) + 1 # One `bincount` over row-offset labels: a `scatter_add_` into a handful of bins # has every thread fighting for the same addresses, where bincount accumulates # in shared memory per block. - offset = flat + torch.arange(batch, device=flat.device).view(-1, 1) * n_bins - return torch.bincount(offset.reshape(-1), minlength=batch * n_bins).view(batch, n_bins) + offset = flat.clamp_min(0) + torch.arange(batch, device=flat.device).view(-1, 1) * n_bins + return torch.bincount(offset[keep], minlength=batch * n_bins).view(batch, n_bins) def _em_gmm_1d(x_fit: torch.Tensor, n_components: int, n_iters: int, eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: @@ -805,6 +818,13 @@ def em_subdivide_labels( label_counts = _label_counts_per_sample(label_map) totals = label_counts.sum(dim=0) parents = totals.nonzero().flatten().tolist() # sorted, batch-wide + if not parents: + # No non-negative label anywhere: every voxel is outside the annotated region. + # One undivided region, no EM, and a merge map that resolves -- rather than + # indexing an empty `parents` in the return statement at the end of this + # function. Previously unreachable, because the histogram clamped the negative + # labels into the background bin. + return torch.zeros_like(label_map), [0], [int(background_label)] counts_per_sample = label_counts.tolist() parent_to_idx = {p: i for i, p in enumerate(parents)} lo, hi = int(background_clusters_range[0]), int(background_clusters_range[1]) @@ -814,7 +834,7 @@ def em_subdivide_labels( # decoding shifted every label by +1, so the real background is no longer 0), fall back # to the largest-area label, so it still receives the richer [min, max] background # split rather than the 2-cluster foreground split. - if background_label not in parents and len(parents) > 0: + if background_label not in parents: background_label = int(totals.argmax()) bg_n_shared = int(torch.randint(lo, hi + 1, (1,), device=device)) if same_on_batch else None diff --git a/unit_tests/test_synthseg_em_ignore_label.py b/unit_tests/test_synthseg_em_ignore_label.py new file mode 100644 index 0000000..e09cbb5 --- /dev/null +++ b/unit_tests/test_synthseg_em_ignore_label.py @@ -0,0 +1,102 @@ +"""EM label completion on a label map that carries nnU-Net's ignore label. + +nnU-Net writes -1 for voxels outside the annotated region, and the preprocessed +volumes of a partially-annotated dataset are full of them. `em_subdivide_labels` +took each label's voxel count from a histogram of `label_map.clamp_min(0)`, which +counts every -1 as background, and then selected that label's voxels with +`label_map == p`, which does not. The count was larger than the selection, so the +`randint(0, count)` subsample used to fit the background GMM could draw an index +past the end of the gathered values. + +On CPU that is an `IndexError`. On CUDA it is an asynchronous `index out of bounds` +device-side assert, attributed to whichever kernel happens to run next, and it +only fires when an oversized index is actually drawn -- for a 128^3 patch with six +ignore voxels and 100 000 fit draws, about two runs in five. The symptom was an +occasional crash that pointed at innocent code. + +The first test is the invariant itself and is deterministic; the rest exercise the +call that used to raise, and the all-ignore patch that the clamp made +accidentally representable. +""" + +from __future__ import annotations + +import torch + +from smauglab.transforms.synthseg import functional as FN +from unit_tests.helpers import SmaugLabTestCase + +SHAPE = (16, 16, 16) +VOXELS = SHAPE[0] * SHAPE[1] * SHAPE[2] +FOREGROUND = 8 * 8 * 8 +#: Smaller than the number of background voxels, so the subsample path is the one +#: under test: with `max_fit_voxels >= count` the function uses every voxel and +#: never draws an index at all. +MAX_FIT = 64 +#: Enough ignored voxels that an out-of-range draw is a near-certainty rather than a +#: coin flip: the old histogram would report 3584 background voxels where only 1584 +#: can be selected, so P(no bad index in 64 draws) = (1584/3584)^64 < 1e-23. +N_IGNORE = 2000 + + +def pair_with_ignore(n_ignore: int = N_IGNORE, batch: int = 2): + """A label map with background, one structure and `n_ignore` ignored voxels.""" + depth, height, width = SHAPE + labels = torch.zeros(batch, 1, depth, height, width, dtype=torch.long) + labels[..., 4:12, 4:12, 4:12] = 1 + labels.reshape(batch, -1)[:, :n_ignore] = -1 + image = torch.rand(batch, 1, depth, height, width) + return labels, image + + +class EmIgnoreLabelTests(SmaugLabTestCase): + def test_counts_match_the_selection_they_are_used_for(self): + """The histogram must count exactly the voxels `label_map == p` selects. + + This is the invariant the crash violated, checked directly so a future + rewrite of the histogram cannot reintroduce it without failing here. + """ + labels, _ = pair_with_ignore() + counts = FN._label_counts_per_sample(labels) + for b in range(labels.shape[0]): + for p in counts.sum(dim=0).nonzero().flatten().tolist(): + with self.subTest(sample=b, label=p): + self.assertEqual(int(counts[b][p]), int((labels[b, 0] == p).sum())) + + def test_ignore_voxels_are_not_counted_as_background(self): + labels, _ = pair_with_ignore() + counts = FN._label_counts_per_sample(labels) + # The structure is partly overwritten by the ignore block, so the foreground + # count is read back rather than assumed. + foreground = int((labels[0, 0] == 1).sum()) + self.assertEqual(int(counts[0][0]), VOXELS - foreground - N_IGNORE) + self.assertEqual(int(counts.sum()), labels.numel() - N_IGNORE * labels.shape[0]) + + def test_subdivide_survives_the_ignore_label(self): + """The call that used to raise. `n_iters=1` keeps it quick -- the index is + drawn before the first iteration, so the path under test is covered.""" + labels, image = pair_with_ignore() + fine, gen, out = FN.em_subdivide_labels( + image, + labels, + n_foreground_clusters=2, + background_clusters_range=(3, 3), + background_label=0, + n_iters=1, + max_fit_voxels=MAX_FIT, + ) + self.assertEqual(fine.shape, labels.shape) + self.assertEqual(len(gen), len(out)) + self.assertTrue(set(out).issubset({0, 1}), out) + # The ignored voxels keep generation label 0, which merges back to the + # background parent: they are treated as background, as they were before. + self.assertEqual(int(fine.reshape(labels.shape[0], -1)[0, :N_IGNORE].abs().sum()), 0) + + def test_all_ignore_patch_returns_one_region(self): + """A patch with no annotated voxel at all resolves instead of raising.""" + labels = torch.full((1, 1, *SHAPE), -1, dtype=torch.long) + image = torch.rand(1, 1, *SHAPE) + fine, gen, out = FN.em_subdivide_labels(image, labels, background_label=0, n_iters=1, max_fit_voxels=MAX_FIT) + self.assertEqual(gen, [0]) + self.assertEqual(out, [0]) + self.assertEqual(int(fine.abs().sum()), 0)