From 92d8574426b3be5570cadd2d217983863bf040e5 Mon Sep 17 00:00:00 2001 From: iback Date: Tue, 13 Oct 2026 08:16:02 +0000 Subject: [PATCH] feat: add the main GPU normalization schemes and a random choice between them ZscoreNormalizationGPU was the only normalization a GPU config could reach. The other schemes a medical imaging pipeline uses -- min-max, the CT-style percentile clip, a robust median/IQR scaling, divide-by-max, divide-by-mean -- existed only in nnU-Net's preprocessing, which runs once per dataset. All five land in a new gpu/normalization.py behind a shared base that owns the clone, the channel guard, the statistics mask, both degeneracy guards and the region selection, so each scheme is an affine map plus an optional clip. At their defaults MinMaxNormalizationGPU reproduces nnU-Net's RescaleTo01Normalization and PercentileNormalizationGPU, given a clip_range, mean and std, reproduces CTNormalization; test_normalization.py pins both against the installed nnU-Net. The group is new: AugType.NORM, because its members are alternatives rather than additions -- a min-max rescale on top of a z-score discards the z-score. The builder therefore puts them in a RandomChooseXTransformsGPU with num_transforms=1 as the last stage in every pipeline mode, which makes each scheme's p its weight (P(scheme) = p / n) and leaves the remainder as no normalization at all. A config naming a single normalization is untouched, so no shipped pipeline changes shape. ZscoreNormalizationGPU moves from GE to NORM, which is what puts it in the same choice as the rest. In random_order mode that takes it out of the shuffled GE bucket and makes it a trailing stage; every shipped config sets its p to 0.0, so nothing fires differently, but the GE bucket no longer spends a draw on it and random_order runs will not reproduce earlier volumes bit for bit under one seed. Co-Authored-By: Claude Opus 5 --- README.md | 62 ++- smauglab/cli.py | 2 +- smauglab/configs/all_augmentations.json | 85 ++++ smauglab/registry.py | 25 +- smauglab/transforms/__init__.py | 1 + smauglab/transforms/build.py | 68 ++- smauglab/transforms/cpu/contrast.py | 2 +- smauglab/transforms/gpu/base.py | 7 +- smauglab/transforms/gpu/contrast.py | 2 +- smauglab/transforms/gpu/normalization.py | 569 +++++++++++++++++++++++ unit_tests/test_channel_guard.py | 17 +- unit_tests/test_normalization.py | 449 ++++++++++++++++++ unit_tests/test_transforms_gpu.py | 1 + 13 files changed, 1273 insertions(+), 17 deletions(-) create mode 100644 smauglab/transforms/gpu/normalization.py create mode 100644 unit_tests/test_normalization.py diff --git a/README.md b/README.md index 4755ed2..5a67c16 100644 --- a/README.md +++ b/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 ` @@ -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` | — | diff --git a/smauglab/cli.py b/smauglab/cli.py index 0cf470d..637fa8f 100644 --- a/smauglab/cli.py +++ b/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) diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json index 8be9b9b..eb351ed 100644 --- a/smauglab/configs/all_augmentations.json +++ b/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": { diff --git a/smauglab/registry.py b/smauglab/registry.py index 36805bc..48604e3 100644 --- a/smauglab/registry.py +++ b/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: ( diff --git a/smauglab/transforms/__init__.py b/smauglab/transforms/__init__.py index f3a9ece..d97887e 100644 --- a/smauglab/transforms/__init__.py +++ b/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 diff --git a/smauglab/transforms/build.py b/smauglab/transforms/build.py index 63fe621..d0a906e 100644 --- a/smauglab/transforms/build.py +++ b/smauglab/transforms/build.py @@ -76,8 +76,8 @@ 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, @@ -85,16 +85,61 @@ def _buckets(built: list[tuple[Any, AugEntry]], *, hoist_sequential: bool) -> tu 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 diff --git a/smauglab/transforms/cpu/contrast.py b/smauglab/transforms/cpu/contrast.py index 57ed290..a4963d0 100644 --- a/smauglab/transforms/cpu/contrast.py +++ b/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): """ diff --git a/smauglab/transforms/gpu/base.py b/smauglab/transforms/gpu/base.py index 65490ad..0079c5b 100644 --- a/smauglab/transforms/gpu/base.py +++ b/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 = [] diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index ef6c3c8..d9bb38e 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/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. diff --git a/smauglab/transforms/gpu/normalization.py b/smauglab/transforms/gpu/normalization.py new file mode 100644 index 0000000..db188f1 --- /dev/null +++ b/smauglab/transforms/gpu/normalization.py @@ -0,0 +1,569 @@ +"""Whole-volume intensity standardizations, and the schemes nnU-Net only has offline. + +`ZscoreNormalizationGPU` in `gpu/contrast.py` was the only normalization reachable from +a GPU config. Every other scheme a medical imaging pipeline uses -- min-max, the CT-style +percentile clip, a robust median/IQR scaling, divide-by-max, divide-by-mean -- existed +only inside nnU-Net's preprocessing +(`nnunetv2/preprocessing/normalization/default_normalization_schemes.py`), which runs once +per dataset, long before a patch reaches a training step. + +Putting them here makes the *choice of scheme* an augmentation: the builder puts this +group in a bucket that picks one member per draw (see `AugType.NORM` and +`build._normalization_stage`), so the network cannot come to depend on one fixed intensity +standardization. Two consequences of running at augmentation time rather than in +preprocessing are worth knowing, because they shape the parameters below: + +* The patch has already been normalized once, by preprocessing. These schemes therefore + act on standardized intensities, not on HU. A CT window in HU means nothing here unless + its bounds are given explicitly, which is why `PercentileNormalizationGPU` takes both an + explicit `clip_range` (nnU-Net parity, for a config that knows its dataset) and + per-sample `percentiles` (the useful augmentation). +* The global mean of z-scored data is ~0, so dividing by it amplifies rounding noise by + 1/mean. Hence `min_scale` -- a sample whose scale is negligible against its own range is + left alone -- and `stats_from="nonzero"` as `MeanScaleNormalizationGPU`'s default. + +Every scheme is a per-sample affine map with an optional clip in front of it, so the base +class owns the clone, the channel guard, the statistics mask, both degeneracy guards and +the region selection, and a leaf supplies only `_clip_bounds` and `_statistics`. That is +the same split `contrast.py` already uses for its gamma, convolution and function +families. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch +from torch import Tensor + +from smauglab.registry import AugId, AugType, Backend, register +from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from + +# Private to contrast.py, imported rather than copied: these four are the house rules for +# an intensity transform (clone, reject an out-of-range channel, apply the in_seg/out_seg +# region selection, refuse a non-finite result) and a second copy would be a second +# behaviour. The dependency is one-way -- contrast.py needs nothing from here. +from smauglab.transforms.gpu.contrast import _check_channel, _foreground, _select_and_check + +#: Which voxels a scheme reads its statistics from. `"seg"` is the labelled anatomy, via +#: the mask the pipeline injects; `"nonzero"` is everything the preprocessing did not zero +#: out, which is the closest thing here to nnU-Net's `use_mask_for_norm`. +STATS_SOURCES = ("all", "seg", "nonzero") + +#: Floor on a divisor's magnitude. Matches nnU-Net's `max(std, 1e-8)` so the parity tests +#: compare the same arithmetic, not an approximation of it. +_EPS = 1e-8 + +#: IQR and MAD of a standard normal. Dividing by these is what makes the robust scaling +#: comparable to a standard deviation instead of ~1.35x or ~0.67x of one. +_NORMAL_IQR = 1.34898 +_NORMAL_MAD = 0.67449 + + +def _broadcast(values: Tensor, like: Tensor) -> Tensor: + """Reshape a per-sample `[N]` tensor to broadcast over `like`'s spatial dims.""" + return values.view(-1, *([1] * (like.dim() - 1))) + + +def _resolve_mask(transform: _NormalizationBaseGPU, channel: Tensor, seg_mask: Tensor | None) -> Tensor | None: + """The voxels this transform's statistics come from, or None for all of them. + + A sample whose mask selects nothing has no statistics to read, and letting that + through would produce inf, trip the non-finite guard and print a warning for what is + just an unlabelled patch. Those rows fall back to the whole volume instead. The + fallback is unconditional rather than behind an `if`: the test is a device tensor and + reading it back to the host would stall the stream on every batch. + """ + if transform.stats_from == "all": + return None + if transform.stats_from == "nonzero": + mask = channel.abs() > transform.nonzero_threshold + else: # "seg" + if seg_mask is None: + return None + mask = _foreground(seg_mask, 1) + reduce_dims = tuple(range(1, mask.dim())) + return mask | _broadcast(~mask.any(dim=reduce_dims), mask) + + +def _masked_extrema(channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor]: + """Per-sample min and max over the masked voxels, as `[N]` tensors.""" + flat = channel.reshape(channel.shape[0], -1) + if mask is None: + return flat.amin(dim=1), flat.amax(dim=1) + flat_mask = mask.reshape(channel.shape[0], -1) + return ( + flat.masked_fill(~flat_mask, float("inf")).amin(dim=1), + flat.masked_fill(~flat_mask, float("-inf")).amax(dim=1), + ) + + +def _masked_moments(channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor]: + """Per-sample mean and population std over the masked voxels, as `[N]` tensors. + + `unbiased=False`, because nnU-Net normalizes with `np.std` (ddof=0) and so does + `ZscoreNormalizationGPU`; an unbiased estimate here would make the parity tests fail + on a difference that is not a bug. + """ + reduce_dims = tuple(range(1, channel.dim())) + if mask is None: + return channel.mean(dim=reduce_dims), channel.std(dim=reduce_dims, unbiased=False) + weights = mask.to(channel.dtype) + count = weights.sum(dim=reduce_dims).clamp_min(1.0) + mean = (channel * weights).sum(dim=reduce_dims) / count + deviation = channel - _broadcast(mean, channel) + variance = (deviation.pow(2) * weights).sum(dim=reduce_dims) / count + return mean, variance.sqrt() + + +def _masked_quantiles(channel: Tensor, mask: Tensor | None, probabilities: Sequence[float]) -> list[Tensor]: + """Per-sample quantiles over the masked voxels, in one batched sort. + + The other half of the problem `contrast._batched_quantiles` solves: there the + probability varies per sample over the whole volume, here the probabilities are shared + and the *population* varies per sample. Out-of-mask voxels are pushed to +inf so they + sort to the end, and each row is then indexed by its own mask count, with + `torch.quantile`'s linear interpolation (`low + frac * (high - low)`) so the values are + the ones it would return. No per-sample Python loop and no host read. + """ + batch = channel.shape[0] + flat = channel.reshape(batch, -1) + if mask is None: + ordered, _ = torch.sort(flat, dim=1) + counts = torch.full((batch,), flat.shape[1], device=channel.device, dtype=torch.long) + else: + flat_mask = mask.reshape(batch, -1) + ordered, _ = torch.sort(flat.masked_fill(~flat_mask, float("inf")), dim=1) + counts = flat_mask.sum(dim=1) + last = (counts - 1).clamp_min(0) + + picked = [] + for probability in probabilities: + position = last.to(ordered.dtype) * float(probability) + low_idx = position.floor().long() + high_idx = position.ceil().long() + frac = (position - low_idx.to(position.dtype)).unsqueeze(1) + low_val = ordered.gather(1, low_idx.unsqueeze(1)) + high_val = ordered.gather(1, high_idx.unsqueeze(1)) + picked.append((low_val + frac * (high_val - low_val)).squeeze(1)) + return picked + + +class _NormalizationBaseGPU(ImageOnlyTransform): + """Clone, mask, clip, affine-map, guard, write back -- everything but the scheme. + + A leaf implements `_statistics` (and `_clip_bounds` if it clips), each returning + per-sample `[N]` tensors; the base broadcasts them. `same_on_batch` is absent for the + same reason it is absent from `ZscoreNormalizationGPU`: a normalization draws no random + parameters, so there is nothing for the flag to hold the same. + + Args: + apply_to_channel: Channels to normalize. + stats_from: Which voxels the statistics come from; one of `STATS_SOURCES`. + nonzero_threshold: `|x|` above which a voxel counts when `stats_from="nonzero"`. + min_scale: Leave a sample alone when its divisor is below this fraction of the + sample's standard deviation. The guard that makes dividing by a near-zero mean + safe. + in_seg: Probability of applying the result inside the segmentation only. + out_seg: Probability of applying it outside the segmentation only. + p: Application probability. + p_batch: Batch-wise application probability. + keepdim: Keep the input's dimensions (passed to the base). + """ + + def __init__( + self, + apply_to_channel: Sequence[int] = (0,), + stats_from: str = "all", + nonzero_threshold: float = 1e-6, + min_scale: float = 1e-3, + in_seg: float = 0.0, + out_seg: float = 0.0, + p: float = 1.0, + p_batch: float = 1.0, + keepdim: bool = True, + ) -> None: + super().__init__(p=p, p_batch=p_batch, same_on_batch=False, keepdim=keepdim) + if stats_from not in STATS_SOURCES: + raise ValueError(f"{type(self).__name__}: stats_from must be one of {STATS_SOURCES}, got {stats_from!r}.") + self.apply_to_channel = apply_to_channel + self.stats_from = stats_from + self.nonzero_threshold = float(nonzero_threshold) + self.min_scale = float(min_scale) + self.in_seg = in_seg + self.out_seg = out_seg + # Not a constructor parameter: these schemes have no mix_in_out knob and the + # registry derives a config's accepted keys from __init__, so adding one there + # would invent a setting. Set here so the region-selection helper's contract holds. + self.mix_in_out = False + + def _clip_bounds(self, channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor] | None: + """Per-sample `(lo, hi)` to clamp to before the affine map, or None for no clip.""" + return None + + def _statistics(self, channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor, Tensor]: + """Per-sample `(center, scale, offset)` for `y = (x - center) / scale + offset`.""" + raise NotImplementedError + + @torch.no_grad() + def apply_transform( + self, + input: Tensor, + params: dict[str, Tensor], + flags: dict[str, Any], + transform: Tensor | None = None, + ) -> Tensor: + # input: (N, C, [D,] H, W) + # Clone: kornia passes the caller's tensor through; this writes channels back. + input = input.clone() + seg_mask = segmentation_from(params) + for c in self.apply_to_channel: + _check_channel(self, c, input.shape[1]) + channel = input[:, c] + orig = channel.clone() + mask = _resolve_mask(self, channel, seg_mask) + + flat = channel.reshape(channel.shape[0], -1) + spread = _broadcast(flat.amax(dim=1) - flat.amin(dim=1), channel) + deviation = _broadcast(_masked_moments(channel, mask)[1], channel) + + bounds = self._clip_bounds(channel, mask) + clipped = channel if bounds is None else torch.clamp(channel, _broadcast(bounds[0], channel), _broadcast(bounds[1], channel)) + center, scale, offset = (_broadcast(t, channel) for t in self._statistics(clipped, mask)) + + # Two guards, both "leave this sample as it is": + # + # * No range at all. There is no normalization of a volume with no variance, + # and the clamp below would turn floating-point residue into an O(1) value + # decided by rounding -- `ZscoreNormalizationGPU` returned a constant -1.0 + # for a constant -2.709 CT-air patch before it grew the same guard. Tested on + # the range rather than the scale: `amax - amin` is exactly zero for a + # constant tensor while a summed standard deviation is not. + # * A divisor negligible against the sample's own standard deviation. This is + # what a mean-scaling scheme hits on already-z-scored data, whose mean is ~0: + # the result is finite, so nothing downstream rejects it, and it is still + # noise times 1e8. + # + # The second guard is measured against the standard deviation and not against + # `spread`, which was the first thing tried and is wrong for the robust + # schemes: one 1e4 voxel puts the range four orders of magnitude above the IQR, + # so a range-relative floor declined to run on exactly the volume + # `RobustZscoreNormalizationGPU` exists for. The standard deviation moves with + # the outlier (240 rather than 2.9 on that volume) and so stays in proportion + # to any scale a scheme might derive from the same data. + applies = (spread > 0) & (scale.abs() >= self.min_scale * deviation) + # Sign-preserving floor, not `clamp_min`: a negative scale is meaningful + # (dividing by a negative mean inverts the volume, which is what the scheme + # says), and `clamp_min(-0.5, 1e-8)` would return 1e-8 and flip it. + safe = torch.where(scale.abs() < _EPS, torch.full_like(scale, _EPS), scale) + normalized = torch.where(applies, (clipped - center) / safe + offset, channel) + + checked = _select_and_check(self, orig, normalized, seg_mask) + if checked is None: + continue + input[:, c] = checked + + return input + + +@register( + aug_id=AugId.NORM_MINMAX, + backend=Backend.GPU, + group=AugType.NORM, +) +class MinMaxNormalizationGPU(_NormalizationBaseGPU): + """Rescale each sample's intensity range onto a fixed interval. + + At its defaults this is nnU-Net's `RescaleTo01Normalization`, including its floored + denominator. `clip_percentiles` clips to a percentile window first, which is what makes + the result robust to a single hot voxel -- and a different augmentation, since the + voxels outside the window are flattened onto the interval's ends. + + Args: + out_range: Interval to map onto. Must be non-empty. + clip_percentiles: `[lower, upper]` in [0, 1] to clip to first, or None for the + sample's own min and max. + """ + + def __init__( + self, + out_range: Sequence[float] = (0.0, 1.0), + clip_percentiles: Sequence[float] | None = None, + apply_to_channel: Sequence[int] = (0,), + stats_from: str = "all", + nonzero_threshold: float = 1e-6, + min_scale: float = 1e-3, + in_seg: float = 0.0, + out_seg: float = 0.0, + p: float = 1.0, + p_batch: float = 1.0, + keepdim: bool = True, + ) -> None: + super().__init__( + apply_to_channel=apply_to_channel, + stats_from=stats_from, + nonzero_threshold=nonzero_threshold, + min_scale=min_scale, + in_seg=in_seg, + out_seg=out_seg, + p=p, + p_batch=p_batch, + keepdim=keepdim, + ) + low, high = float(out_range[0]), float(out_range[1]) + if high <= low: + raise ValueError(f"MinMaxNormalizationGPU: out_range must be increasing, got {tuple(out_range)!r}.") + self.out_range = (low, high) + self.clip_percentiles = None if clip_percentiles is None else (float(clip_percentiles[0]), float(clip_percentiles[1])) + + def _clip_bounds(self, channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor] | None: + if self.clip_percentiles is None: + return None + low, high = _masked_quantiles(channel, mask, self.clip_percentiles) + return low, high + + def _statistics(self, channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor, Tensor]: + low, high = _masked_extrema(channel, mask) + out_low, out_high = self.out_range + return low, (high - low) / (out_high - out_low), torch.full_like(low, out_low) + + +@register( + aug_id=AugId.NORM_PERCENTILE, + backend=Backend.GPU, + group=AugType.NORM, +) +class PercentileNormalizationGPU(_NormalizationBaseGPU): + """Clip to a percentile window, then standardize. nnU-Net's CT scheme, generalized. + + With `clip_range`, `mean` and `std` supplied this is exactly + `CTNormalization`: nnU-Net reads those three from the dataset's intensity properties, + which nothing hands a transform at training time, so a config that knows them states + them. Left as None they are measured per sample -- the same scheme, estimated from the + patch, which is the form that works as an augmentation. + + Args: + percentiles: `[lower, upper]` in [0, 1] to clip to when `clip_range` is None. + clip_range: Explicit `[lower, upper]` intensities to clip to, overriding + `percentiles`. + mean: Center to subtract, or None to use the clipped sample's own mean. + std: Divisor, or None to use the clipped sample's own standard deviation. + """ + + def __init__( + self, + percentiles: Sequence[float] = (0.005, 0.995), + clip_range: Sequence[float] | None = None, + mean: float | None = None, + std: float | None = None, + apply_to_channel: Sequence[int] = (0,), + stats_from: str = "all", + nonzero_threshold: float = 1e-6, + min_scale: float = 1e-3, + in_seg: float = 0.0, + out_seg: float = 0.0, + p: float = 1.0, + p_batch: float = 1.0, + keepdim: bool = True, + ) -> None: + super().__init__( + apply_to_channel=apply_to_channel, + stats_from=stats_from, + nonzero_threshold=nonzero_threshold, + min_scale=min_scale, + in_seg=in_seg, + out_seg=out_seg, + p=p, + p_batch=p_batch, + keepdim=keepdim, + ) + self.percentiles = (float(percentiles[0]), float(percentiles[1])) + self.clip_range = None if clip_range is None else (float(clip_range[0]), float(clip_range[1])) + self.mean = None if mean is None else float(mean) + self.std = None if std is None else float(std) + + def _clip_bounds(self, channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor] | None: + if self.clip_range is not None: + batch = (channel.shape[0],) + return ( + channel.new_full(batch, self.clip_range[0]), + channel.new_full(batch, self.clip_range[1]), + ) + low, high = _masked_quantiles(channel, mask, self.percentiles) + return low, high + + def _statistics(self, channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor, Tensor]: + center, scale = _masked_moments(channel, mask) + if self.mean is not None: + center = torch.full_like(center, self.mean) + if self.std is not None: + scale = torch.full_like(scale, self.std) + return center, scale, torch.zeros_like(center) + + +@register( + aug_id=AugId.NORM_ROBUST_ZSCORE, + backend=Backend.GPU, + group=AugType.NORM, +) +class RobustZscoreNormalizationGPU(_NormalizationBaseGPU): + """Z-score on robust statistics: the median, scaled by the IQR or the MAD. + + The scheme to reach for when the patch has outliers a mean and standard deviation + would chase -- metal artefacts, a bright rim, a fat saturation band. `unit_variance` + divides the spread by its value for a standard normal, so the output is on the same + scale as an ordinary z-score rather than ~1.35x or ~0.67x of it. + + Args: + percentiles: `[lower, upper]` defining the IQR. Ignored when `spread="mad"`. + spread: `"iqr"` or `"mad"`. + unit_variance: Rescale the spread so Gaussian input lands at unit variance. + """ + + SPREADS = ("iqr", "mad") + + def __init__( + self, + percentiles: Sequence[float] = (0.25, 0.75), + spread: str = "iqr", + unit_variance: bool = True, + apply_to_channel: Sequence[int] = (0,), + stats_from: str = "all", + nonzero_threshold: float = 1e-6, + min_scale: float = 1e-3, + in_seg: float = 0.0, + out_seg: float = 0.0, + p: float = 1.0, + p_batch: float = 1.0, + keepdim: bool = True, + ) -> None: + super().__init__( + apply_to_channel=apply_to_channel, + stats_from=stats_from, + nonzero_threshold=nonzero_threshold, + min_scale=min_scale, + in_seg=in_seg, + out_seg=out_seg, + p=p, + p_batch=p_batch, + keepdim=keepdim, + ) + if spread not in self.SPREADS: + raise ValueError(f"RobustZscoreNormalizationGPU: spread must be one of {self.SPREADS}, got {spread!r}.") + self.percentiles = (float(percentiles[0]), float(percentiles[1])) + self.spread = spread + self.unit_variance = bool(unit_variance) + + def _statistics(self, channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor, Tensor]: + (median,) = _masked_quantiles(channel, mask, (0.5,)) + if self.spread == "iqr": + low, high = _masked_quantiles(channel, mask, self.percentiles) + scale = high - low + if self.unit_variance: + scale = scale / _NORMAL_IQR + else: + deviation = (channel - _broadcast(median, channel)).abs() + (scale,) = _masked_quantiles(deviation, mask, (0.5,)) + if self.unit_variance: + scale = scale / _NORMAL_MAD + return median, scale, torch.zeros_like(median) + + +@register( + aug_id=AugId.NORM_MAX_SCALE, + backend=Backend.GPU, + group=AugType.NORM, +) +class MaxScaleNormalizationGPU(_NormalizationBaseGPU): + """Divide by the sample's maximum, leaving zero where it is. + + The MRI convention that a z-score does not cover: it fixes the scale without moving + the origin, so a background that preprocessing set to zero stays at zero. `percentile` + uses a high quantile instead of the maximum, which keeps one hot voxel from deciding + the scale for the whole volume -- the values above it then exceed the nominal 1. + + Args: + use_abs: Scale by `max|x|` rather than by `max x`, so negative data (a z-scored + patch) does not get a scale smaller than its own amplitude. + percentile: Quantile in [0, 1] to use in place of the maximum, or None. + """ + + def __init__( + self, + use_abs: bool = True, + percentile: float | None = None, + apply_to_channel: Sequence[int] = (0,), + stats_from: str = "all", + nonzero_threshold: float = 1e-6, + min_scale: float = 1e-3, + in_seg: float = 0.0, + out_seg: float = 0.0, + p: float = 1.0, + p_batch: float = 1.0, + keepdim: bool = True, + ) -> None: + super().__init__( + apply_to_channel=apply_to_channel, + stats_from=stats_from, + nonzero_threshold=nonzero_threshold, + min_scale=min_scale, + in_seg=in_seg, + out_seg=out_seg, + p=p, + p_batch=p_batch, + keepdim=keepdim, + ) + self.use_abs = bool(use_abs) + self.percentile = None if percentile is None else float(percentile) + + def _statistics(self, channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor, Tensor]: + source = channel.abs() if self.use_abs else channel + if self.percentile is None: + _, scale = _masked_extrema(source, mask) + else: + (scale,) = _masked_quantiles(source, mask, (self.percentile,)) + return torch.zeros_like(scale), scale, torch.zeros_like(scale) + + +@register( + aug_id=AugId.NORM_MEAN_SCALE, + backend=Backend.GPU, + group=AugType.NORM, +) +class MeanScaleNormalizationGPU(_NormalizationBaseGPU): + """Divide by the sample's mean, leaving zero where it is. + + The other scale-only MRI scheme, and the one that needs its statistics masked: the + mean of a z-scored volume is ~0 over every voxel, so `stats_from` defaults to + `"nonzero"` here. Where the mean is still negligible against the sample's range the + base class leaves that sample alone rather than dividing by it; see `min_scale`. + """ + + def __init__( + self, + apply_to_channel: Sequence[int] = (0,), + stats_from: str = "nonzero", + nonzero_threshold: float = 1e-6, + min_scale: float = 1e-3, + in_seg: float = 0.0, + out_seg: float = 0.0, + p: float = 1.0, + p_batch: float = 1.0, + keepdim: bool = True, + ) -> None: + super().__init__( + apply_to_channel=apply_to_channel, + stats_from=stats_from, + nonzero_threshold=nonzero_threshold, + min_scale=min_scale, + in_seg=in_seg, + out_seg=out_seg, + p=p, + p_batch=p_batch, + keepdim=keepdim, + ) + + def _statistics(self, channel: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor, Tensor]: + mean, _ = _masked_moments(channel, mask) + return torch.zeros_like(mean), mean, torch.zeros_like(mean) diff --git a/unit_tests/test_channel_guard.py b/unit_tests/test_channel_guard.py index 57868ae..6cd251f 100644 --- a/unit_tests/test_channel_guard.py +++ b/unit_tests/test_channel_guard.py @@ -13,9 +13,24 @@ import torch from smauglab.transforms.gpu.contrast import RandomBiasFieldGPU, RandomBrightnessGPU, ZscoreNormalizationGPU +from smauglab.transforms.gpu.normalization import ( + MaxScaleNormalizationGPU, + MeanScaleNormalizationGPU, + MinMaxNormalizationGPU, + PercentileNormalizationGPU, + RobustZscoreNormalizationGPU, +) from unit_tests.helpers import SmaugLabTestCase -GUARDED = (RandomBiasFieldGPU, ZscoreNormalizationGPU) +GUARDED = ( + RandomBiasFieldGPU, + ZscoreNormalizationGPU, + MinMaxNormalizationGPU, + PercentileNormalizationGPU, + RobustZscoreNormalizationGPU, + MaxScaleNormalizationGPU, + MeanScaleNormalizationGPU, +) class TestOutOfRangeChannelRaises(SmaugLabTestCase): diff --git a/unit_tests/test_normalization.py b/unit_tests/test_normalization.py new file mode 100644 index 0000000..058ac1b --- /dev/null +++ b/unit_tests/test_normalization.py @@ -0,0 +1,449 @@ +"""The GPU normalization schemes, and the choice the builder makes between them. + +`test_transforms_gpu.py` already runs every registered transform over a tiny volume and +checks it stays finite and leaves the mask alone. That is not enough here: a normalization +whose arithmetic is subtly wrong still produces a finite volume of the right shape, and +three of these schemes claim to reproduce an nnU-Net preprocessing scheme exactly. So this +file asserts the maths, the parity, and the two guards -- the constant patch and the +negligible divisor -- that decide when a scheme declines to run. + +The bucket tests are here rather than in `test_builder.py` because the thing being tested +is semantic, not structural: with `num_transforms=1` a scheme's `p` stops being "how often +do I fire" and becomes "what share of the choice is mine", and nothing else in the +repository works that way. +""" + +from __future__ import annotations + +import unittest +from typing import ClassVar + +import torch + +from smauglab.config import OrderSource, PipelineMode +from smauglab.registry import AugType, Backend, names +from smauglab.transforms.build import build_gpu_pipeline +from smauglab.transforms.gpu.contrast import ZscoreNormalizationGPU +from smauglab.transforms.gpu.normalization import ( + MaxScaleNormalizationGPU, + MeanScaleNormalizationGPU, + MinMaxNormalizationGPU, + PercentileNormalizationGPU, + RobustZscoreNormalizationGPU, +) +from smauglab.transforms.gpu.transforms_list import RandomChooseXTransformsGPU +from unit_tests.helpers import SmaugLabTestCase + +#: Every scheme in the group, at defaults. Built from the registry rather than hand-listed +#: so a sixth one cannot be added without the guard tests below seeing it. +SCHEMES = ( + MinMaxNormalizationGPU, + PercentileNormalizationGPU, + RobustZscoreNormalizationGPU, + MaxScaleNormalizationGPU, + MeanScaleNormalizationGPU, +) + + +def apply(transform, volume, seg=None): + """Call a transform the way its pipeline would, without the kornia container.""" + params = {} if seg is None else {"seg": seg} + return transform.apply_transform(volume.clone(), params, transform.flags) + + +class TestTheGroupIsComplete(unittest.TestCase): + def test_every_registered_normalization_is_covered(self): + """A new scheme joins SCHEMES, or this says so.""" + registered = set(names(Backend.GPU, group=AugType.NORM)) + covered = {cls.__name__ for cls in SCHEMES} | {ZscoreNormalizationGPU.__name__} + self.assertEqual(registered, covered) + + +class TestEachSchemeDoesWhatItsNameSays(SmaugLabTestCase): + """One assertion per scheme on the property the class name claims.""" + + def volume(self): + # Offset and scaled away from zero mean and unit variance, so a scheme that + # silently did nothing would not pass by accident. + return torch.randn(2, 1, 12, 12, 12) * 3.0 + 5.0 + + def test_minmax_lands_on_its_output_range(self): + for out_range in ((0.0, 1.0), (-1.0, 1.0)): + with self.subTest(out_range=out_range): + out = apply(MinMaxNormalizationGPU(out_range=out_range, p=1.0), self.volume()) + self.assertAlmostEqual(float(out.amin()), out_range[0], places=5) + self.assertAlmostEqual(float(out.amax()), out_range[1], places=5) + + def test_percentile_standardizes_what_it_clipped(self): + out = apply(PercentileNormalizationGPU(p=1.0), self.volume()) + self.assertAlmostEqual(float(out.mean()), 0.0, places=4) + self.assertAlmostEqual(float(out.std(unbiased=False)), 1.0, places=3) + + def test_percentile_clips_to_the_window_it_was_given(self): + """The clip is the scheme, not a detail: a wider window is a weaker transform.""" + volume = self.volume() + narrow = apply(PercentileNormalizationGPU(clip_range=[0.0, 6.0], mean=3.0, std=1.0, p=1.0), volume) + self.assertAlmostEqual(float(narrow.amin()), -3.0, places=5) + self.assertAlmostEqual(float(narrow.amax()), 3.0, places=5) + + def test_robust_zscore_centres_the_median(self): + for spread in RobustZscoreNormalizationGPU.SPREADS: + with self.subTest(spread=spread): + out = apply(RobustZscoreNormalizationGPU(spread=spread, p=1.0), self.volume()) + # Per sample: each one is centred on its own median, and the median of two + # pooled distributions is not the median of either. Compared against + # `torch.quantile`, not `torch.median`: on an even number of voxels the + # transform interpolates between the two middle values, as quantile does, + # while `median` returns the lower of them. + for sample in out: + self.assertAlmostEqual(float(torch.quantile(sample.flatten(), 0.5)), 0.0, places=5) + # Gaussian input, so the unit_variance rescaling should land near 1. + self.assertAlmostEqual(float(sample.std(unbiased=False)), 1.0, places=1) + + def test_robust_zscore_ignores_an_outlier_a_mean_would_chase(self): + """What "robust" buys, stated as the comparison that fails without it. + + One 1e4 voxel takes the sample's standard deviation from ~3 to ~240, so an ordinary + z-score divides the anatomy by eighty times too much and flattens it. The IQR does + not move, so the robust scheme still returns a unit-variance volume. + """ + volume = self.volume()[:1] + volume[0, 0, 0, 0, 0] = 1e4 + interior = (slice(None), slice(None), slice(1, None)) + + robust = apply(RobustZscoreNormalizationGPU(p=1.0), volume)[interior] + plain = apply(ZscoreNormalizationGPU(p=1.0), volume)[interior] + self.assertAlmostEqual(float(robust.std(unbiased=False)), 1.0, places=1) + self.assertLess(float(plain.std(unbiased=False)), 0.1) + + def test_max_scale_puts_the_maximum_at_one_and_leaves_zero_alone(self): + volume = self.volume() + volume[:, :, :2] = 0.0 + out = apply(MaxScaleNormalizationGPU(p=1.0), volume) + self.assertAlmostEqual(float(out.abs().amax()), 1.0, places=5) + self.assertEqual(float(out[:, :, :2].abs().amax()), 0.0) + + def test_max_scale_percentile_leaves_the_hot_voxels_above_one(self): + out = apply(MaxScaleNormalizationGPU(percentile=0.99, p=1.0), self.volume()) + self.assertGreater(float(out.abs().amax()), 1.0) + + def test_mean_scale_puts_the_mean_at_one(self): + out = apply(MeanScaleNormalizationGPU(stats_from="all", p=1.0), self.volume()) + self.assertAlmostEqual(float(out.mean()), 1.0, places=4) + + +class TestParityWithNnunetPreprocessing(SmaugLabTestCase): + """Three schemes claim an nnU-Net equivalent; this is what makes the claim checkable. + + Exact equality wherever the arithmetic is fully determined -- the schemes were written + against nnU-Net's own, including its floored denominators + (`np.clip(image.max(), 1e-8)`, `max(std, 1e-8)`) and its population standard + deviation. The one case with a tolerance is the per-sample mean and standard + deviation, where numpy and torch sum 1000 float32 values in a different order; the + gap there is one ulp and is not something a test should pin to zero. + + Not asserted: the masked variants. nnU-Net masks on `seg >= 0`, i.e. everything inside + the body including unlabelled background, while `stats_from="seg"` here means the + labelled anatomy only. The two are different populations on purpose. + """ + + #: Intensity properties of the shape nnU-Net's CT scheme reads from its dataset + #: fingerprint. A transform is never handed these at training time, which is why + #: `PercentileNormalizationGPU` takes them as parameters. + PROPERTIES: ClassVar[dict] = {"mean": 38.5, "std": 115.0, "percentile_00_5": -150.0, "percentile_99_5": 260.0} + + def setUp(self): + super().setUp() + try: + from nnunetv2.preprocessing.normalization import default_normalization_schemes + except ImportError as missing: # pragma: no cover -- an nnU-Net that moved the module + self.skipTest(f"nnU-Net normalization schemes not importable: {missing}") + self.schemes = default_normalization_schemes + # Hounsfield-ish, so the CT scheme's percentile window actually bites. + self.volume = torch.randn(1, 1, 10, 10, 10) * 120.0 + 40.0 + self.image = self.volume[0, 0].numpy().copy() + + def _compare(self, reference, transform, atol=0.0): + got = apply(transform, self.volume)[0, 0].numpy() + self.assertTrue( + torch.allclose(torch.from_numpy(reference), torch.from_numpy(got), rtol=0.0, atol=atol), + f"{type(transform).__name__} diverged from nnU-Net by up to {abs(reference - got).max()}", + ) + + def test_minmax_matches_rescale_to_01(self): + reference = self.schemes.RescaleTo01Normalization(use_mask_for_norm=False, intensityproperties={}).run(self.image.copy()) + self._compare(reference, MinMaxNormalizationGPU(p=1.0)) + + def test_percentile_matches_ct_normalization(self): + reference = self.schemes.CTNormalization(use_mask_for_norm=False, intensityproperties=self.PROPERTIES).run(self.image.copy()) + self._compare( + reference, + PercentileNormalizationGPU( + clip_range=[self.PROPERTIES["percentile_00_5"], self.PROPERTIES["percentile_99_5"]], + mean=self.PROPERTIES["mean"], + std=self.PROPERTIES["std"], + p=1.0, + ), + ) + + def test_percentile_without_a_window_matches_zscore(self): + """The degenerate case of the CT scheme: clip to the full range, measure per sample.""" + reference = self.schemes.ZScoreNormalization(use_mask_for_norm=False, intensityproperties={}).run(self.image.copy()) + # Both reductions, so float32 accumulation order shows up; see the class docstring. + self._compare(reference, PercentileNormalizationGPU(percentiles=[0.0, 1.0], p=1.0), atol=1e-6) + + +class TestTheGuardsDecline(SmaugLabTestCase): + """Both guards, and the paired assertion that neither disables the scheme outright.""" + + def test_a_constant_channel_is_returned_unchanged(self): + """No scheme has anything to say about a volume with no variance. + + `ZscoreNormalizationGPU` used to return a constant -1.0 for a constant -2.709 CT-air + patch, an O(1) value decided entirely by rounding; the same trap is waiting in every + divisor here. + """ + for cls in SCHEMES: + for value in (-2.709, 0.0, 1.0): + with self.subTest(scheme=cls.__name__, value=value): + volume = self.constant_volume(value) + out = apply(cls(p=1.0), volume) + self.assertTrue(bool(torch.equal(out, volume)), f"a constant {value} volume was rewritten") + + def test_a_normal_volume_is_still_normalized(self): + """The guard against passing the test above by doing nothing at all.""" + volume = torch.randn(1, 1, 12, 12, 12) * 2.0 + 4.0 + for cls in SCHEMES: + with self.subTest(scheme=cls.__name__): + out = apply(cls(p=1.0), volume) + self.assertFalse(bool(torch.equal(out, volume)), "the scheme left a perfectly normal volume alone") + self.assertTrue(bool(torch.isfinite(out).all())) + + def test_mean_scaling_declines_on_already_zscored_data(self): + """The reason `min_scale` exists. + + The incoming patch has been z-scored by preprocessing, so its mean over every voxel + is ~0 and `x / mean` is rounding noise times a hundred million. The result is finite, + so nothing downstream rejects it. + """ + volume = torch.randn(1, 1, 12, 12, 12) + volume = volume - volume.mean() + out = apply(MeanScaleNormalizationGPU(stats_from="all", p=1.0), volume) + self.assertTrue(bool(torch.equal(out, volume)), "divided by a mean of ~0") + + def test_mean_scaling_still_runs_where_the_mean_is_real(self): + volume = torch.rand(1, 1, 12, 12, 12) + 1.0 + out = apply(MeanScaleNormalizationGPU(stats_from="all", p=1.0), volume) + self.assertAlmostEqual(float(out.mean()), 1.0, places=4) + + def test_an_unknown_stats_source_is_rejected_at_construction(self): + with self.assertRaises(ValueError): + MinMaxNormalizationGPU(stats_from="foreground") + + def test_an_empty_out_range_is_rejected_at_construction(self): + with self.assertRaises(ValueError): + MinMaxNormalizationGPU(out_range=(1.0, 1.0)) + + +class TestMaskedStatistics(SmaugLabTestCase): + """`stats_from` changes which voxels are measured, never which ones are written.""" + + def _bordered(self): + """Interior in [1, 2], border exactly zero -- what masked preprocessing leaves.""" + volume = torch.zeros(1, 1, 12, 12, 12) + volume[:, :, 2:10, 2:10, 2:10] = torch.rand(1, 1, 8, 8, 8) + 1.0 + return volume + + def test_nonzero_statistics_ignore_the_zeroed_border(self): + volume = self._bordered() + interior = (slice(None), slice(None), slice(2, 10), slice(2, 10), slice(2, 10)) + + out = apply(MinMaxNormalizationGPU(stats_from="nonzero", p=1.0), volume) + self.assertAlmostEqual(float(out[interior].amin()), 0.0, places=5) + self.assertAlmostEqual(float(out[interior].amax()), 1.0, places=5) + # The border was still written, it just was not measured: 0 sits below the + # interior minimum, so it maps outside the output range. + self.assertLess(float(out.amin()), 0.0) + + def test_the_unmasked_scheme_measures_the_border_too(self): + """The comparison that makes the test above about masking rather than luck.""" + out = apply(MinMaxNormalizationGPU(stats_from="all", p=1.0), self._bordered()) + self.assertAlmostEqual(float(out.amin()), 0.0, places=5) + + def test_seg_statistics_come_from_the_labelled_region(self): + volume = torch.rand(1, 1, 12, 12, 12) + volume[:, :, 4:8, 4:8, 4:8] += 4.0 + labelled = (slice(None), slice(None), slice(4, 8), slice(4, 8), slice(4, 8)) + + out = apply(MinMaxNormalizationGPU(stats_from="seg", p=1.0), volume, self.tiny_seg_like(volume)) + self.assertAlmostEqual(float(out[labelled].amin()), 0.0, places=5) + self.assertAlmostEqual(float(out[labelled].amax()), 1.0, places=5) + + def test_an_empty_mask_falls_back_to_the_whole_volume(self): + """An unlabelled patch is not an error, and must not print a non-finite warning.""" + volume = torch.rand(1, 1, 12, 12, 12) + out = apply(MinMaxNormalizationGPU(stats_from="seg", p=1.0), volume, torch.zeros_like(volume)) + self.assertAlmostEqual(float(out.amin()), 0.0, places=5) + self.assertAlmostEqual(float(out.amax()), 1.0, places=5) + + def tiny_seg_like(self, volume): + seg = torch.zeros_like(volume) + seg[:, :, 4:8, 4:8, 4:8] = 1.0 + return seg + + +class TestTheBuilderMakesAChoiceNotASequence(SmaugLabTestCase): + """Where the normalization stage lands, and when it is a bucket at all.""" + + ONE: ClassVar[dict] = {"RandomFlipTransformGPU": {"p": 1.0}, "ZscoreNormalizationGPU": {"p": 1.0}} + MANY: ClassVar[dict] = { + "RandomFlipTransformGPU": {"p": 1.0}, + "ZscoreNormalizationGPU": {"p": 1.0}, + "MinMaxNormalizationGPU": {"p": 1.0}, + "MaxScaleNormalizationGPU": {"p": 1.0}, + } + + def test_a_single_normalization_is_not_wrapped(self): + """Wrapping it would move it, and gate it behind a second probability.""" + for mode in PipelineMode: + with self.subTest(mode=mode.value): + built = build_gpu_pipeline(self.ONE, mode=mode) + self.assertIsInstance(built[-1], ZscoreNormalizationGPU) + + def test_a_single_normalization_keeps_its_configured_position(self): + """`pipeline.order: config` put it first; the normalization stage must not hoist it.""" + built = build_gpu_pipeline(self.ONE, order_source=OrderSource.CONFIG) + self.assertEqual([type(t).__name__ for t in built], ["RandomFlipTransformGPU", "ZscoreNormalizationGPU"]) + + def test_several_normalizations_become_one_trailing_choice(self): + for mode in PipelineMode: + with self.subTest(mode=mode.value): + built = build_gpu_pipeline(self.MANY, mode=mode) + bucket = built[-1] + self.assertIsInstance(bucket, RandomChooseXTransformsGPU) + self.assertEqual(bucket.provenance_label, "norm") + self.assertEqual(bucket.num_transforms, 1) + self.assertEqual(len(bucket.transforms_list), 3) + self.assertTrue(bucket.random_order, "a fixed order would pin the choice to one scheme") + others = [t for t in built[:-1] if isinstance(t, RandomChooseXTransformsGPU)] + self.assertNotIn("norm", [t.provenance_label for t in others], "only one normalization stage") + + def test_no_normalization_emits_no_stage(self): + built = build_gpu_pipeline({"RandomFlipTransformGPU": {"p": 1.0}}) + self.assertEqual([type(t).__name__ for t in built], ["RandomFlipTransformGPU"]) + + def test_the_bucket_probability_comes_from_the_pipeline_options(self): + built = build_gpu_pipeline(self.MANY, options={"norm_probability": 0.25}) + self.assertAlmostEqual(built[-1].p, 0.25) + + +class TestOneDrawPerBatchOrPerSample(SmaugLabTestCase): + """Which `same_on_batch` decides, because it is not the bucket's own argument. + + `_normalization_stage` builds the bucket with `same_on_batch=False`, and then the + container overwrites it: `pipeline.same_on_batch` defaults to True and kornia forces a + non-None value onto every stage it holds. So the shipped default is one scheme for the + whole batch, as it already is for the `ta` and `ge` buckets. This was written the wrong + way round in a docstring first; the assertion is here so the next reader does not have + to run a batch to find out. + """ + + SECTION: ClassVar[dict] = { + "MinMaxNormalizationGPU": {"p": 1.0}, + "MaxScaleNormalizationGPU": {"p": 1.0}, + } + + def _bucket(self, *, same_on_batch): + from smauglab.transforms.gpu.base import AugmentationSequentialCustom + + stages = build_gpu_pipeline(self.SECTION) + container = AugmentationSequentialCustom( + *stages, + data_keys=["input", "mask"], + same_on_batch=True if same_on_batch else None, + ) + return container[-1] + + def test_the_bucket_is_built_per_sample(self): + self.assertFalse(RandomChooseXTransformsGPU(transforms_list=[], num_transforms=1).same_on_batch) + + def test_the_container_default_makes_the_draw_batch_wide(self): + self.assertTrue(self._bucket(same_on_batch=True).same_on_batch) + + def test_opting_out_leaves_the_draw_per_sample(self): + self.assertFalse(self._bucket(same_on_batch=False).same_on_batch) + + def test_a_per_sample_bucket_can_give_two_samples_different_schemes(self): + """The property the per-sample setting buys, on identical inputs.""" + torch.manual_seed(0) + bucket = self._bucket(same_on_batch=False) + volume = torch.rand(1, 1, 8, 8, 8) + 2.0 + batch = torch.cat([volume, volume], dim=0) + + differed = sum(not torch.allclose(*apply(bucket, batch), atol=1e-6) for _ in range(40)) + self.assertGreater(differed, 5, "the two samples always drew the same scheme") + + +class TestProbabilityIsTheSchemesWeight(SmaugLabTestCase): + """Inside a `num_transforms=1` bucket, `p` stops meaning "how often" and starts + meaning "what share". + + `RandomChooseXTransformsGPU` draws one child uniformly and then spends that child's own + `p` (`transforms_list.py`), so `P(scheme) = p / n` and the leftover is no normalization + at all. That is the identity arm of the random choice, and it exists without a + registered no-op transform -- which is the whole reason there isn't one. + """ + + DRAWS = 400 + + def test_a_scheme_at_half_probability_takes_half_its_share(self): + torch.manual_seed(0) + # Two schemes, one of them switched off: the bucket should normalize on ~half the + # draws (minmax chosen, p=1.0) and leave the volume alone on the rest (the p=0 + # scheme chosen, nothing applied). + bucket = RandomChooseXTransformsGPU( + transforms_list=[MinMaxNormalizationGPU(p=1.0), MaxScaleNormalizationGPU(p=0.0)], + num_transforms=1, + p=1.0, + ) + volume = torch.rand(1, 1, 8, 8, 8) + 2.0 + + normalized = sum(not torch.equal(apply(bucket, volume), volume) for _ in range(self.DRAWS)) + expected = 0.5 * self.DRAWS + self.assertGreater(normalized, 0.65 * expected, "the switched-off scheme is swallowing draws") + self.assertLess(normalized, 1.35 * expected, "the switched-off scheme is not taking its share") + + def test_every_scheme_in_the_bucket_gets_chosen(self): + """A uniform draw over the members, not a favourite. + + The three schemes are told apart by what they do to a [2, 3] volume: min-max puts + the maximum at 1 and the minimum at 0, max-scale puts the maximum at 1 and leaves + the minimum positive, and the identity leaves the maximum above 1. + """ + torch.manual_seed(1) + bucket = RandomChooseXTransformsGPU( + transforms_list=[MinMaxNormalizationGPU(p=1.0), MaxScaleNormalizationGPU(p=1.0), MinMaxNormalizationGPU(p=0.0)], + num_transforms=1, + p=1.0, + ) + volume = torch.rand(1, 1, 8, 8, 8) + 2.0 + + seen = {"minmax": 0, "maxscale": 0, "identity": 0} + for _ in range(self.DRAWS): + out = apply(bucket, volume) + if float(out.amax()) > 1.5: + seen["identity"] += 1 + elif float(out.amin()) < 1e-6: + seen["minmax"] += 1 + else: + seen["maxscale"] += 1 + + expected = self.DRAWS / 3 + for scheme, count in seen.items(): + with self.subTest(scheme=scheme): + self.assertGreater(count, 0.65 * expected, f"{scheme} is under-represented: {seen}") + self.assertLess(count, 1.35 * expected, f"{scheme} is over-represented: {seen}") + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_transforms_gpu.py b/unit_tests/test_transforms_gpu.py index d370f68..1b1ff7d 100644 --- a/unit_tests/test_transforms_gpu.py +++ b/unit_tests/test_transforms_gpu.py @@ -25,6 +25,7 @@ "smauglab.transforms.gpu.contrast", "smauglab.transforms.gpu.spatial", "smauglab.transforms.gpu.fromSeg", + "smauglab.transforms.gpu.normalization", "smauglab.transforms.gpu.domain_transfer", "smauglab.transforms.gpu.palette.transform", ]