diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 221a5ff..b9c912c 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -51,7 +51,15 @@ jobs: - uses: actions/setup-python@v5 with: - python-version: "3.12" + # The floor of the supported range, matching `python_version = "3.10"` in + # [tool.mypy] -- not the newest interpreter. mypy applies its target version + # when parsing *every* file, third-party stubs included, so a runner that + # resolves a dependency too new for that target breaks the job with an error + # in someone else's .pyi. That is what numpy did: 2.3+ requires Python >=3.11 + # and its stubs use PEP 695 `type` statements, so a 3.12 runner installed + # numpy 2.5 and mypy rejected numpy/__init__.pyi as 3.12-only syntax. + # Installing on 3.10 resolves numpy 2.2.x, whose stubs parse under the target. + python-version: "3.10" cache: pip cache-dependency-path: pyproject.toml diff --git a/smauglab/transforms/gpu/spatial.py b/smauglab/transforms/gpu/spatial.py index 5ac9e5d..6801040 100644 --- a/smauglab/transforms/gpu/spatial.py +++ b/smauglab/transforms/gpu/spatial.py @@ -187,9 +187,10 @@ def apply_transform_mask( Convert "resample" arguments to "nearest" by default. """ - resample_method: Resample | None - if "resample" in flags: - resample_method = flags["resample"] + # `resample_method` was declared but only *assigned* inside the `if`, so the + # restore below raised UnboundLocalError whenever flags carried no "resample". + resample_method: Resample | None = flags.get("resample") + if resample_method is not None: flags["resample"] = Resample.get("nearest") output = self.apply_transform(input, params, flags, transform) if resample_method is not None: @@ -293,6 +294,24 @@ def apply_transform_mask( return output +def _choose_axis(batch_size: int, device: torch.device, same_on_batch: bool) -> torch.Tensor: + """Pick the single axis to act on, per batch element. Returns `[B]` indices. + + Drawn here, from `forward`, rather than in `make_samplers`. kornia calls + `make_samplers` once and caches the samplers it builds, so an axis picked there was + fixed for the transform's lifetime -- "degrade a random axis" degraded the *same* + axis for a whole training run. + """ + keep = torch.randint(0, 3, (1 if same_on_batch else batch_size,), device=device) + return keep.expand(batch_size) if same_on_batch else keep + + +def _keep_one_axis(values: torch.Tensor, keep: torch.Tensor, neutral: float) -> torch.Tensor: + """Keep column `keep[b]` of a `[B, 3]` draw and set the other two to `neutral`.""" + selected = torch.arange(3, device=values.device).unsqueeze(0) == keep.unsqueeze(1) + return torch.where(selected, values, torch.full_like(values, neutral)) + + class ScaleGenerator3D(RandomGeneratorBase): def __init__(self, scale: tuple[float, float], one_dim: bool = False) -> None: super().__init__() @@ -301,13 +320,6 @@ def __init__(self, scale: tuple[float, float], one_dim: bool = False) -> None: def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None: scale = _tuple_range_reader(self.scale, 3, device, dtype) - if self.one_dim: - # Pick a random dimension to apply scaling - dim = torch.randint(0, 3, (1,)).item() - for i in range(3): - if i != dim: - scale[i, 0] = 1.0 - scale[i, 1] = 1.0 self.scalex_sampler = UniformDistribution(scale[0, 0], scale[0, 1], validate_args=False) self.scaley_sampler = UniformDistribution(scale[1, 0], scale[1, 1], validate_args=False) self.scalez_sampler = UniformDistribution(scale[2, 0], scale[2, 1], validate_args=False) @@ -322,6 +334,10 @@ def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> scalez = _adapted_rsampling((batch_size,), self.scalez_sampler, same_on_batch) scale = torch.stack([scalex, scaley, scalez], dim=1) + if self.one_dim: + # A scale of 1.0 leaves an axis at full resolution. + scale = _keep_one_axis(scale, _choose_axis(batch_size, scale.device, same_on_batch), 1.0) + return {"scale": torch.as_tensor(scale, device=_device, dtype=_dtype)} @@ -449,16 +465,24 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ batch_size, C, D, H, W = input.shape - # Expect params to contain 'flip' tensor of shape [B, 3] with 0/1 values - # flips = None - # if params is not None and 'flip' in params: - # flips = params['flip'] + # params["flip"] is [B, 3] of 0/1 flags over (z, y, x), produced by + # FlipGenerator3D. Reading it is what makes this transform random: the loop + # below used to recompute the same `flip_axis`-derived list for every b and + # ignore the sampled flags entirely, so every call flipped all configured axes + # identically -- three seeded calls gave byte-identical output, and the + # generator (including its "at least one axis" guarantee) was dead code. + flips = params.get("flip") out = input.clone() - # For each batch element, build list of spatial dims to flip (D,H,W -> dims 2,3,4) + # For each batch element, build list of spatial dims to flip. `input[b]` is + # [C, D, H, W], so spatial axis i sits at dim 1 + i. for b in range(batch_size): - # fb expected as length-3 tensor for (z,y,x) - flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis] + if flips is None: + # No sampled flags (a caller invoking apply_transform directly): fall + # back to flipping every configured axis. + flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis] + else: + flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis and bool(flips[b, axis])] if len(flip_dims) > 0: out[b] = torch.flip(input[b], dims=tuple(flip_dims)) @@ -649,25 +673,11 @@ def __init__(self, crop: tuple[float, float], pos: tuple[float, float], one_dim: def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None: crop = _tuple_range_reader(self.crop, 3, device, dtype) - if self.one_dim: - # Pick a random dimension to apply cropping - dim = torch.randint(0, 3, (1,)).item() - for i in range(3): - if i != dim: - crop[i, 0] = 1.0 - crop[i, 1] = 1.0 self.cropx_sampler = UniformDistribution(crop[0, 0], crop[0, 1], validate_args=False) self.cropy_sampler = UniformDistribution(crop[1, 0], crop[1, 1], validate_args=False) self.cropz_sampler = UniformDistribution(crop[2, 0], crop[2, 1], validate_args=False) pos = _tuple_range_reader(self.pos, 3, device, dtype) - if self.one_dim: - # Pick a random dimension to apply cropping - dim = torch.randint(0, 3, (1,)).item() - for i in range(3): - if i != dim: - pos[i, 0] = 1.0 - pos[i, 1] = 1.0 self.posx_sampler = UniformDistribution(pos[0, 0], pos[0, 1], validate_args=False) self.posy_sampler = UniformDistribution(pos[1, 0], pos[1, 1], validate_args=False) self.posz_sampler = UniformDistribution(pos[2, 0], pos[2, 1], validate_args=False) @@ -689,4 +699,17 @@ def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> posz = _adapted_rsampling((batch_size,), self.posz_sampler, same_on_batch) pos = torch.stack([posx, posy, posz], dim=1) + if self.one_dim: + # One axis for both: `make_samplers` drew a separate `dim` for crop and for + # pos, so the crop could be taken along one axis while the position that + # placed it was randomised along another. + keep = _choose_axis(batch_size, crop.device, same_on_batch) + # A crop fraction of 1.0 keeps the whole axis. The *position*, though, is + # the crop centre as a fraction of the axis, so its neutral value is 0.5 + # (centred) -- the previous code copied the crop's 1.0 onto it, which put + # the box centre on the far edge and left the crop flush against it after + # clamping. + crop = _keep_one_axis(crop, keep, 1.0) + pos = _keep_one_axis(pos, keep, 0.5) + return {"crop": torch.as_tensor(crop, device=_device, dtype=_dtype), "pos": torch.as_tensor(pos, device=_device, dtype=_dtype)} diff --git a/unit_tests/test_spatial_sampling.py b/unit_tests/test_spatial_sampling.py new file mode 100644 index 0000000..592acf2 --- /dev/null +++ b/unit_tests/test_spatial_sampling.py @@ -0,0 +1,193 @@ +"""Spatial transforms that were advertised as random but were not. + +The first two groups fail against the implementation that preceded them: + +* `RandomFlipTransformGPU` never read the flip flags its own generator sampled, so it + flipped every configured axis, identically, on every call. +* The single-axis generators drew their axis in `make_samplers`, which kornia calls + once and caches -- so "degrade a random axis" degraded the same axis for a whole + training run, and `CropGenerator3D` neutralised the crop *position* to the far edge + instead of the centre. + +`TestMaskResampleRestore` is a guard rather than a regression test; see its docstring. +""" + +import torch +from kornia.constants import Resample + +from smauglab.transforms.gpu.base import AugmentationSequentialCustom +from smauglab.transforms.gpu.spatial import CropGenerator3D, RandomAffine3DCustom, RandomFlipTransformGPU, ScaleGenerator3D +from unit_tests.helpers import SmaugLabTestCase, first_output + + +class TestFlipIsActuallyRandom(SmaugLabTestCase): + def _pipeline(self, **kwargs): + return AugmentationSequentialCustom( + RandomFlipTransformGPU(p=1.0, **kwargs), + data_keys=["input", "mask"], + same_on_batch=False, + ) + + def test_two_seeds_give_two_different_flips(self): + volume, seg = self.tiny_volume(), self.tiny_seg() + pipeline = self._pipeline(flip_axis=(0, 1, 2)) + + outputs = [] + for seed in range(12): + torch.manual_seed(seed) + outputs.append(first_output(pipeline(volume.clone(), seg.clone())).clone()) + + distinct = {tuple(out.flatten()[:64].tolist()) for out in outputs} + self.assertGreater(len(distinct), 1, "every seed produced the same flip -- params['flip'] is being ignored") + + def test_batch_elements_flip_independently(self): + """The generator samples [B, 3]; the loop used to discard it and flip all of them.""" + torch.manual_seed(0) + volume = torch.rand(8, 1, 8, 8, 8) + seg = torch.zeros(8, 1, 8, 8, 8) + pipeline = self._pipeline(flip_axis=(0, 1, 2)) + + out = first_output(pipeline(volume.clone(), seg.clone())) + flipped = [bool(torch.allclose(out[b], torch.flip(volume[b], dims=(1, 2, 3)))) for b in range(8)] + + self.assertIn(False, flipped, "every batch element got the identical all-axis flip") + + def test_the_mask_is_flipped_the_same_way_as_the_image(self): + """Image and mask read the same params['flip'], so they cannot disagree.""" + torch.manual_seed(3) + volume = torch.rand(4, 1, 8, 8, 8) + seg = (volume > 0.5).float() + + result = self._pipeline(flip_axis=(0, 1, 2))(volume.clone(), seg.clone()) + image, mask = result[0], result[1] + + self.assertTrue(torch.equal((image > 0.5).float(), mask), "the mask was flipped differently from the image") + + def test_only_configured_axes_are_ever_flipped(self): + torch.manual_seed(1) + volume = torch.rand(6, 1, 8, 8, 8) + seg = torch.zeros(6, 1, 8, 8, 8) + + out = first_output(self._pipeline(flip_axis=(0,))(volume.clone(), seg.clone())) + for b in range(6): + unchanged = torch.allclose(out[b], volume[b]) + flipped_axis0 = torch.allclose(out[b], torch.flip(volume[b], dims=(1,))) + self.assertTrue(unchanged or flipped_axis0, f"batch element {b} was flipped along an axis that was not configured") + + def test_apply_transform_without_params_still_flips_every_configured_axis(self): + """The fallback for callers that reach past the generator into apply_transform.""" + transform = RandomFlipTransformGPU(p=1.0, flip_axis=(0, 1, 2)) + volume = torch.rand(2, 1, 8, 8, 8) + + out = transform.apply_transform(volume.clone(), {}, {}, transform=None) + + self.assertTrue(torch.equal(out, torch.flip(volume, dims=(2, 3, 4)))) + + +class TestSingleAxisIsRedrawnEveryCall(SmaugLabTestCase): + """The axis used to be chosen in make_samplers, which kornia calls once.""" + + def test_the_scale_generator_does_not_pin_one_axis_forever(self): + torch.manual_seed(0) + generator = ScaleGenerator3D(scale=(0.3, 1.0), one_dim=True) + generator.make_samplers(torch.device("cpu"), torch.float32) + + degraded_axes = set() + for _ in range(20): + scale = generator((4,), same_on_batch=False)["scale"] + # Exactly one axis per row is scaled; the rest sit at the neutral 1.0. + for row in scale: + self.assertEqual(int((row != 1.0).sum()), 1, "more than one axis was degraded") + degraded_axes.add(int((row != 1.0).nonzero().item())) + + self.assertGreater(len(degraded_axes), 1, "the same axis was degraded on every call") + + def test_same_on_batch_degrades_one_shared_axis(self): + torch.manual_seed(0) + generator = ScaleGenerator3D(scale=(0.3, 1.0), one_dim=True) + generator.make_samplers(torch.device("cpu"), torch.float32) + + scale = generator((5,), same_on_batch=True)["scale"] + chosen = {int((row != 1.0).nonzero().item()) for row in scale} + self.assertEqual(len(chosen), 1, "same_on_batch should pick one axis for the whole batch") + + def test_isotropic_scaling_is_left_alone(self): + """one_dim=False must keep every axis independent, as before.""" + torch.manual_seed(0) + generator = ScaleGenerator3D(scale=(0.3, 0.9), one_dim=False) + generator.make_samplers(torch.device("cpu"), torch.float32) + + scale = generator((4,), same_on_batch=False)["scale"] + self.assertTrue(bool((scale != 1.0).all()), "one_dim=False should not neutralise any axis") + + def test_the_crop_generator_uses_one_axis_for_crop_and_position(self): + """make_samplers drew a separate axis for each, so they could disagree.""" + torch.manual_seed(0) + generator = CropGenerator3D(crop=(0.5, 0.9), pos=(0.2, 0.8), one_dim=True) + generator.make_samplers(torch.device("cpu"), torch.float32) + + params = generator((4,), same_on_batch=False) + crop, pos = params["crop"], params["pos"] + + for b in range(4): + cropped = (crop[b] != 1.0).nonzero().flatten().tolist() + positioned = (pos[b] != 0.5).nonzero().flatten().tolist() + self.assertEqual(len(cropped), 1) + self.assertEqual(cropped, positioned, "the crop and its position were placed on different axes") + + def test_the_crop_generator_neutralises_position_at_the_centre(self): + """The old code copied the crop's neutral 1.0 onto `pos`, i.e. the far edge.""" + torch.manual_seed(0) + generator = CropGenerator3D(crop=(0.5, 0.9), pos=(0.2, 0.8), one_dim=True) + generator.make_samplers(torch.device("cpu"), torch.float32) + + params = generator((4,), same_on_batch=False) + crop, pos = params["crop"], params["pos"] + + for b in range(4): + kept = (crop[b] != 1.0).nonzero().flatten().tolist() + untouched = [axis for axis in range(3) if axis != kept[0]] + for axis in untouched: + self.assertAlmostEqual(float(pos[b, axis]), 0.5, places=5, msg="a non-cropped axis was not centred") + + +class TestMaskResampleRestore(SmaugLabTestCase): + """Guard tests for `RandomAffine3DCustom.apply_transform_mask`. + + Unlike the rest of this file these pass before the change too: `resample_method` + was annotated but assigned only inside the `if`, so the restore below it could + read an unbound local -- except that `apply_transform` indexes `flags["resample"]` + unguarded and raises `KeyError` first, which makes the unbound local unreachable + rather than harmless. `flags.get(...)` removes the hazard; these tests pin the + behaviour that has to survive it. + """ + + def _transform(self): + return RandomAffine3DCustom(p=1.0, degrees=5, align_corners=True) + + def _flags(self, resample: str = "bilinear"): + transform = self._transform() + flags = dict(transform.flags) + flags["resample"] = Resample.get(resample) + return transform, flags + + def test_a_mask_is_resampled_and_keeps_its_shape(self): + transform, flags = self._flags() + seg = self.tiny_seg() + params = transform.forward_parameters(seg.shape) + matrix = transform.compute_transformation(seg, params, flags) + + out = transform.apply_transform_mask(seg.clone(), params, flags, transform=matrix) + + self.assertEqual(out.shape, seg.shape) + + def test_the_callers_resample_mode_is_restored(self): + """The method flips the flag to "nearest" for the mask and must put it back.""" + transform, flags = self._flags() + seg = self.tiny_seg() + params = transform.forward_parameters(seg.shape) + matrix = transform.compute_transformation(seg, params, flags) + + transform.apply_transform_mask(seg.clone(), params, flags, transform=matrix) + + self.assertEqual(flags["resample"], Resample.get("bilinear"))