diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json index 8be9b9b..fab8399 100644 --- a/smauglab/configs/all_augmentations.json +++ b/smauglab/configs/all_augmentations.json @@ -318,6 +318,7 @@ "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, + "relative": false, "same_on_batch": false, "std": 1.0 }, diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index ef6c3c8..3802b6a 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -718,7 +718,15 @@ class RandomGaussianNoiseGPU(ImageOnlyTransform): Args: mean (float): Mean of the Gaussian noise. Default is 0.0. - std (float): Standard deviation of the Gaussian noise. Default is 0.1. + std (float): Upper bound of the noise standard deviation, which is drawn from + `U(0, std)` per sample (or once, with `same_on_batch`). Default is 1.0. + relative (bool): Read `std` as a fraction of each sample's own standard deviation + in that channel rather than in absolute intensity units. The absolute form + makes the same config a different augmentation on every intensity scale: a + z-scored MR patch (std ~0.9), a CT patch under nnU-Net's CT normalisation + (std ~0.17) and a `[0, 1]` SynthSeg synthesis (std ~0.18) see noise at the + same absolute std, so the last two are mostly noise. `mean` stays absolute. + Default False. same_on_batch (bool): Apply the same transformation across the batch. Default is False. p (float): Probability of applying the transform. Default is 1.0. keepdim (bool): Whether to keep the number of dimensions. Default is False. @@ -731,6 +739,7 @@ def __init__( self, mean: float = 0.0, std: float = 1.0, + relative: bool = False, apply_to_channel: Sequence[int] = (0,), # Apply to first channel by default same_on_batch: bool = False, in_seg: float = 0.0, @@ -744,6 +753,7 @@ def __init__( self.apply_to_channel = apply_to_channel self.mean = mean self.std = std + self.relative = relative self.in_seg = in_seg self.out_seg = out_seg self.mix_in_out = mix_in_out @@ -757,13 +767,16 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ if self.same_on_batch: std = torch.rand(1, device=input.device, dtype=input.dtype) * self.std noise = torch.randn_like(input[:, c], device=input.device, dtype=input.dtype) - noise = noise * std + self.mean else: std = torch.rand(input.shape[0], device=input.device, dtype=input.dtype) * self.std noise = torch.randn_like(input[:, c], device=input.device, dtype=input.dtype) - # Broadcast the per-sample std instead of writing one row at a time: - # same draws, same arithmetic, one kernel instead of N. - noise = noise * std.view(-1, *([1] * (noise.dim() - 1))) + self.mean + if self.relative: + # Scaled after both draws, so the random stream is the absolute form's. + std = std * input[:, c].reshape(input.shape[0], -1).std(dim=1) + # Broadcast the per-sample std instead of writing one row at a time: + # same draws, same arithmetic, one kernel instead of N. A one-element std + # (`same_on_batch`, absolute) broadcasts the same way it did unreshaped. + noise = noise * std.view(-1, *([1] * (noise.dim() - 1))) + self.mean orig = input[:, c] x = orig + noise diff --git a/unit_tests/test_gaussian_noise_relative.py b/unit_tests/test_gaussian_noise_relative.py new file mode 100644 index 0000000..cd89b64 --- /dev/null +++ b/unit_tests/test_gaussian_noise_relative.py @@ -0,0 +1,71 @@ +"""`RandomGaussianNoiseGPU(relative=True)` scales the noise with the image. + +The std is drawn from `U(0, std)` in absolute intensity units, so the same config is a +different augmentation on every intensity scale: a z-scored MR patch (std ~0.9), a CT +patch under nnU-Net's CT normalisation (std ~0.17) and a `[0, 1]` SynthSeg synthesis +(std ~0.18) all get noise of the same absolute std. `relative=True` reads `std` as a +fraction of each sample's own std, which makes the transform scale-equivariant. +""" + +from __future__ import annotations + +import torch + +from smauglab.registry import Backend, accepted_params, get +from smauglab.transforms.gpu.contrast import RandomGaussianNoiseGPU +from unit_tests.helpers import SmaugLabTestCase + +SHAPE = (4, 1, 12, 12, 12) + + +def noised(x: torch.Tensor, relative: bool, same_on_batch: bool, seed: int = 0) -> torch.Tensor: + transform = RandomGaussianNoiseGPU(std=0.3, relative=relative, same_on_batch=same_on_batch, p=1.0) + torch.manual_seed(seed) + return transform(x.clone()) + + +def image() -> torch.Tensor: + torch.manual_seed(123) + # Four samples on four different intensity scales. + return torch.randn(*SHAPE) * torch.tensor([0.17, 0.9, 2.5, 30.0]).view(-1, 1, 1, 1, 1) + + +class TestRelativeGaussianNoise(SmaugLabTestCase): + def test_the_absolute_form_is_unchanged(self): + """Same draws in the same order, same arithmetic as before the option existed.""" + x = image() + for same_on_batch in (False, True): + with self.subTest(same_on_batch=same_on_batch): + torch.manual_seed(0) + n = 1 if same_on_batch else SHAPE[0] + std = torch.rand(n) * 0.3 + noise = torch.randn_like(x[:, 0]) + scale = std if same_on_batch else std.view(-1, 1, 1, 1) + expected = x[:, 0] + (noise * scale + 0.0) + out = noised(x, relative=False, same_on_batch=same_on_batch) + self.assertTrue(torch.equal(out[:, 0], expected)) + + def test_relative_noise_is_scale_equivariant(self): + x = image() + for same_on_batch in (False, True): + with self.subTest(same_on_batch=same_on_batch): + out = noised(x, relative=True, same_on_batch=same_on_batch) + out_scaled = noised(7.0 * x, relative=True, same_on_batch=same_on_batch) + torch.testing.assert_close(out_scaled, 7.0 * out, rtol=1e-5, atol=1e-5) + + def test_relative_is_the_absolute_draw_times_the_sample_std(self): + """No extra random draw: the noise field is the absolute one, rescaled per sample. + + In float64, because the noise is read back as `out - x`, which in float32 loses + about ulp(x) -- 2e-6 at the std-30 sample, then multiplied by 30. + """ + x = image().double() + sample_std = x.reshape(SHAPE[0], -1).std(dim=1).view(-1, 1, 1, 1, 1) + for same_on_batch in (False, True): + with self.subTest(same_on_batch=same_on_batch): + absolute = noised(x, relative=False, same_on_batch=same_on_batch) - x + relative = noised(x, relative=True, same_on_batch=same_on_batch) - x + torch.testing.assert_close(relative, absolute * sample_std, rtol=1e-4, atol=1e-5) + + def test_the_config_surface_accepts_it(self): + self.assertIn("relative", accepted_params(get("RandomGaussianNoiseGPU", Backend.GPU)))