From 0369fc8115bacfd5d97aef59acd8081cdd9f95e7 Mon Sep 17 00:00:00 2001 From: iback Date: Fri, 13 Nov 2026 14:25:51 +0000 Subject: [PATCH] feat: make RandomRedistributeSegGPU's bump-width floor configurable The remap adds a signed Gaussian bump per label at that label's mean intensity and rescales the sum to a peak-to-peak of 2 in the [0, 1] min-max units it works in. Each bump's width is floored at a hardcoded 0.01. A bump of width s rescaled that way has a slope of about 1.2/s, so at the floor a thin band of intensities swings by more than the whole input range and voxel noise inside the band comes out as salt-and-pepper. Measured on a ramp with the floor binding: maximum slope 123 / 26 / 13 at floors 0.01 / 0.05 / 0.1. On real spine patches the high-frequency share of the image at least doubled in 24% of applications. `min_redist_std` exposes the floor. Default 0.01 is the historical value. Co-Authored-By: Claude Opus 5.5 --- smauglab/configs/all_augmentations.json | 1 + smauglab/transforms/gpu/fromSeg.py | 13 ++++- unit_tests/test_redistribute_std_floor.py | 67 +++++++++++++++++++++++ 3 files changed, 80 insertions(+), 1 deletion(-) create mode 100644 unit_tests/test_redistribute_std_floor.py diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json index 8be9b9b..bfadb91 100644 --- a/smauglab/configs/all_augmentations.json +++ b/smauglab/configs/all_augmentations.json @@ -218,6 +218,7 @@ ], "in_seg": 0.2, "keepdim": true, + "min_redist_std": 0.01, "p": 1.0, "p_batch": 1.0, "retain_stats": true, diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index a53155a..8002437 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -265,6 +265,13 @@ class RandomRedistributeSegGPU(ImageOnlyTransform): only changes which setting you get by default. `transform_params_hybrid.json` and `transform_params_hybrid_TAGE.json` ask for False explicitly and are unaffected. + + `min_redist_std` floors the width of each label's Gaussian bump, in the `[0, 1]` + min-max units the remap is built in. The bumps are rescaled to a peak-to-peak of 2, + so a narrow one maps a thin band of intensities onto a swing larger than the whole + input range, and voxel noise inside that band comes out as salt-and-pepper. The + default 0.01 is the historical floor; 0.05 keeps the remap's slope at about a fifth + of that worst case. """ def __init__( @@ -278,13 +285,17 @@ def __init__( keepdim: bool = True, std_noise_range: Sequence[float] = (0.1, 0.3), dilation_iterations_range: Sequence[int] = (1, 3), + min_redist_std: float = 0.01, ) -> None: super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) + if min_redist_std <= 0: + raise ValueError(f"min_redist_std must be positive, got {min_redist_std}") self.in_seg = in_seg self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats self.std_noise_range = std_noise_range self.dilation_iterations_range = dilation_iterations_range + self.min_redist_std = min_redist_std @torch.no_grad() def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -393,7 +404,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ ) redist_std = torch.maximum( torch.rand(R, device=input.device) * std_noise_range + 0.4 * torch.abs((means - dil_means) * stds / (dil_stds + 1e-6)), - torch.full((R,), 0.01, device=input.device, dtype=input.dtype), + torch.full((R,), self.min_redist_std, device=input.device, dtype=input.dtype), ) # Accumulated in place rather than stacked: the global branch built a list diff --git a/unit_tests/test_redistribute_std_floor.py b/unit_tests/test_redistribute_std_floor.py new file mode 100644 index 0000000..2cb94e2 --- /dev/null +++ b/unit_tests/test_redistribute_std_floor.py @@ -0,0 +1,67 @@ +"""`RandomRedistributeSegGPU(min_redist_std=...)` bounds how steep the remap gets. + +The remap adds a signed Gaussian bump per label, centred on that label's mean intensity, +and rescales the sum to a peak-to-peak of 2 in the `[0, 1]` min-max units it works in. +Each bump's width is floored at `min_redist_std`. A bump of width s, rescaled that way, +has a slope of about 1.2 / s, so the historical floor of 0.01 lets a thin band of +intensities swing by more than the whole input range; voxel noise inside the band then +comes out as salt-and-pepper. The floor binds whenever a label's mean is close to its +dilation ring's and the random part of the width is drawn small. +""" + +from __future__ import annotations + +import torch + +from smauglab.registry import Backend, accepted_params, get +from smauglab.transforms.gpu.fromSeg import RandomRedistributeSegGPU +from unit_tests.helpers import SmaugLabTestCase + +N = 16 + + +def inputs() -> tuple[torch.Tensor, torch.Tensor]: + """An intensity ramp and a voxel-random label map. + + Random labels make every region's dilation ring look like the region itself, so the + data-driven part of the width is about zero and the floor is what sets it. + """ + torch.manual_seed(0) + image = torch.linspace(0, 1, N**3).view(1, 1, N, N, N) + labels = torch.randint(0, 3, (1, 1, N, N, N)).float() + return image, labels + + +def max_slope(floor: float, seed: int) -> float: + image, labels = inputs() + torch.manual_seed(seed) + transform = RandomRedistributeSegGPU(p=1.0, in_seg=0.0, retain_stats=False, std_noise_range=(0.0, 0.0), min_redist_std=floor) + out = transform.apply_transform(image.clone(), {"seg": labels}, transform.flags).flatten() + x = image.flatten() + return float(((out[1:] - out[:-1]) / (x[1:] - x[:-1])).abs().max()) + + +class TestRedistributeStdFloor(SmaugLabTestCase): + def test_the_slope_follows_the_floor(self): + for floor in (0.01, 0.05, 0.1): + with self.subTest(floor=floor): + slopes = [max_slope(floor, seed) for seed in range(10)] + # 1 from the identity plus about 1.2 / floor from the steepest bump. + self.assertLess(max(slopes), 1.5 * (1 + 1.2 / floor)) + self.assertGreater(max(slopes), 0.5 * (1 + 1.2 / floor)) + + def test_the_default_is_the_historical_floor(self): + image, labels = inputs() + outs = [] + for kwargs in ({}, {"min_redist_std": 0.01}): + torch.manual_seed(1) + transform = RandomRedistributeSegGPU(p=1.0, **kwargs) + outs.append(transform.apply_transform(image.clone(), {"seg": labels}, transform.flags)) + self.assertTrue(torch.equal(outs[0], outs[1])) + + def test_a_non_positive_floor_is_refused(self): + with self.assertRaises(ValueError): + RandomRedistributeSegGPU(min_redist_std=0.0) + + def test_the_config_surface_accepts_it(self): + self.assertIn("min_redist_std", accepted_params(get("RandomRedistributeSegGPU", Backend.GPU)))