feat: add the main GPU normalization schemes and a random choice between them

README.md

@@ -187,6 +187,61 @@ 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.
+* 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.
+
 ### Paper configs
 
 Two shipped configs are the setups behind the paper's tables. `smauglab hash <config>`
@@ -231,8 +286,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/configs/all_augmentations.json

@@ -531,6 +531,91 @@
             "out_seg": 0.0,
             "p": 1.0,
             "p_batch": 1.0
+        },
+        "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": {

smauglab/registry.py

@@ -65,15 +65,23 @@ class Backend(str, Enum):
 class AugType(str, Enum):
     """Augmentation strength class.
 
-    Member names are verbatim from segtransferaug's `AUG2GROUP`, which this
+    GEO, GE and TA are verbatim from segtransferaug's `AUG2GROUP`, which this
     replaces, so downstream `AugType.TA` keeps resolving. The GPU random-order
     pipelines bucket transforms by this to decide what goes inside each
     `RandomChooseXTransformsGPU`.
+
+    NORM has no `AUG2GROUP` ancestor. It is the group of whole-volume intensity
+    standardizations, which are mutually exclusive by nature: chaining a min-max
+    rescale onto a z-score does not compose two augmentations, it discards the
+    first. So the builder bucket for this group picks exactly *one* of its members
+    per sample rather than running them all, which is a different question from the
+    "which order, how many" one GE and TA ask. See `build.build_gpu_pipeline`.
     """
 
     GEO = "GEO"  # geometric: applied in sequence, never bucketed
     GE = "GE"  # general enhancement, usually light
     TA = "TA"  # transfer augmentation, usually heavy
+    NORM = "NORM"  # intensity standardization: mutually exclusive, one per sample
 
 
 class AugId(str, Enum):
@@ -106,8 +114,14 @@ class AugId(str, Enum):
     CLAMP = "clamp"
     LOW_RES = "low_res"
     ACQ = "acq"
-    ZSCORE = "zscore"
     MIRROR = "mirror"
+    # -- intensity standardization, one concept each so key <-> class stays 1:1
+    ZSCORE = "zscore"
+    NORM_MINMAX = "norm_minmax"
+    NORM_PERCENTILE = "norm_percentile"
+    NORM_ROBUST_ZSCORE = "norm_robust_zscore"
+    NORM_MAX_SCALE = "norm_max_scale"
+    NORM_MEAN_SCALE = "norm_mean_scale"
     # -- transfer augmentation
     SCHARR = "scharr"
     LAPLACE = "laplace"
@@ -231,7 +245,14 @@ def __init__(self, source: str, problems: list[str]) -> None:
             "RandomAcqTransformGPU",  # AcqTransform
             "RandomCropTransformGPU",  # CropTransform
             "RandomBiasFieldGPU",  # BiasFieldTransform
+            # The normalization group, last: it standardizes whatever the intensity
+            # transforms above it did. Only one of these runs per sample -- see AugType.NORM.
             "ZscoreNormalizationGPU",  # ZscoreNormalizationTransform
+            "MinMaxNormalizationGPU",
+            "PercentileNormalizationGPU",
+            "RobustZscoreNormalizationGPU",
+            "MaxScaleNormalizationGPU",
+            "MeanScaleNormalizationGPU",
         ),
         # From the ladder in cpu/transforms.py, top to bottom.
         Backend.CPU: (

smauglab/transforms/__init__.py

@@ -19,6 +19,7 @@ class carries an `@register(...)` decorator, so the registry is complete once th
 from smauglab.transforms.gpu import contrast as gpu_contrast  # noqa: F401
 from smauglab.transforms.gpu import domain_transfer  # noqa: F401
 from smauglab.transforms.gpu import fromSeg as gpu_fromSeg
+from smauglab.transforms.gpu import normalization as gpu_normalization  # noqa: F401
 from smauglab.transforms.gpu import palette as gpu_palette  # noqa: F401
 from smauglab.transforms.gpu import spatial as gpu_spatial  # noqa: F401
 from smauglab.transforms.synthseg import transforms as synthseg_transforms  # noqa: F401

smauglab/transforms/build.py

@@ -76,25 +76,70 @@ def build_transforms(
     return sorted(built, key=lambda pair: registry.pipeline_position(pair[1]))
 
 
-def _buckets(built: list[tuple[Any, AugEntry]], *, hoist_sequential: bool) -> tuple[list, list, list]:
-    """Split into (sequential, transfer, general-enhancement).
+def _buckets(built: list[tuple[Any, AugEntry]], *, hoist_sequential: bool) -> tuple[list, list, list, list]:
+    """Split into (sequential, transfer, general-enhancement, normalization).
 
     Geometry always leads, in order. `force_sequential` means "never inside the GE
     bucket" and so only bites when there *is* one: with both groups bucketed,
     `RandomLowResTransformGPU` is hoisted up front alongside the geometry, but when
     only the transfer group is bucketed it stays in its ordinary GE position. That
     is exactly what the two hand-written pipelines did, and the difference is
     invisible unless you line their two ladders up side by side.
+
+    Normalization is its own list because it is the one group whose members are
+    alternatives rather than additions -- see `AugType.NORM` and `_normalization_stage`.
     """
-    sequential, transfer, enhancement = [], [], []
+    sequential, transfer, enhancement, normalization = [], [], [], []
     for transform, entry in built:
         if entry.group is AugType.GEO or (hoist_sequential and entry.force_sequential):
             sequential.append(transform)
         elif entry.group is AugType.TA:
             transfer.append(transform)
+        elif entry.group is AugType.NORM:
+            normalization.append(transform)
         else:
             enhancement.append(transform)
-    return sequential, transfer, enhancement
+    return sequential, transfer, enhancement, normalization
+
+
+def _normalization_stage(normalization: list, *, options: dict[str, Any]) -> list:
+    """Wrap the normalization group so that exactly one of its members runs.
+
+    Chaining normalizations does not compose them -- a min-max rescale on top of a
+    z-score simply discards the z-score -- so a config naming several of them is asking
+    for a *choice*, not for a sequence. `RandomChooseXTransformsGPU` with
+    `num_transforms=1` is that choice already: it draws one child uniformly and then
+    spends that child's own `p`, which makes each scheme's `p` its weight
+    (`P(scheme) = p / n`) and leaves the remainder as "no normalization at all".
+    `random_order=True` in every mode, including the sequential one: with one transform
+    to pick there is no order to shuffle, and `False` would pin the choice to whichever
+    scheme sorts first.
+
+    One draw per *batch*, not per sample, under the default `pipeline.same_on_batch`.
+    That setting defaults to True and the container forces it onto every stage it holds,
+    this bucket included, which is what makes a bucket take one decision for the whole
+    batch -- the same is true of the `ta` and `ge` buckets today. A config wanting each
+    sample to draw its own scheme sets `pipeline.same_on_batch: false`, which is also
+    what makes every transform's `p` per-sample; see `SmaugConfig.same_on_batch`.
+
+    Fewer than two members and nothing is wrapped. A single normalization keeps its
+    pipeline position verbatim -- wrapping it would move it, which `pipeline.order:
+    config` can notice, and would gate it behind a second probability for no gain.
+    """
+    if len(normalization) < 2:
+        return normalization
+
+    from smauglab.transforms.gpu.transforms_list import RandomChooseXTransformsGPU
+
+    return [
+        RandomChooseXTransformsGPU(
+            transforms_list=normalization,
+            num_transforms=1,
+            p=options.get("norm_probability", 1.0),
+            random_order=True,
+            label="norm",
+        )
+    ]
 
 
 def build_gpu_pipeline(
@@ -107,13 +152,20 @@ def build_gpu_pipeline(
 ) -> list[Any]:
     """The GPU transform list for a given pipeline mode."""
     built = build_transforms(section, Backend.GPU, source=source, order_source=order_source)
+    options = options or {}
     if mode is PipelineMode.SEQUENTIAL:
-        return [transform for transform, _ in built]
+        # Order is left exactly as `build_transforms` produced it, including under
+        # `pipeline.order: config`. The only rearrangement is the normalization choice,
+        # and only when there is a choice to make: a lone normalization is not touched.
+        normalization = [transform for transform, entry in built if entry.group is AugType.NORM]
+        if len(normalization) < 2:
+            return [transform for transform, _ in built]
+        kept = [transform for transform, entry in built if entry.group is not AugType.NORM]
+        return kept + _normalization_stage(normalization, options=options)
 
     from smauglab.transforms.gpu.transforms_list import RandomChooseXTransformsGPU
 
-    options = options or {}
-    sequential, transfer, enhancement = _buckets(built, hoist_sequential=mode is PipelineMode.RANDOM_ORDER)
+    sequential, transfer, enhancement, normalization = _buckets(built, hoist_sequential=mode is PipelineMode.RANDOM_ORDER)
 
     out = list(sequential)
     out.append(
@@ -139,6 +191,8 @@ def build_gpu_pipeline(
         # RANDOM_ORDER_TA shuffles only the transfer augmentations; the general
         # enhancements stay in pipeline order, as they always have.
         out.extend(enhancement)
+    # Last in every mode: it standardizes whatever everything above it did.
+    out.extend(_normalization_stage(normalization, options=options))
     return out
 
 

smauglab/transforms/cpu/contrast.py

@@ -366,7 +366,7 @@ def apply_filter(x: torch.Tensor, kernel: torch.Tensor, **kwargs) -> torch.Tenso
 @register(
     aug_id=AugId.ZSCORE,
     backend=Backend.CPU,
-    group=AugType.GE,
+    group=AugType.NORM,
 )
 class ZscoreNormalization(ImageOnlyTransform):
     """

smauglab/transforms/gpu/base.py

@@ -695,9 +695,10 @@ class name, which tensor it ran on, whether that tensor changed, and its scalar
       -- which is this wrapper, since the patch is an instance attribute. One
       application per record; the identical `params` on the pair is the proof that
       image and mask moved together.
-    * `RandomChooseXTransformsGPU` appears twice in `random_order` mode because the
-      builder emits two of them, the transfer bucket and the general-enhancement one.
-      `label` tells them apart.
+    * `RandomChooseXTransformsGPU` appears more than once because the builder emits one
+      per bucket: the transfer bucket, the general-enhancement one in `random_order`
+      mode, and the normalization choice whenever a config names more than one
+      normalization scheme. `label` tells them apart ("ta", "ge", "norm").
     """
     fired: list[dict] = []
     patched = []

smauglab/transforms/gpu/contrast.py

@@ -1819,7 +1819,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[
 @register(
     aug_id=AugId.ZSCORE,
     backend=Backend.GPU,
-    group=AugType.GE,
+    group=AugType.NORM,
 )
 class ZscoreNormalizationGPU(ImageOnlyTransform):
     """Apply z-score normalization to selected channels.