feat: restore a jittered intensity profile around the whole GPU pipeline

.gitignore

@@ -50,6 +50,9 @@ MANIFEST
 # configs beside it are tracked only because they predate the blanket rule above;
 # a new one has to say so.
 !smauglab/configs/transform_params_paper-synthseg.json
+# The statistics bracket's reference config (see "Statistics restore" in README.md),
+# shipped so `unit_tests/test_configs.py` builds and runs it like any other.
+!smauglab/configs/transform_params_gpu_stats-restore.json
 *.yaml
 *.nii
 *.nii.gz

README.md

@@ -187,6 +187,129 @@ Scripts developped in this repository use JSON files to specify image and segmen
 
 To track parameters used during data augmentation, JSON files are also used: see this [example](https://github.com/neuropoly/SmaugLab/blob/16653a84e031c40e25a72e946c2724494606b21c/smauglab/configs/transform_params.json)
 
+### Random normalization
+
+The `NORM` group in the table below -- z-score, min-max, the CT-style percentile clip, a
+robust median/IQR scaling, divide-by-max, divide-by-mean -- is the intensity
+standardization step. Its members are alternatives, not additions: a min-max rescale on
+top of a z-score does not apply two augmentations, it discards the first. So a config
+naming more than one of them does not get a sequence. The builder puts them in a single
+stage at the end of the pipeline that picks **one scheme per draw**, which makes the
+choice of normalization itself an augmentation.
+
+```json
+{
+  "GPU": {
+    "ZscoreNormalizationGPU": { "p": 1.0 },
+    "MinMaxNormalizationGPU": { "out_range": [0.0, 1.0], "p": 1.0 },
+    "PercentileNormalizationGPU": { "percentiles": [0.005, 0.995], "p": 1.0 },
+    "RobustZscoreNormalizationGPU": { "spread": "iqr", "p": 1.0 },
+    "MaxScaleNormalizationGPU": { "p": 0.5 }
+  },
+  "pipeline": { "random_choose": { "norm_probability": 1.0 } }
+}
+```
+
+Inside that stage a scheme's `p` is its **weight**, not its frequency: one scheme is drawn
+uniformly from the `n` named, and its own `p` is then spent, so `P(scheme) = p / n`. The
+config above gives each of the first four a 1-in-5 share and `MaxScaleNormalizationGPU` a
+1-in-10 one, and in the remaining 1-in-10 of samples nothing is applied -- which is how
+"leave it at the preprocessing normalization" is expressed. There is deliberately no
+no-op transform to name; a scheme at `p: 0.0` is a share of the choice that does nothing.
+
+| key | default | what it does |
+| --- | --- | --- |
+| `stats_from` | `"all"` | Which voxels the statistics are read from: `"all"`, `"seg"` (the labelled anatomy) or `"nonzero"` (everything preprocessing did not zero out, the closest thing here to nnU-Net's `use_mask_for_norm`). It changes what is *measured*, never what is written -- `in_seg` / `out_seg` still decide that. |
+| `min_scale` | `0.001` | Leave a sample alone when the divisor is below this fraction of its standard deviation. Without it, `MeanScaleNormalizationGPU` on an already-z-scored patch divides rounding noise by a mean of ~0. |
+| `norm_probability` | `1.0` | Probability that the stage runs at all, under `pipeline.random_choose`. |
+
+A draw is one per **batch** at the default `pipeline.same_on_batch: true`, which the
+pipeline forces onto every stage it holds -- the `ta` and `ge` buckets behave the same
+way. Set `pipeline.same_on_batch: false` and each sample draws its own scheme, which is
+also what makes every transform's `p` per-sample rather than batch-wide.
+
+Three things to know before turning this on:
+
+* The patch has already been normalized once, by nnU-Net's preprocessing, and inference
+  re-applies only that fixed preprocessing. Randomizing the final standardization is
+  therefore a deliberate train/test difference in input scale, on the same logic as every
+  other appearance augmentation here: the network is being asked to stop depending on one
+  scale. There is no compensating re-normalization at the end -- if that is what you want,
+  it is the statistics bracket below, which runs *before* this stage.
+* The schemes act on preprocessed intensities, not on HU. A CT window in Hounsfield units
+  means nothing after preprocessing unless its bounds are stated, which is why
+  `PercentileNormalizationGPU` takes an explicit `clip_range`, `mean` and `std` -- given
+  all three it reproduces nnU-Net's `CTNormalization` exactly, and
+  `MinMaxNormalizationGPU` at its defaults reproduces `RescaleTo01Normalization`.
+  `unit_tests/test_normalization.py` pins both against the installed nnU-Net.
+
+### Statistics restore (`pipeline.retain_stats`)
+
+The per-transform `retain_stats` flag puts a channel's mean and standard deviation back
+after one transform. This is the same idea for the whole pipeline, with two more
+statistics and a jitter: the profile is snapshotted once after the geometry and restored
+once at the end, so the augmentations are free to do what they like in between and the
+network still sees an intensity profile that is plausible -- but never exactly the one the
+patch arrived with.
+
+```json
+{
+  "pipeline": {
+    "retain_stats": {
+      "stats": ["mean", "std", "p1", "p99"],
+      "percentiles": [1.0, 99.0],
+      "std_scale": [0.9, 1.1],
+      "percentile_scale": [0.9, 1.1],
+      "stats_from": "all",
+      "p": 1.0
+    }
+  }
+}
+```
+
+An absent block is off and spends no random numbers. A block that is present but empty is
+on with every default, which restores the snapshot exactly.
+
+**Two of the four, drawn per batch.** The snapshot is the per-sample, per-channel mean,
+standard deviation, 1st and 99th percentile. An affine map `y = ax + b` has two degrees of
+freedom, so all four cannot be restored at once -- each pass therefore draws **two** of
+them and solves for the map that puts exactly those two on target, leaving the other two
+to fall where the augmentations left them. The six pairs span the two useful extremes:
+
+| drawn pair | what it pins | what it is |
+| --- | --- | --- |
+| mean, std | the moments | what per-transform `retain_stats` does |
+| p1, p99 | the robust dynamic range | immune to the outliers an edge filter or a random convolution manufactures, which are exactly what make the standard deviation a bad target |
+| mean, p1 / mean, p99 / std, p1 / std, p99 | one of each | the mixed cases in between |
+
+`stats` is the pool the pair is drawn from, so `["mean", "std"]` reproduces the old
+behaviour at pipeline level and the default draws from all six pairs.
+
+**The jitter.** `std_scale` multiplies the standard-deviation target; `percentile_scale`
+multiplies each percentile's distance from the mean, independently, so the span can widen,
+narrow or go asymmetric. Both are a scalar or a `[lo, hi]` range, and `1.0` restores the
+snapshot exactly. The mean target is never jittered -- it is the anchor every pair shares,
+and moving it as well would put two unrelated augmentations behind one knob. `p` is the
+per-sample probability that the restore applies at all.
+
+**Where it sits.** With the block on, the pipeline is
+
+```
+[geometry] [snapshot] [augmentations] [restore] [crop] [normalization]
+```
+
+in all three pipeline modes. Two groups move to the tail: the crop, because it occludes a
+random box with the volume's own minimum rather than changing intensities, so a restore
+measured before it would be undoing the occlusion; and the normalization stage, because
+deciding the scale the network sees is its entire job. In `random_order` mode the snapshot
+lands after the geometry but *before* the hoisted `RandomLowResTransformGPU`, which is a
+general enhancement that merely runs early.
+
+A config should pick one mechanism or the other: the two do not compose, and a transform
+that restores its own moments inside the bracket is doing work the restore will redo.
+`smauglab/configs/transform_params_gpu_stats-restore.json` is the paper's `Ours` setup with
+the bracket on and every per-transform `retain_stats` turned off.
+
 ### Paper configs
 
 Two shipped configs are the setups behind the paper's tables. `smauglab hash <config>`
@@ -231,8 +354,13 @@ column -- see "Run Monai training with SmaugLab augmentations" above. Regenerate
 | clamp | GE | `RandomClampGPU` | — |
 | low_res | GE | `RandomLowResTransformGPU` | `SimulateLowResolutionTransform` |
 | acq | GE | `RandomAcqTransformGPU` | — |
-| zscore | GE | `ZscoreNormalizationGPU` | `ZscoreNormalization` |
 | mirror | GEO | — | `MirrorTransform` |
+| zscore | NORM | `ZscoreNormalizationGPU` | `ZscoreNormalization` |
+| norm_minmax | NORM | `MinMaxNormalizationGPU` | — |
+| norm_percentile | NORM | `PercentileNormalizationGPU` | — |
+| norm_robust_zscore | NORM | `RobustZscoreNormalizationGPU` | — |
+| norm_max_scale | NORM | `MaxScaleNormalizationGPU` | — |
+| norm_mean_scale | NORM | `MeanScaleNormalizationGPU` | — |
 | scharr | TA | `RandomScharrGPU` | `ScharrConvTransform` |
 | laplace | TA | `RandomLaplaceGPU` | `LaplaceConvTransform` |
 | unsharp_mask | TA | `RandomUnsharpMaskGPU` | — |

smauglab/cli.py

@@ -232,7 +232,7 @@ def build_parser() -> argparse.ArgumentParser:
 
     p = sub.add_parser("list", help="registered augmentations, in pipeline order")
     p.add_argument("--backend", choices=["gpu", "cpu"])
-    p.add_argument("--group", choices=["geo", "ge", "ta"])
+    p.add_argument("--group", choices=["geo", "ge", "ta", "norm"])
     p.add_argument("--json", action="store_true")
     p.set_defaults(func=cmd_list)
 

smauglab/config.py

@@ -6,7 +6,7 @@
       "_comment": "...",              // any key starting with '_' is ignored
       "GPU":      { "<ClassName>": { <constructor kwargs> }, ... },
       "CPU":      { ... },
-      "pipeline": { "random_choose": { ... } }
+      "pipeline": { "random_choose": { ... }, "retain_stats": { ... } }
     }
 
 A key is a class name, exactly, and a parameter is a constructor argument, exactly. Both are
@@ -96,7 +96,7 @@ def validate_section(section: dict, backend: Backend, *, source: str = "<config>
 RESERVED_SECTIONS = ("pipeline",)
 
 #: Keys the `pipeline` section may hold.
-PIPELINE_KEYS = ("mode", "order", "random_choose", "same_on_batch")
+PIPELINE_KEYS = ("mode", "order", "random_choose", "retain_stats", "same_on_batch")
 
 #: What to do about a config that has no backend section. Quoted by both errors that
 #: can report it, so the two cannot drift apart.
@@ -137,6 +137,22 @@ def pipeline_options(self, name: str) -> dict[str, Any]:
         """
         return copy.deepcopy(self.payload.get("pipeline", {}).get(name, {}))
 
+    def retain_stats_spec(self) -> Any:
+        """The parsed `pipeline.retain_stats` block, or None when there is none.
+
+        Returns a `gpu.stats.StatsSpec`, imported here rather than at module scope: that
+        module reaches kornia, and a CPU-only caller validating a config should not have
+        to. Typed loosely for the same reason -- `build_gpu_pipeline` has the real
+        annotation under `TYPE_CHECKING`.
+        """
+        from smauglab.transforms.gpu.stats import parse_spec
+
+        # Read off the payload rather than through `pipeline_options`, which cannot tell
+        # an absent block from an empty one: it always hands back a dict, and here the
+        # difference is the difference between off and on-with-defaults.
+        block = self.payload.get("pipeline", {}).get("retain_stats")
+        return parse_spec(copy.deepcopy(block))
+
     def pipeline_mode(self) -> PipelineMode:
         """How the GPU pipeline arranges its transforms.
 
@@ -239,6 +255,14 @@ def _pipeline_problems(self) -> list[str]:
         if same_on_batch is not None and not isinstance(same_on_batch, bool):
             problems.append(f"pipeline.same_on_batch: expected true or false, got {same_on_batch!r}")
 
+        retain_stats = pipeline.get("retain_stats")
+        if retain_stats is not None:
+            # Imported inside the check so a config with no statistics bracket still
+            # validates without importing the GPU transform stack.
+            from smauglab.transforms.gpu.stats import spec_problems
+
+            problems.extend(spec_problems(retain_stats))
+
         mode = pipeline.get("mode")
         if mode is not None:
             valid = [m.value for m in PipelineMode]

smauglab/configs/all_augmentations.json

@@ -803,6 +803,91 @@
             "seg_feather": 0.0,
             "seg_unit": "whole",
             "seg_unit_prob": 0.5
+        },
+        "MinMaxNormalizationGPU": {
+            "apply_to_channel": [
+                0
+            ],
+            "clip_percentiles": null,
+            "in_seg": 0.0,
+            "keepdim": true,
+            "min_scale": 0.001,
+            "nonzero_threshold": 1e-06,
+            "out_range": [
+                0.0,
+                1.0
+            ],
+            "out_seg": 0.0,
+            "p": 1.0,
+            "p_batch": 1.0,
+            "stats_from": "all"
+        },
+        "PercentileNormalizationGPU": {
+            "apply_to_channel": [
+                0
+            ],
+            "clip_range": null,
+            "in_seg": 0.0,
+            "keepdim": true,
+            "mean": null,
+            "min_scale": 0.001,
+            "nonzero_threshold": 1e-06,
+            "out_seg": 0.0,
+            "p": 1.0,
+            "p_batch": 1.0,
+            "percentiles": [
+                0.005,
+                0.995
+            ],
+            "stats_from": "all",
+            "std": null
+        },
+        "RobustZscoreNormalizationGPU": {
+            "apply_to_channel": [
+                0
+            ],
+            "in_seg": 0.0,
+            "keepdim": true,
+            "min_scale": 0.001,
+            "nonzero_threshold": 1e-06,
+            "out_seg": 0.0,
+            "p": 1.0,
+            "p_batch": 1.0,
+            "percentiles": [
+                0.25,
+                0.75
+            ],
+            "spread": "iqr",
+            "stats_from": "all",
+            "unit_variance": true
+        },
+        "MaxScaleNormalizationGPU": {
+            "apply_to_channel": [
+                0
+            ],
+            "in_seg": 0.0,
+            "keepdim": true,
+            "min_scale": 0.001,
+            "nonzero_threshold": 1e-06,
+            "out_seg": 0.0,
+            "p": 1.0,
+            "p_batch": 1.0,
+            "percentile": null,
+            "stats_from": "all",
+            "use_abs": true
+        },
+        "MeanScaleNormalizationGPU": {
+            "apply_to_channel": [
+                0
+            ],
+            "in_seg": 0.0,
+            "keepdim": true,
+            "min_scale": 0.001,
+            "nonzero_threshold": 1e-06,
+            "out_seg": 0.0,
+            "p": 1.0,
+            "p_batch": 1.0,
+            "stats_from": "nonzero"
         }
     },
     "CPU": {