From e38293807d82598a7e7501e65149290ef0d62ba3 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 20 Aug 2026 08:00:40 +0000 Subject: [PATCH 1/2] fix: make the spatial transforms sample the way they claim to Three defects in gpu/spatial.py, all silent -- nothing crashes and no test fails, the pipeline just does something other than what the config asked for. * RandomFlipTransformGPU never read the flip flags its own generator sampled. The loop recomputed the same `flip_axis`-derived list for every batch element, so it flipped all configured axes, identically, on every call: three seeded calls gave byte-identical output, and FlipGenerator3D -- including its "at least one axis" guarantee -- was dead code. It now reads params["flip"], falling back to the old all-configured-axes behaviour only for callers that reach into apply_transform directly with no sampled params. * The single-axis generators drew their "random" axis in make_samplers, which kornia calls once and caches. The same axis was therefore degraded for a whole training run. The draw moves to forward(), via _choose_axis/_keep_one_axis. * CropGenerator3D additionally drew separate axes for the crop and for its position, so the crop could be taken along one axis while the position placing it was randomised along another. It also neutralised the position to 1.0 using the crop's neutral value; the position is the crop centre as a fraction of the axis, so its neutral value is 0.5 (centred), not the far edge. Also `flags.get("resample")` instead of `"resample" in flags` plus a bare assignment in apply_transform_mask: `resample_method` was annotated but only assigned inside the `if`, so the restore below could read an unbound local. It is unreachable today -- apply_transform indexes flags["resample"] and raises KeyError first -- so this is removing a hazard, not fixing an observed failure, and its two tests are guards that pass either way. unit_tests/test_spatial_sampling.py covers all of it; the five regression tests fail against the previous implementation. Models trained before this change saw the old behaviour and will not reproduce against it. No config key, parameter or default changed. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/gpu/spatial.py | 85 +++++++----- unit_tests/test_spatial_sampling.py | 193 ++++++++++++++++++++++++++++ 2 files changed, 247 insertions(+), 31 deletions(-) create mode 100644 unit_tests/test_spatial_sampling.py 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")) From dcd133dadea8e4ca5f8a8dc06def3bfd73bd36ce Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 20 Aug 2026 11:16:51 +0000 Subject: [PATCH 2/2] ci: run the mypy job on 3.10, the version it type-checks against The typecheck job started failing on an error in numpy's own stub: numpy/__init__.pyi:737: error: Type statement is only supported in Python 3.12 and greater [syntax] Found 1 error in 1 file (errors prevented further checking) mypy applies its target version when parsing *every* file, third-party stubs included. [tool.mypy] deliberately sets python_version = "3.10" -- the floor of the range `requires-python` promises, rather than whatever CI happens to run -- but the job itself ran on 3.12. pip therefore resolved numpy 2.5, which requires Python >=3.12 and writes PEP 695 `type` statements in its stubs, and mypy rejected them as too new for the declared target. Nothing in smauglab/ was ever reached: the run aborted at the parse error. Running the job on 3.10 makes the interpreter and the target agree, which is what the config comment already said was intended. It also resolves numpy 2.2.x, the newest release that supports 3.10, whose stubs parse under the target. The test job has been installing this same dependency set on 3.10 all along, so the install path is already exercised. Unrelated to the augmentation fix in this branch -- a dependency released since the last run on main, which was green on 2026-08-10. It is here because the branch is red without it; happy to split it out if you would rather review it on its own. Not reproducible locally: numpy 2.5 cannot be installed below 3.12. Verified instead that numpy 2.2.6 (the 3.10 resolution) type-checks clean against python_version = "3.10", and that numpy 2.3.4 does not yet carry the offending syntax. Co-Authored-By: Claude Opus 5 --- .github/workflows/lint.yml | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) 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