diff --git a/smauglab/transforms/cpu/fromSeg.py b/smauglab/transforms/cpu/fromSeg.py index c52b42c..30e9f2c 100644 --- a/smauglab/transforms/cpu/fromSeg.py +++ b/smauglab/transforms/cpu/fromSeg.py @@ -73,7 +73,10 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) # Convert to NumPy for dilation operations (not supported in PyTorch) l_mask_np = l_mask.cpu().numpy() - struct = ndi.iterate_structure(ndi.generate_binary_structure(3, 1), 3) + # Rank from the data, not hardcoded 3: scipy requires the structuring element + # to match the input's rank, so a 2-D image raised + # "structure rank must match input rank" here. + struct = ndi.iterate_structure(ndi.generate_binary_structure(l_mask_np.ndim, 1), 3) l_mask_dilate_np = ndi.binary_dilation(l_mask_np, structure=struct) # Convert back to PyTorch diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index c26f8ea..358734b 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -30,6 +30,26 @@ def _choose_region_mode(p_in: float, p_out: float, seg_mask: torch.Tensor | None return "all" +def _foreground(mask: torch.Tensor, dim: int) -> torch.Tensor: + """Which voxels the segmentation covers, reducing over the class axis. + + This used to be `torch.argmax(mask, dim) > 0`, which is only "is anything labelled + here" if class 0 is background -- and it is not: + + * For a single-channel [B, 1, D, H, W] mask (an ordinary nnU-Net target, and what + the tests build) `argmax` over a length-1 axis is always 0, so the result was + **all False**. `in_seg` then applied the transform nowhere and `out_seg` applied + it everywhere: the two knobs did nothing and the opposite of nothing. + * For a one-hot mask, this repository's convention (`collapse_onehot_to_index` in + gpu/fromSeg.py) is that channel `c` encodes label `c + 1` with background + implicit, so `argmax == 0` is a real foreground class and was being dropped. + + `amax > 0` asks the question that was meant, and matches what + `collapse_onehot_to_index` already does with `seg_raw.any(dim=1)`. + """ + return mask.amax(dim=dim) > 0 + + def _apply_region_mode( orig: torch.Tensor, transformed: torch.Tensor, @@ -67,7 +87,7 @@ def _apply_region_mode( o = torch.randint(0, 2, (seg_mask.shape[1],), device=seg_mask.device, dtype=seg_mask.dtype) m[i] = m[i] * o.view(-1, 1, 1, 1) # Broadcasting o to match the dimensions of m - m = torch.argmax(m, dim=1) > 0 + m = _foreground(m, dim=1) m = m.to(transformed.dtype) if mode == "out": m = 1.0 - m @@ -85,7 +105,7 @@ def _apply_region_mode( # Create a tensor with random one and zero o = torch.randint(0, 2, (seg_mask.shape[0],), device=seg_mask.device, dtype=seg_mask.dtype) m = m * o.view(-1, 1, 1, 1) # Broadcasting o to match the dimensions of m - m = torch.argmax(m, dim=0) > 0 + m = _foreground(m, dim=0) m = m.to(transformed.dtype) if mode == "out": m = 1.0 - m @@ -779,8 +799,17 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ orig_means = x.mean(dim=reduce_dims) orig_stds = x.std(dim=reduce_dims) - # Normalize to make values >=0 - x = (x - x.min()) / (x.max() - x.min() + 0.00001) + # Normalize to make values >=0, per sample. + # + # This used to be a bare `x.min()` / `x.max()`, which reduces over the whole + # [N, ...] slab: an image's augmentation then depended on which other images + # happened to share its batch, so the same volume augmented twice in + # different batches came out differently. Every other transform in this file + # reduces over `dim=reduce_dims` per sample. + keep_dims = tuple(range(1, x.dim())) + x_min = x.amin(dim=keep_dims, keepdim=True) + x_max = x.amax(dim=keep_dims, keepdim=True) + x = (x - x_min) / (x_max - x_min + 0.00001) # Apply function x = self.func(x) @@ -932,7 +961,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # Apply histogram equalization transform seg_mask = params.get("seg") for c in self.apply_to_channel: - channel_data = input[:, c] # shape [N, ...spatial...] + # `.clone()`, not the bare `input[:, c]` view this used to take: the loop + # below assigns into `channel_data[b]`, which through a view writes straight + # into `input`. The non-finite guard at the bottom would then `continue` + # over values that were already in the batch -- the guard skipped nothing. + channel_data = input[:, c].clone() # shape [N, ...spatial...] orig = channel_data.clone() if self.retain_stats: diff --git a/unit_tests/test_region_and_stats.py b/unit_tests/test_region_and_stats.py new file mode 100644 index 0000000..caea01c --- /dev/null +++ b/unit_tests/test_region_and_stats.py @@ -0,0 +1,246 @@ +"""Region selection and per-sample statistics. + +Three defects, each of which made a transform depend on something it should not: + +* `_apply_region_mode` reduced the mask's class axis with `torch.argmax(mask, dim) > 0`. + For the ordinary single-channel mask that is always False, so `in_seg` applied the + transform nowhere and `out_seg` applied it everywhere -- both knobs silently + inoperative. For a one-hot mask it dropped the first foreground class, because this + repository encodes channel `c` as label `c + 1`. +* The elementwise function transforms normalised with a slab-wide `x.min()`/`x.max()`, + so a volume's augmentation depended on which other volumes shared its batch. +* `RandomHistogramEqualizationGPU` wrote through an `input[:, c]` view, so its + non-finite guard skipped over values already in the batch. +""" + +import numpy as np +import torch + +from smauglab.transforms.gpu.base import AugmentationSequentialCustom +from smauglab.transforms.gpu.contrast import ( + RandomConvTransformGPU, + RandomFunctionGPU, + RandomHistogramEqualizationGPU, + _apply_region_mode, + _foreground, +) +from unit_tests.helpers import SmaugLabTestCase, first_output + + +class TestForegroundReduction(SmaugLabTestCase): + def test_a_single_channel_mask_is_not_collapsed_to_nothing(self): + """argmax over a length-1 axis is always 0, so `> 0` was always False.""" + mask = torch.zeros(1, 1, 4, 4, 4) + mask[0, 0, 1:3, 1:3, 1:3] = 1.0 + + found = _foreground(mask, dim=1) + + self.assertEqual(int(found.sum()), 8, "the labelled block was not recognised as foreground") + self.assertTrue(bool(found[0, 1, 1, 1])) + self.assertFalse(bool(found[0, 0, 0, 0])) + + def test_the_old_argmax_reduction_really_did_collapse_it(self): + """Control: what the previous expression produced on the same input.""" + mask = torch.zeros(1, 1, 4, 4, 4) + mask[0, 0, 1:3, 1:3, 1:3] = 1.0 + self.assertEqual(int((torch.argmax(mask, dim=1) > 0).sum()), 0) + + def test_the_first_one_hot_channel_counts_as_foreground(self): + """Channel 0 encodes label 1 -- see collapse_onehot_to_index in gpu/fromSeg.py.""" + mask = torch.zeros(1, 3, 4, 4, 4) + mask[0, 0, 0, 0, 0] = 1.0 # only the first class is present here + mask[0, 2, 3, 3, 3] = 1.0 + + found = _foreground(mask, dim=1) + + self.assertTrue(bool(found[0, 0, 0, 0]), "the first one-hot class was dropped") + self.assertTrue(bool(found[0, 3, 3, 3])) + self.assertEqual(int(found.sum()), 2) + + def test_background_stays_background(self): + self.assertEqual(int(_foreground(torch.zeros(1, 4, 5, 5, 5), dim=1).sum()), 0) + + +class TestApplyRegionMode(SmaugLabTestCase): + """The behaviour a config actually asks for when it sets in_seg or out_seg.""" + + def setUp(self): + super().setUp() + # orig/transformed are one channel of the batch: [N, D, H, W]. + self.orig = torch.zeros(1, 4, 4, 4) + self.transformed = torch.ones(1, 4, 4, 4) + # A single-channel mask covering one corner: [N, 1, D, H, W]. + self.mask = torch.zeros(1, 1, 4, 4, 4) + self.mask[0, 0, :2, :2, :2] = 1.0 + self.inside = (slice(None), slice(0, 2), slice(0, 2), slice(0, 2)) + + def test_mode_in_changes_only_the_masked_voxels(self): + out = _apply_region_mode(self.orig, self.transformed, self.mask, "in") + + self.assertTrue(bool((out[self.inside] == 1.0).all()), "'in' did not apply the transform inside the mask") + self.assertEqual(int(out.sum()), 8, "'in' leaked outside the mask") + + def test_mode_out_changes_only_the_unmasked_voxels(self): + out = _apply_region_mode(self.orig, self.transformed, self.mask, "out") + + self.assertTrue(bool((out[self.inside] == 0.0).all()), "'out' applied the transform inside the mask") + self.assertEqual(int(out.sum()), 64 - 8) + + def test_in_and_out_partition_the_volume(self): + inside = _apply_region_mode(self.orig, self.transformed, self.mask, "in") + outside = _apply_region_mode(self.orig, self.transformed, self.mask, "out") + self.assertTrue(torch.equal(inside + outside, self.transformed)) + + def test_mode_all_ignores_the_mask(self): + out = _apply_region_mode(self.orig, self.transformed, self.mask, "all") + self.assertTrue(torch.equal(out, self.transformed)) + + def test_a_missing_mask_means_apply_everywhere(self): + out = _apply_region_mode(self.orig, self.transformed, None, "in") + self.assertTrue(torch.equal(out, self.transformed)) + + def test_the_unbatched_3d_path_behaves_the_same(self): + orig = torch.zeros(4, 4, 4) + transformed = torch.ones(4, 4, 4) + mask = torch.zeros(2, 4, 4, 4) + mask[0, :2, :2, :2] = 1.0 + + out = _apply_region_mode(orig, transformed, mask, "in") + self.assertEqual(int(out.sum()), 8) + + def test_an_unsupported_rank_is_rejected(self): + with self.assertRaises(ValueError): + _apply_region_mode(torch.rand(2, 2), torch.rand(2, 2), torch.rand(1, 2, 2), "in") + + +class TestRegionModeReachesTheRealTransforms(SmaugLabTestCase): + """End to end: a GPU transform with in_seg=1.0 must respect the mask.""" + + def test_in_seg_confines_a_scharr_transform_to_the_mask(self): + image = self.tiny_volume() + seg = self.tiny_seg() # single channel, a centred cube + + # Driven through the container, as AugTransformsGPU does: that is what routes + # the mask into params["seg"], which is where _apply_region_mode reads it. + pipeline = AugmentationSequentialCustom( + RandomConvTransformGPU(kernel_type="Scharr", p=1.0, in_seg=1.0, out_seg=0.0, mix_prob=0.0), + data_keys=["input", "mask"], + same_on_batch=True, + ) + out = first_output(pipeline(image.clone(), seg.clone())) + + outside = ~(seg[0, 0] > 0) + self.assertTrue( + torch.allclose(out[0, 0][outside], image[0, 0][outside], atol=1e-5), + "in_seg=1.0 changed voxels outside the segmentation", + ) + self.assertFalse( + torch.allclose(out[0, 0], image[0, 0], atol=1e-5), + "in_seg=1.0 changed nothing at all -- this is what the argmax bug did", + ) + + +class TestFunctionTransformIsPerSample(SmaugLabTestCase): + """The normalisation used a slab-wide min/max, coupling every volume in the batch.""" + + def _run(self, volume: torch.Tensor) -> torch.Tensor: + torch.manual_seed(0) + transform = RandomFunctionGPU(func=torch.sqrt, p=1.0) + return transform.apply_transform(volume.clone(), {}, {}, transform=None) + + def test_a_volume_is_augmented_the_same_alone_and_in_a_batch(self): + subject = torch.rand(1, 1, 8, 8, 8) + # A neighbour with a much wider range: under a slab-wide min/max it drags the + # subject's normalisation with it. + neighbour = torch.rand(1, 1, 8, 8, 8) * 100.0 + batch = torch.cat([subject, neighbour], dim=0) + + alone = self._run(subject) + together = self._run(batch) + + self.assertTrue( + torch.allclose(alone[0, 0], together[0, 0], atol=1e-5), + "the same volume was augmented differently depending on its batch neighbours", + ) + + def test_each_sample_is_normalised_onto_its_own_range(self): + batch = torch.cat([torch.rand(1, 1, 6, 6, 6), torch.rand(1, 1, 6, 6, 6) * 50.0], dim=0) + + out = self._run(batch) + + for b in range(2): + with self.subTest(sample=b): + # sqrt of a [0, 1]-normalised sample still spans close to the full range. + self.assertAlmostEqual(float(out[b, 0].max()), 1.0, places=3) + + +class TestNonFiniteGuard(SmaugLabTestCase): + """Guard tests for `RandomHistogramEqualizationGPU`, not regression tests. + + The `.clone()` this pins is unobservable from outside today, in the same way as + the `resample` change in the previous PR. `channel_data` was `input[:, c]`, a view, + and the loop assigns into `channel_data[b]` -- so by the time the non-finite guard + at the bottom `continue`d, every value it meant to withhold was already in the + batch, and the closing `input[:, c] = channel_data` was a no-op. With the clone the + guard does what it says. Reaching it needs the equalisation to *produce* a + non-finite value from finite input, which no input constructed here manages: a NaN + supplied by the caller raises out of `torch.histc` first. So these tests pin the + behaviour the clone must not disturb, and pass either way. + """ + + def test_a_degenerate_constant_channel_stays_finite(self): + transform = RandomHistogramEqualizationGPU(p=1.0, mix_prob=0.0) + # A constant channel makes img_max == img_min, the degenerate histogram case. + volume = torch.ones(1, 1, 8, 8, 8) + + out = transform.apply_transform(volume.clone(), {}, {}, transform=None) + + self.assertTrue(bool(torch.isfinite(out).all()), "a non-finite result reached the batch") + + def test_a_non_finite_input_is_still_rejected_loudly(self): + """Documents why the guard cannot be exercised: histc rejects the range first.""" + transform = RandomHistogramEqualizationGPU(p=1.0, mix_prob=0.0) + volume = torch.rand(1, 1, 8, 8, 8) + volume[0, 0, 0, 0, 0] = float("nan") + + with self.assertRaises(RuntimeError): + transform.apply_transform(volume, {}, {}, transform=None) + + def test_equalisation_still_changes_the_image(self): + transform = RandomHistogramEqualizationGPU(p=1.0, mix_prob=0.0) + volume = torch.rand(1, 1, 8, 8, 8) + + out = transform.apply_transform(volume.clone(), {}, {}, transform=None) + + self.assertFalse(torch.allclose(out, volume, atol=1e-6)) + + +class TestDilationRankFollowsTheData(SmaugLabTestCase): + """scipy requires the structuring element's rank to match the input's. + + `aug_redistribute_seg` is called per channel (`img[c], seg[c]`), so it sees a bare + spatial array -- 3-D for a volume, 2-D for an image -- and returns (img, seg). + """ + + def test_a_2d_image_does_not_raise(self): + from smauglab.transforms.cpu.fromSeg import aug_redistribute_seg + + image = torch.rand(12, 12) + seg = torch.zeros(12, 12) + seg[3:8, 3:8] = 1.0 + + out, _ = aug_redistribute_seg(image.clone(), seg, in_seg=1.0) + + self.assertEqual(tuple(out.shape), tuple(image.shape)) + self.assertTrue(bool(np.isfinite(out.numpy()).all())) + + def test_the_3d_path_is_unchanged(self): + from smauglab.transforms.cpu.fromSeg import aug_redistribute_seg + + image = torch.rand(10, 10, 10) + seg = torch.zeros(10, 10, 10) + seg[2:6, 2:6, 2:6] = 1.0 + + out, _ = aug_redistribute_seg(image.clone(), seg, in_seg=1.0) + + self.assertEqual(tuple(out.shape), tuple(image.shape))