feat: make RandomRedistributeSegGPU's bump-width floor configurable

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,

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

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)))