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", ]