From 92d8574426b3be5570cadd2d217983863bf040e5 Mon Sep 17 00:00:00 2001 From: iback Date: Tue, 13 Oct 2026 08:16:02 +0000 Subject: [PATCH 1/2] 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", ] From 869f706c4b6c2f899997612dd496e0ec24504018 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 15 Oct 2026 14:50:05 +0000 Subject: [PATCH 2/2] feat: snapshot the intensity profile before the augmentations and restore it after `retain_stats` restores a channel's mean and standard deviation one transform at a time, which pays for N restores to end up where the last one puts it, knows only two statistics -- so a transform that manufactures extreme voxels drags the standard deviation and the restore with it -- and restores the input's profile exactly, which is the opposite of what a domain-generalising pipeline wants. This is the same idea for the whole pipeline. `pipeline.retain_stats` brackets the augmentation stack: the snapshot runs after the geometry and records the per-sample, per-channel mean, standard deviation, 1st and 99th percentile; the restore runs before the trailing crop and normalization. Four statistics over-determine an affine map, so each pass draws *two* of them and solves for the map that puts exactly those two on target, leaving the other two where the augmentations left them. Writing each statistic as `t = a * u + b * c`, with `c = 0` for the offset-invariant standard deviation, makes all six pairs one 2x2 solve. The drawn pair is itself an augmentation: it spans "keep the moments", which is what `retain_stats` did, and "keep the robust dynamic range", which the outliers an edge filter manufactures cannot move. `std_scale` and `percentile_scale` then jitter the targets; the mean is never jittered, being the anchor every pair shares. Crop moves after the restore because it occludes a random box with the volume's own minimum rather than changing intensities, and the normalization stage stays last because deciding the scale the network sees is its whole job. With no `pipeline.retain_stats` block nothing is emitted and the RNG stream is untouched; with one, the restore spends a fixed five uniforms per channel whichever pair comes out, so a jitter sweep is comparable against its own baseline at a seed. `retain_stats` itself is untouched. The new shipped config turns it off on every transform, because the two do not compose, and leaves all 33 existing configs byte-identical. Co-Authored-By: Claude Opus 5 --- .gitignore | 3 + README.md | 70 ++- smauglab/config.py | 28 +- .../transform_params_gpu_stats-restore.json | 263 +++++++++++ smauglab/trainers/nnUNetTrainerDAExt.py | 1 + smauglab/transforms/build.py | 119 +++-- smauglab/transforms/gpu/normalization.py | 20 +- smauglab/transforms/gpu/stats.py | 437 ++++++++++++++++++ smauglab/transforms/gpu/transforms.py | 1 + unit_tests/test_pipeline_stats.py | 426 +++++++++++++++++ 10 files changed, 1328 insertions(+), 40 deletions(-) create mode 100644 smauglab/configs/transform_params_gpu_stats-restore.json create mode 100644 smauglab/transforms/gpu/stats.py create mode 100644 unit_tests/test_pipeline_stats.py diff --git a/.gitignore b/.gitignore index 794750d..32dedd4 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,9 @@ MANIFEST # configs beside it are tracked only because they predate the blanket rule above; # a new one has to say so. !smauglab/configs/transform_params_paper-synthseg.json +# The statistics bracket's reference config (see "Statistics restore" in README.md), +# shipped so `unit_tests/test_configs.py` builds and runs it like any other. +!smauglab/configs/transform_params_gpu_stats-restore.json *.yaml *.nii *.nii.gz diff --git a/README.md b/README.md index 5a67c16..05929f3 100644 --- a/README.md +++ b/README.md @@ -234,7 +234,8 @@ Three things to know before turning this on: 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. + scale. There is no compensating re-normalization at the end -- if that is what you want, + it is the statistics bracket below, which runs *before* this stage. * The schemes act on preprocessed intensities, not on HU. A CT window in Hounsfield units means nothing after preprocessing unless its bounds are stated, which is why `PercentileNormalizationGPU` takes an explicit `clip_range`, `mean` and `std` -- given @@ -242,6 +243,73 @@ Three things to know before turning this on: `MinMaxNormalizationGPU` at its defaults reproduces `RescaleTo01Normalization`. `unit_tests/test_normalization.py` pins both against the installed nnU-Net. +### Statistics restore (`pipeline.retain_stats`) + +The per-transform `retain_stats` flag puts a channel's mean and standard deviation back +after one transform. This is the same idea for the whole pipeline, with two more +statistics and a jitter: the profile is snapshotted once after the geometry and restored +once at the end, so the augmentations are free to do what they like in between and the +network still sees an intensity profile that is plausible -- but never exactly the one the +patch arrived with. + +```json +{ + "pipeline": { + "retain_stats": { + "stats": ["mean", "std", "p1", "p99"], + "percentiles": [1.0, 99.0], + "std_scale": [0.9, 1.1], + "percentile_scale": [0.9, 1.1], + "stats_from": "all", + "p": 1.0 + } + } +} +``` + +An absent block is off and spends no random numbers. A block that is present but empty is +on with every default, which restores the snapshot exactly. + +**Two of the four, drawn per batch.** The snapshot is the per-sample, per-channel mean, +standard deviation, 1st and 99th percentile. An affine map `y = ax + b` has two degrees of +freedom, so all four cannot be restored at once -- each pass therefore draws **two** of +them and solves for the map that puts exactly those two on target, leaving the other two +to fall where the augmentations left them. The six pairs span the two useful extremes: + +| drawn pair | what it pins | what it is | +| --- | --- | --- | +| mean, std | the moments | what per-transform `retain_stats` does | +| p1, p99 | the robust dynamic range | immune to the outliers an edge filter or a random convolution manufactures, which are exactly what make the standard deviation a bad target | +| mean, p1 / mean, p99 / std, p1 / std, p99 | one of each | the mixed cases in between | + +`stats` is the pool the pair is drawn from, so `["mean", "std"]` reproduces the old +behaviour at pipeline level and the default draws from all six pairs. + +**The jitter.** `std_scale` multiplies the standard-deviation target; `percentile_scale` +multiplies each percentile's distance from the mean, independently, so the span can widen, +narrow or go asymmetric. Both are a scalar or a `[lo, hi]` range, and `1.0` restores the +snapshot exactly. The mean target is never jittered -- it is the anchor every pair shares, +and moving it as well would put two unrelated augmentations behind one knob. `p` is the +per-sample probability that the restore applies at all. + +**Where it sits.** With the block on, the pipeline is + +``` +[geometry] [snapshot] [augmentations] [restore] [crop] [normalization] +``` + +in all three pipeline modes. Two groups move to the tail: the crop, because it occludes a +random box with the volume's own minimum rather than changing intensities, so a restore +measured before it would be undoing the occlusion; and the normalization stage, because +deciding the scale the network sees is its entire job. In `random_order` mode the snapshot +lands after the geometry but *before* the hoisted `RandomLowResTransformGPU`, which is a +general enhancement that merely runs early. + +A config should pick one mechanism or the other: the two do not compose, and a transform +that restores its own moments inside the bracket is doing work the restore will redo. +`smauglab/configs/transform_params_gpu_stats-restore.json` is the paper's `Ours` setup with +the bracket on and every per-transform `retain_stats` turned off. + ### Paper configs Two shipped configs are the setups behind the paper's tables. `smauglab hash ` diff --git a/smauglab/config.py b/smauglab/config.py index 93bd7ca..892a2ed 100644 --- a/smauglab/config.py +++ b/smauglab/config.py @@ -6,7 +6,7 @@ "_comment": "...", // any key starting with '_' is ignored "GPU": { "": { }, ... }, "CPU": { ... }, - "pipeline": { "random_choose": { ... } } + "pipeline": { "random_choose": { ... }, "retain_stats": { ... } } } A key is a class name, exactly, and a parameter is a constructor argument, exactly. Both are @@ -96,7 +96,7 @@ def validate_section(section: dict, backend: Backend, *, source: str = " RESERVED_SECTIONS = ("pipeline",) #: Keys the `pipeline` section may hold. -PIPELINE_KEYS = ("mode", "order", "random_choose", "same_on_batch") +PIPELINE_KEYS = ("mode", "order", "random_choose", "retain_stats", "same_on_batch") #: What to do about a config that has no backend section. Quoted by both errors that #: can report it, so the two cannot drift apart. @@ -137,6 +137,22 @@ def pipeline_options(self, name: str) -> dict[str, Any]: """ return copy.deepcopy(self.payload.get("pipeline", {}).get(name, {})) + def retain_stats_spec(self) -> Any: + """The parsed `pipeline.retain_stats` block, or None when there is none. + + Returns a `gpu.stats.StatsSpec`, imported here rather than at module scope: that + module reaches kornia, and a CPU-only caller validating a config should not have + to. Typed loosely for the same reason -- `build_gpu_pipeline` has the real + annotation under `TYPE_CHECKING`. + """ + from smauglab.transforms.gpu.stats import parse_spec + + # Read off the payload rather than through `pipeline_options`, which cannot tell + # an absent block from an empty one: it always hands back a dict, and here the + # difference is the difference between off and on-with-defaults. + block = self.payload.get("pipeline", {}).get("retain_stats") + return parse_spec(copy.deepcopy(block)) + def pipeline_mode(self) -> PipelineMode: """How the GPU pipeline arranges its transforms. @@ -239,6 +255,14 @@ def _pipeline_problems(self) -> list[str]: if same_on_batch is not None and not isinstance(same_on_batch, bool): problems.append(f"pipeline.same_on_batch: expected true or false, got {same_on_batch!r}") + retain_stats = pipeline.get("retain_stats") + if retain_stats is not None: + # Imported inside the check so a config with no statistics bracket still + # validates without importing the GPU transform stack. + from smauglab.transforms.gpu.stats import spec_problems + + problems.extend(spec_problems(retain_stats)) + mode = pipeline.get("mode") if mode is not None: valid = [m.value for m in PipelineMode] diff --git a/smauglab/configs/transform_params_gpu_stats-restore.json b/smauglab/configs/transform_params_gpu_stats-restore.json new file mode 100644 index 0000000..65f8972 --- /dev/null +++ b/smauglab/configs/transform_params_gpu_stats-restore.json @@ -0,0 +1,263 @@ +{ + "_comment": "Ours, with the per-transform retain_stats replaced by the pipeline-level statistics bracket. The snapshot runs after the geometry and the restore before the trailing crop and normalization: each batch draws two of mean/std/p1/p99 per channel and solves for the affine map that puts those two back, with the std and the two percentile targets scaled by a factor drawn from [0.9, 1.1] so the profile the network sees is never exactly the one the patch arrived with. Every transform's own retain_stats is false here: the two mechanisms do not compose, and the pipeline-level one is the point of this config. Derived from transform_params_one-sequence-to-segment-them-all.json, which is unchanged.", + "GPU": { + "RandomFlipTransformGPU": { + "flip_axis": [ + 0 + ], + "same_on_batch": false, + "keepdim": true, + "p": 0.5 + }, + "RandomAffineGPU": { + "degrees": 5, + "translate": [ + 0.1, + 0.1, + 0.1 + ], + "scale": [ + 0.7, + 1.4 + ], + "shears": [ + -5, + 5, + -5, + 5, + -5, + 5 + ], + "resample": "bilinear", + "p": 0.0 + }, + "RandomInverseGPU": { + "retain_stats": false, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "mix_prob": 0.0, + "p": 0.25 + }, + "RandomHistogramEqualizationGPU": { + "retain_stats": false, + "in_seg": 0, + "out_seg": 0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "mix_prob": 0.6, + "p": 0.25 + }, + "RandomRedistributeSegGPU": { + "confine_prob": 0.25, + "retain_stats": false, + "p": 0.4 + }, + "RandomScharrGPU": { + "absolute": true, + "retain_stats": false, + "mix_prob": 0.5, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "p": 0.25 + }, + "RandomUnsharpMaskGPU": { + "sigma": 1.0, + "unsharp_amount": 1.5, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "mix_prob": 0.5, + "p": 0.25 + }, + "RandomRandConvGPU": { + "kernel_sizes": [ + 3, + 5, + 7 + ], + "retain_stats": false, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "mix_prob": 0.5, + "p": 0.2 + }, + "RandomClampGPU": { + "max_clamp_amount": 0.2, + "retain_stats": false, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "p": 0.0 + }, + "RandomGaussianNoiseGPU": { + "mean": 0.0, + "std": 1.0, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "p": 0.1 + }, + "RandomGaussianBlurGPU": { + "sigma": 1.0, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "p": 0.2 + }, + "RandomBrightnessGPU": { + "brightness_range": [ + 0.75, + 1.25 + ], + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "p": 0.15 + }, + "RandomGammaGPU": { + "gamma_range": [ + 0.7, + 1.5 + ], + "retain_stats": false, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "p": 0.3 + }, + "RandomInvGammaGPU": { + "gamma_range": [ + 0.7, + 1.5 + ], + "retain_stats": false, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "p": 0.1 + }, + "RandomContrastGPU": { + "contrast_range": [ + 0.75, + 1.25 + ], + "retain_stats": false, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "p": 0.15 + }, + "RandomLog1pGPU": { + "retain_stats": false, + "in_seg": 0, + "out_seg": 0, + "p": 0.1 + }, + "RandomSqrtGPU": { + "retain_stats": false, + "in_seg": 0, + "out_seg": 0, + "p": 0.1 + }, + "RandomSinGPU": { + "retain_stats": false, + "in_seg": 0, + "out_seg": 0, + "p": 0.1 + }, + "RandomExpGPU": { + "retain_stats": false, + "in_seg": 0, + "out_seg": 0, + "p": 0.1 + }, + "RandomSigmoidGPU": { + "retain_stats": false, + "in_seg": 0, + "out_seg": 0, + "p": 0.1 + }, + "RandomLowResTransformGPU": { + "scale": [ + 0.5, + 1.0 + ], + "same_on_batch": false, + "p": 0.25 + }, + "RandomAcqTransformGPU": { + "scale": [ + 0.6, + 1.0 + ], + "same_on_batch": false, + "p": 0.0 + }, + "RandomBiasFieldGPU": { + "retain_stats": false, + "coefficients": 0.2, + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": "label", + "seg_unit_prob": 0.5, + "p": 0.15 + }, + "ZscoreNormalizationGPU": { + "p": 0.0 + } + }, + "CPU": { + "SpatialTransform": { + "patch_center_dist_from_border": 80, + "random_crop": false, + "p_elastic_deform": 0.0, + "p_rotation": 0.2, + "p_scaling": 0.2, + "scaling": [ + 0.7, + 1.4 + ], + "p_synchronize_scaling_across_axes": 1, + "bg_style_seg_sampling": false, + "mode_seg": "nearest" + } + }, + "pipeline": { + "retain_stats": { + "stats": [ + "mean", + "std", + "p1", + "p99" + ], + "percentiles": [ + 1.0, + 99.0 + ], + "std_scale": [ + 0.9, + 1.1 + ], + "percentile_scale": [ + 0.9, + 1.1 + ], + "stats_from": "all", + "p": 1.0 + } + } +} diff --git a/smauglab/trainers/nnUNetTrainerDAExt.py b/smauglab/trainers/nnUNetTrainerDAExt.py index a9dac85..cb7b42b 100644 --- a/smauglab/trainers/nnUNetTrainerDAExt.py +++ b/smauglab/trainers/nnUNetTrainerDAExt.py @@ -194,6 +194,7 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic options=config.pipeline_options("random_choose"), source=config.source, order_source=config.order_source(), + retain_stats=config.retain_stats_spec(), ), data_keys=["input", "mask"], # See AugTransformsGPU: True forces the flag onto every child and diff --git a/smauglab/transforms/build.py b/smauglab/transforms/build.py index d0a906e..5e214fc 100644 --- a/smauglab/transforms/build.py +++ b/smauglab/transforms/build.py @@ -14,11 +14,14 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any from smauglab import registry from smauglab.config import NON_AUGMENTATION_KEYS, OrderSource, PipelineMode, validate_section -from smauglab.registry import AugEntry, AugType, Backend, InvalidConfigError +from smauglab.registry import AugEntry, AugId, AugType, Backend, InvalidConfigError + +if TYPE_CHECKING: # the spec is a plain dataclass, but importing it pulls in kornia + from smauglab.transforms.gpu.stats import StatsSpec def _instantiate(entry: AugEntry, params: dict, context: dict[str, Any]) -> Any: @@ -142,6 +145,24 @@ def _normalization_stage(normalization: list, *, options: dict[str, Any]) -> lis ] +def _after_geometry(core: list[Any], geometry: list[Any], snapshot: Any) -> list[Any]: + """`core` with `snapshot` inserted after the last geometric transform it holds. + + Positions are found by identity rather than by index arithmetic because the three + pipeline modes assemble `core` differently -- geometry leads in all of them, but what + follows it is a bucket in two and a flat run in the third -- and because + `pipeline.order: config` can put a geometric transform anywhere. Insert-after-the-last + is the one rule that reads the same in every case. + + In `random_order` mode this deliberately lands *before* the hoisted `force_sequential` + entry: `RandomLowResTransformGPU` is a general enhancement that happens to run early, + and blurring a patch moves its standard deviation, so it belongs inside the bracket. + """ + identities = {id(transform) for transform in geometry} + cut = max((index for index, transform in enumerate(core) if id(transform) in identities), default=-1) + 1 + return [*core[:cut], snapshot, *core[cut:]] + + def build_gpu_pipeline( section: dict, *, @@ -149,51 +170,79 @@ def build_gpu_pipeline( options: dict[str, Any] | None = None, source: str = "", order_source: OrderSource = OrderSource.REGISTRY, + retain_stats: StatsSpec | None = None, ) -> list[Any]: - """The GPU transform list for a given pipeline mode.""" + """The GPU transform list for a given pipeline mode. + + `retain_stats` is the parsed `pipeline.retain_stats` block. When it is given, the list + gains the statistics bracket from `gpu/stats.py` and two groups move to the tail: + + [geometry] [SNAPSHOT] [augmentations] [RESTORE] [crop] [normalization] + + Crop goes after the restore because it is not an intensity augmentation at all -- it + occludes a random box with the volume's own minimum, so restoring a profile measured + before it would be undoing the occlusion rather than the augmentations. Normalization + goes last, as it already does, because its entire job is to decide the scale the + network sees; a restore after it would overwrite that decision. With `retain_stats` + absent the list is exactly what it was, including the order a lone normalization sits + in and the random numbers the pipeline spends. + """ built = build_transforms(section, Backend.GPU, source=source, order_source=order_source) options = options or {} + + bracket = None + staged: list[Any] = [] + if retain_stats is not None: + from smauglab.transforms.gpu.stats import make_bracket + + bracket = make_bracket(retain_stats) + staged = [transform for transform, entry in built if entry.aug_id is AugId.CROP] + built = [(transform, entry) for transform, entry in built if entry.aug_id is not AugId.CROP] + geometry = [transform for transform, entry in built if entry.group is AugType.GEO] + if mode is PipelineMode.SEQUENTIAL: - # 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: + if bracket is None and len(normalization) < 2: + # 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. 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 + core = [transform for transform, entry in built if entry.group is not AugType.NORM] + else: + from smauglab.transforms.gpu.transforms_list import RandomChooseXTransformsGPU - sequential, transfer, enhancement, normalization = _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( - RandomChooseXTransformsGPU( - transforms_list=transfer, - num_transforms=len(transfer), - p=options.get("ta_probability", 1.0), - random_order=options.get("ta_random_order", True) if mode is PipelineMode.RANDOM_ORDER else False, - label="ta", - ) - ) - if mode is PipelineMode.RANDOM_ORDER: - out.append( + core = list(sequential) + core.append( RandomChooseXTransformsGPU( - transforms_list=enhancement, - num_transforms=len(enhancement), - p=options.get("ge_probability", 1.0), - random_order=options.get("ge_random_order", True), - label="ge", + transforms_list=transfer, + num_transforms=len(transfer), + p=options.get("ta_probability", 1.0), + random_order=options.get("ta_random_order", True) if mode is PipelineMode.RANDOM_ORDER else False, + label="ta", ) ) - else: - # RANDOM_ORDER_TA shuffles only the transfer augmentations; the general - # enhancements stay in pipeline order, as they always have. - out.extend(enhancement) + if mode is PipelineMode.RANDOM_ORDER: + core.append( + RandomChooseXTransformsGPU( + transforms_list=enhancement, + num_transforms=len(enhancement), + p=options.get("ge_probability", 1.0), + random_order=options.get("ge_random_order", True), + label="ge", + ) + ) + else: + # RANDOM_ORDER_TA shuffles only the transfer augmentations; the general + # enhancements stay in pipeline order, as they always have. + core.extend(enhancement) + # Last in every mode: it standardizes whatever everything above it did. - out.extend(_normalization_stage(normalization, options=options)) - return out + if bracket is None: + return core + _normalization_stage(normalization, options=options) + snapshot, restore = bracket + return [*_after_geometry(core, geometry, snapshot), restore, *staged, *_normalization_stage(normalization, options=options)] def build_cpu_pipeline( diff --git a/smauglab/transforms/gpu/normalization.py b/smauglab/transforms/gpu/normalization.py index db188f1..f455be9 100644 --- a/smauglab/transforms/gpu/normalization.py +++ b/smauglab/transforms/gpu/normalization.py @@ -32,7 +32,7 @@ class owns the clone, the channel guard, the statistics mask, both degeneracy gu from __future__ import annotations from collections.abc import Sequence -from typing import Any +from typing import Any, Protocol import torch from torch import Tensor @@ -66,7 +66,23 @@ def _broadcast(values: Tensor, like: Tensor) -> Tensor: return values.view(-1, *([1] * (like.dim() - 1))) -def _resolve_mask(transform: _NormalizationBaseGPU, channel: Tensor, seg_mask: Tensor | None) -> Tensor | None: +class _MaskSource(Protocol): + """What `_resolve_mask` needs off the transform it is handed. + + A Protocol rather than the base class below, for the same reason `Mixing` in + gpu/mixing.py is one: the pipeline-level statistics bracket in gpu/stats.py resolves + its mask with the same rules and the same two settings, and it is not a normalization. + Read-only members, so a plain attribute satisfies them. + """ + + @property + def stats_from(self) -> str: ... + + @property + def nonzero_threshold(self) -> float: ... + + +def _resolve_mask(transform: _MaskSource, 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 diff --git a/smauglab/transforms/gpu/stats.py b/smauglab/transforms/gpu/stats.py new file mode 100644 index 0000000..f7eed0c --- /dev/null +++ b/smauglab/transforms/gpu/stats.py @@ -0,0 +1,437 @@ +"""Snapshot the image's intensity profile before the augmentations, restore it after. + +The per-transform `retain_stats` flag does this one transform at a time: ~15 GPU +transforms each measure their channel's mean and standard deviation, do their work, and +rescale the result back. Three things are wrong with that. It restores once per +transform, so a chain of them pays for N restores to end up where the last one puts it. +It knows only two statistics, so a transform that manufactures extreme voxels +(`RandomInverseGPU`, `RandomScharrGPU`, `RandomRandConvGPU`) drags the standard deviation +and the restore with it. And it restores the input's statistics *exactly*, which teaches +the network that its input scale is reliable -- the opposite of what a domain-generalising +augmentation pipeline is for. + +This module is the pipeline-level replacement: one snapshot after the geometry, one +restore after the augmentations, four statistics instead of two, and a jitter on the +targets so the profile the network sees is never quite the one the patch arrived with. +`smauglab.transforms.build` emits the pair around the stack; see `build_gpu_pipeline`. + +**Two of four.** The snapshot is the per-sample, per-channel mean, standard deviation, +1st and 99th percentile. All four cannot be restored at once: an affine map `y = ax + b` +has two degrees of freedom, and four targets over-determine it. So each pass *draws two* +of the four and solves for the map that puts exactly those on target. That turns the +over-determination into a feature -- the drawn pair is itself an augmentation, and the +six pairs cover the whole range from "keep the moments" (mean, std -- what `retain_stats` +did) to "keep the robust dynamic range" (p1, p99 -- immune to the outliers that make the +moments a bad target after an edge filter). + +Writing each statistic as a linear functional of the map, `t = a * u + b * c` with `c = 1` +for the three that move with an offset and `c = 0` for the standard deviation, which is +offset-invariant, makes all six pairs one 2x2 solve: + + det = u1 * c2 - u2 * c1 + a = (t1 * c2 - t2 * c1) / det + b = t1 - a * u1 (or t2 - a * u2 when the first drawn statistic is the std) + +`det` is a difference of intensities -- a span between two anchors, or the standard +deviation itself -- so a sample whose span has collapsed (a constant patch, or +`p1 == p99` on a nearly empty mask) is left exactly as it is, the same decision +`ZscoreNormalizationGPU` and the normalization schemes make for the same reason. + +**What spends random numbers.** With no `pipeline.retain_stats` block nothing here is +built and the RNG stream is untouched. With one, the restore spends a fixed five uniforms +per channel -- pair, std scale, two percentile scales, gate -- whichever pair comes out +and whatever the ranges are, so two runs that differ only in the configured jitter stay +comparable at a fixed seed. That is invariant 2 of `gpu/regions.py`: the number of draws +is a function of the configuration, never of the data. +""" + +from __future__ import annotations + +import itertools +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch +from torch import Tensor + +from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from +from smauglab.transforms.gpu.contrast import _check_channel +from smauglab.transforms.gpu.mixing import normalise_blend_range +from smauglab.transforms.gpu.normalization import STATS_SOURCES, _broadcast, _masked_moments, _masked_quantiles, _resolve_mask + +#: The four statistics, in the one canonical order. The index into this tuple is what the +#: drawn pairs, the kinds below and the `[4, N]` stacks all agree on. +STAT_NAMES = ("mean", "std", "p1", "p99") + +#: Each statistic's coefficient on the affine map's offset: 1 for the three that move with +#: it, 0 for the standard deviation, which does not. This is the whole reason one solver +#: serves all six pairs. +_STAT_OFFSET_COEFFICIENTS = (1.0, 0.0, 1.0, 1.0) + +#: Default percentile pair, as percentages. The 1st and 99th are what the request asked +#: for and are far enough from the extremes to be a robust span on a patch of any size. +DEFAULT_PERCENTILES = (1.0, 99.0) + +#: Default jitter: none. An absent `pipeline.retain_stats` block is off, and a present one +#: that names no scale restores the snapshot exactly -- the pipeline-level equivalent of +#: `retain_stats: true`, which is the behaviour a config inherits if it asks for nothing. +DEFAULT_SCALE = (1.0, 1.0) + +#: A span below this fraction of the sample's own range makes the solve meaningless. Same +#: value and same intent as `min_scale` in gpu/normalization.py. +_DET_FLOOR = 1e-3 + +#: Sign-preserving floor on the determinant, so the division is finite even for the rows +#: the guard is about to discard anyway. +_EPS = 1e-8 + + +@dataclass(frozen=True) +class StatsSpec: + """A parsed `pipeline.retain_stats` block. + + Built by `parse_spec` from the config, and frozen because both leaves of one bracket + read the same instance: a restore that disagreed with its snapshot about which + percentiles were measured would silently restore the wrong thing. + """ + + stats: tuple[str, ...] = STAT_NAMES + percentiles: tuple[float, float] = DEFAULT_PERCENTILES + std_scale: tuple[float, float] = DEFAULT_SCALE + percentile_scale: tuple[float, float] = DEFAULT_SCALE + stats_from: str = "all" + nonzero_threshold: float = 1e-6 + p: float = 1.0 + #: Which channels the bracket covers. None -- the default -- is every channel, which + #: is what makes this a property of the image rather than of one transform's + #: `apply_to_channel`. A multi-sequence input gets an independent draw per channel. + apply_to_channel: tuple[int, ...] | None = None + + @property + def pairs(self) -> tuple[tuple[int, int], ...]: + """Index pairs into `STAT_NAMES`: the table the per-batch draw indexes into.""" + return tuple(itertools.combinations((STAT_NAMES.index(n) for n in self.stats), 2)) + + @property + def probabilities(self) -> tuple[float, float]: + """The percentile pair as quantile probabilities, for `_masked_quantiles`.""" + return (self.percentiles[0] / 100.0, self.percentiles[1] / 100.0) + + +#: Keys a `pipeline.retain_stats` block may hold. Checked by `smauglab/config.py` so a +#: typo is reported with the rest of a config's problems rather than at build time. +SPEC_KEYS = ("stats", "percentiles", "std_scale", "percentile_scale", "stats_from", "nonzero_threshold", "p", "apply_to_channel") + +_WHERE = "pipeline.retain_stats" + + +def parse_spec(options: dict[str, Any] | None) -> StatsSpec | None: + """A validated `StatsSpec`, or None when the feature is off. + + Off means the block is *absent*, and then `build_gpu_pipeline` emits no leaves at all + -- not a pair of no-ops. The difference is visible in the RNG stream. A block that is + present but empty is on, with every default, which restores the drawn pair exactly; + that is the one reading of `"retain_stats": {}` an author could mean. + + Value checks live here rather than in `smauglab/config.py`, which only ever asks + whether a key is accepted; that is where every other parameter in this package is + checked too. A pipeline is built before the first batch, so a bad range still fails + long before training. + """ + if options is None: + return None + + names = options.get("stats", STAT_NAMES) + if isinstance(names, str) or not isinstance(names, Sequence): + # ValueError, not the TypeError TRY004 asks for: every rejection in this parser is + # one, and a caller should only have to catch one thing for a bad block. + raise ValueError(f"{_WHERE}.stats must be a list of names from {', '.join(STAT_NAMES)}, got {names!r}") # noqa: TRY004 + unknown = [n for n in names if n not in STAT_NAMES] + if unknown: + raise ValueError(f"{_WHERE}.stats: unknown name(s) {', '.join(map(repr, unknown))}. Accepted: {', '.join(STAT_NAMES)}.") + # Canonicalised, not taken in the order written: the pair table is indexed by a draw, + # so the order of the keys in a config would otherwise decide which pair a seed picks. + stats = tuple(n for n in STAT_NAMES if n in set(names)) + if len(stats) < 2: + raise ValueError(f"{_WHERE}.stats needs at least two statistics to determine an affine map, got {list(names)!r}") + + percentiles = options.get("percentiles", DEFAULT_PERCENTILES) + try: + low, high = (float(v) for v in percentiles) + except (TypeError, ValueError) as exc: + raise ValueError(f"{_WHERE}.percentiles must be a [low, high] pair in percent, got {percentiles!r}") from exc + if not 0.0 <= low < high <= 100.0: + raise ValueError(f"{_WHERE}.percentiles must satisfy 0 <= low < high <= 100, got [{low}, {high}]") + + stats_from = options.get("stats_from", "all") + if stats_from not in STATS_SOURCES: + raise ValueError(f"{_WHERE}.stats_from must be one of {STATS_SOURCES}, got {stats_from!r}.") + + probability = options.get("p", 1.0) + if isinstance(probability, bool) or not isinstance(probability, (int, float)) or not 0.0 <= float(probability) <= 1.0: + raise ValueError(f"{_WHERE}.p must be a number in [0, 1], got {probability!r}") + + channels = options.get("apply_to_channel") + if channels is not None: + if isinstance(channels, (str, bool)) or not isinstance(channels, Sequence) or not channels: + raise ValueError(f"{_WHERE}.apply_to_channel must be a non-empty list of channel indices, got {channels!r}") + if any(isinstance(c, bool) or not isinstance(c, int) or c < 0 for c in channels): + raise ValueError(f"{_WHERE}.apply_to_channel must hold non-negative integers, got {list(channels)!r}") + channels = tuple(int(c) for c in channels) + + threshold = float(options.get("nonzero_threshold", 1e-6)) + if threshold <= 0.0: + raise ValueError(f"{_WHERE}.nonzero_threshold must be positive, got {threshold}") + + hint = "It scales the restored target, so 1 restores the snapshot exactly." + return StatsSpec( + stats=stats, + percentiles=(low, high), + std_scale=normalise_blend_range( + options.get("std_scale", DEFAULT_SCALE), key="std_scale", owner=_WHERE, hint=hint, bounds=(0.0, None) + ), + percentile_scale=normalise_blend_range( + options.get("percentile_scale", DEFAULT_SCALE), key="percentile_scale", owner=_WHERE, hint=hint, bounds=(0.0, None) + ), + stats_from=stats_from, + nonzero_threshold=threshold, + p=float(probability), + apply_to_channel=channels, + ) + + +def spec_problems(options: Any) -> list[str]: + """Every problem in a `pipeline.retain_stats` block, for cumulative config reporting.""" + if options is None: + return [] + if not isinstance(options, dict): + return [f"{_WHERE}: expected an object"] + problems = [f"{_WHERE}: unknown key {key!r}. Accepted: {', '.join(SPEC_KEYS)}." for key in options if key not in SPEC_KEYS] + try: + parse_spec(options) + except ValueError as exc: + problems.append(str(exc)) + return problems + + +@dataclass +class _ChannelSnapshot: + """One channel's four statistics, and the voxels they were read from.""" + + #: `[4, N]`, in `STAT_NAMES` order, always float32. + values: Tensor + #: The resolved statistics mask, or None for "all voxels". Kept rather than recomputed + #: at restore time so both ends read the *same* voxel set: with `stats_from="nonzero"` + #: the augmentations themselves move which voxels pass the threshold, and a restore + #: measured over a different population is not a restore. + mask: Tensor | None + + +class StatsSlot: + """The handoff between the two leaves of one bracket. + + A plain object, not a buffer and not a module: the snapshot is per batch, lives for + one forward pass, and has no business being moved by `.to(device)` or saved in a + state dict. One instance per pipeline, held by both leaves. + """ + + __slots__ = ("channels",) + + def __init__(self) -> None: + self.channels: dict[int, _ChannelSnapshot] | None = None + + +def measure(channel: Tensor, mask: Tensor | None, probabilities: tuple[float, float]) -> Tensor: + """The four statistics of a `[N, *spatial]` channel, stacked `[4, N]` and float32. + + float32 whatever the input is: the GPU pipeline runs inside `autocast` + (`nnUNetTrainerDAExt.train_step`), and a half-precision sum over a 128^3 patch loses + enough of the mean to move the restored level visibly. Same reason the region moments + are pinned to fp32. + """ + work = channel.float() + mean, std = _masked_moments(work, mask) + low, high = _masked_quantiles(work, mask, probabilities) + return torch.stack((mean, std, low, high)) + + +def _lerp(bounds: tuple[float, float], draw: Tensor) -> Tensor: + """A uniform draw mapped onto `[lo, hi]`. Exact at `lo == hi`, and always consumed.""" + low, high = bounds + return draw * (high - low) + low + + +def _targets(snapshot: Tensor, spec: StatsSpec, std_draw: Tensor, low_draw: Tensor, high_draw: Tensor) -> Tensor: + """The four jittered targets, `[4, N]`. + + The mean is never jittered: it is the one statistic every pair can anchor on, and + shifting it as well would make the drawn pair decide the output's level as well as its + scale -- two unrelated augmentations behind one knob. The percentiles are jittered on + their *distance from the mean*, independently, so the span can widen, narrow or go + asymmetric while staying a span that contains the mean. + """ + mean, std, low, high = snapshot + return torch.stack( + ( + mean, + std * _lerp(spec.std_scale, std_draw), + mean - (mean - low) * _lerp(spec.percentile_scale, low_draw), + mean + (high - mean) * _lerp(spec.percentile_scale, high_draw), + ) + ) + + +def _solve(sources: Tensor, targets: Tensor, pairs: Tensor, spread: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """`(alpha, beta, valid)` for the drawn pair of each sample. All `[N]`. + + `sources` and `targets` are `[4, N]`; `pairs` is `[N, 2]` of indices into them. + """ + kinds = torch.tensor(_STAT_OFFSET_COEFFICIENTS, device=sources.device, dtype=sources.dtype) + rows = torch.arange(sources.shape[1], device=sources.device) + first, second = pairs[:, 0], pairs[:, 1] + + u1, u2 = sources[first, rows], sources[second, rows] + t1, t2 = targets[first, rows], targets[second, rows] + c1, c2 = kinds[first], kinds[second] + + determinant = u1 * c2 - u2 * c1 + # Sign-preserving, as in gpu/normalization.py: a negative determinant is ordinary (the + # p1/p99 pair always has one) and clamping its magnitude must not flip it. + safe = torch.where(determinant.abs() < _EPS, torch.full_like(determinant, _EPS).copysign(determinant), determinant) + alpha = (t1 * c2 - t2 * c1) / safe + # Whichever of the two carries the offset; at most one of them can be the std. + beta = torch.where(c1 != 0, t1 - alpha * u1, t2 - alpha * u2) + + valid = (spread > 0) & (determinant.abs() >= _DET_FLOOR * spread) & torch.isfinite(alpha) & torch.isfinite(beta) + return alpha, beta, valid + + +class StatsSnapshotGPU(ImageOnlyTransform): + """Records the intensity profile the augmentations are about to change. Not an augmentation. + + Returns its input untouched -- it only reads -- and carries no probability worth + drawing: `p=1.0`, because a snapshot that sometimes did not happen would leave the + restore at the other end of the pipeline with nothing to restore. + + Pipeline structure rather than a registered augmentation, like the 2D/3D converters + `build_cpu_pipeline` emits around `SpatialTransform`: its position is not a choice a + config makes, it is implied by where the geometry ends. + """ + + spec: StatsSpec + slot: StatsSlot + stats_from: str + nonzero_threshold: float + + def __init__(self, spec: StatsSpec, slot: StatsSlot, keepdim: bool = True) -> None: + super().__init__(p=1.0, p_batch=1.0, same_on_batch=False, keepdim=keepdim) + self.spec = spec + self.slot = slot + # Read by `_resolve_mask`; see the `_MaskSource` protocol in gpu/normalization.py. + self.stats_from = spec.stats_from + self.nonzero_threshold = spec.nonzero_threshold + + def channels(self, total: int) -> tuple[int, ...]: + """Which channels this bracket covers, with every index checked against the input.""" + selected = self.spec.apply_to_channel if self.spec.apply_to_channel is not None else tuple(range(total)) + for channel in selected: + _check_channel(self, channel, total) + return selected + + @torch.no_grad() + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: + seg_mask = segmentation_from(params) + snapshot: dict[int, _ChannelSnapshot] = {} + for channel in self.channels(input.shape[1]): + data = input[:, channel] + mask = _resolve_mask(self, data, seg_mask) + snapshot[channel] = _ChannelSnapshot(values=measure(data, mask, self.spec.probabilities), mask=mask) + self.slot.channels = snapshot + return input + + +class StatsRestoreGPU(ImageOnlyTransform): + """Puts two of the four snapshotted statistics back, jittered. Not an augmentation. + + `p=1.0` on the transform; the per-sample probability the config asks for is + `StatsSpec.p` and is spent here, so that a sample which is not restored still costs + the same draws as one that is. + """ + + spec: StatsSpec + slot: StatsSlot + + def __init__(self, spec: StatsSpec, slot: StatsSlot, keepdim: bool = True) -> None: + super().__init__(p=1.0, p_batch=1.0, same_on_batch=False, keepdim=keepdim) + self.spec = spec + self.slot = slot + self._pairs = torch.tensor(spec.pairs, dtype=torch.long) + + def _draw(self, rows: int, device: torch.device) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Pair, three jitter uniforms and the gate -- five draws, unconditionally. + + `torch.rand` scaled into the table rather than `torch.randint`, which is the + fixed-width trick `gpu/regions.py` uses for the same purpose: one uniform per row + whatever the pool's size, so adding a statistic to `stats` does not reshape the + stream. No `dtype=`, following gpu/mixing.py, so a half-precision batch draws from + the same stream as a single-precision one. + """ + count = self._pairs.shape[0] + drawn = (torch.rand(rows, device=device) * count).long().clamp_max_(count - 1) + return ( + drawn, + torch.rand(rows, device=device), + torch.rand(rows, device=device), + torch.rand(rows, device=device), + torch.rand(rows, device=device), + ) + + @torch.no_grad() + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: + taken = self.slot.channels + self.slot.channels = None # consumed: a second pass must not read this one's stats + if taken is None: + raise RuntimeError( + "StatsRestoreGPU ran without a snapshot. The two are emitted as a pair by " + "smauglab.transforms.build.build_gpu_pipeline and the snapshot has to run first." + ) + + batch = input.shape[0] + # One draw shared by the batch under `same_on_batch`, which the container forces + # onto every child at the shipped default -- the same way a bucket takes one + # decision for the whole batch. See SmaugConfig.same_on_batch. + rows = 1 if self.same_on_batch else batch + out = input.clone() + + for channel, recorded in taken.items(): + if recorded.values.shape[1] != batch: + raise RuntimeError( + f"StatsRestoreGPU got {batch} samples but the snapshot holds {recorded.values.shape[1]}. " + "Nothing between the two may change the batch size." + ) + data = out[:, channel] + work = data.float() + sources = measure(work, recorded.mask, self.spec.probabilities) + pair, std_draw, low_draw, high_draw, gate = self._draw(rows, input.device) + targets = _targets(recorded.values, self.spec, std_draw, low_draw, high_draw) + + flat = work.reshape(batch, -1) + spread = flat.amax(dim=1) - flat.amin(dim=1) + pairs = self._pairs.to(input.device)[pair].expand(batch, 2) + alpha, beta, valid = _solve(sources, targets, pairs, spread) + applies = valid & (gate < self.spec.p).expand(batch) + + restored = work * _broadcast(alpha, work) + _broadcast(beta, work) + # Per row, not per batch: one degenerate sample must not cost the others their + # restore. The guard is the same one `_select_and_check` applies globally. + finite = torch.isfinite(restored).reshape(batch, -1).all(dim=1) + out[:, channel] = torch.where(_broadcast(applies & finite, work), restored, work).to(data.dtype) + + return out + + +def make_bracket(spec: StatsSpec) -> tuple[StatsSnapshotGPU, StatsRestoreGPU]: + """The two leaves of one bracket, already sharing a slot.""" + slot = StatsSlot() + return StatsSnapshotGPU(spec, slot), StatsRestoreGPU(spec, slot) diff --git a/smauglab/transforms/gpu/transforms.py b/smauglab/transforms/gpu/transforms.py index c57ed3c..fe5aafe 100644 --- a/smauglab/transforms/gpu/transforms.py +++ b/smauglab/transforms/gpu/transforms.py @@ -24,6 +24,7 @@ def __init__(self, json_path: str): options=config.pipeline_options("random_choose"), source=config.source, order_source=config.order_source(), + retain_stats=config.retain_stats_spec(), ) # None, not False: kornia only overwrites a child's own `same_on_batch` when the # argument is not None, and every config sets the flag per transform. Passing True diff --git a/unit_tests/test_pipeline_stats.py b/unit_tests/test_pipeline_stats.py new file mode 100644 index 0000000..f56834e --- /dev/null +++ b/unit_tests/test_pipeline_stats.py @@ -0,0 +1,426 @@ +"""The pipeline-level statistics bracket: `pipeline.retain_stats`. + +What these pin, in the order the feature is assembled: + +* the solver really does put the two drawn statistics back, for each of the six pairs, + and leaves the other two free -- that is the whole premise, and it is the one thing a + smoke test could not tell apart from a no-op; +* where the builder emits the bracket, in all three pipeline modes: after the geometry, + and before the crop and the normalization that now follow the restore; +* the draw discipline the rest of the GPU package keeps -- a fixed number of uniforms per + channel, independent of the configured ranges and of how many statistics are in the + pool, so a jitter sweep is comparable against its own baseline at a fixed seed; +* the degenerate cases that have broken every stats restore in this repository before: + a constant patch, an unlabelled mask, a channel the solver cannot determine. +""" + +import json + +import torch + +from smauglab.config import PipelineMode, SmaugConfig +from smauglab.registry import AugType, Backend, InvalidConfigError, entries +from smauglab.transforms.build import build_gpu_pipeline +from smauglab.transforms.gpu.base import AugmentationSequentialCustom +from smauglab.transforms.gpu.contrast import RandomGammaGPU, _channel_stats, _restore_stats +from smauglab.transforms.gpu.stats import STAT_NAMES, StatsSpec, make_bracket, measure, parse_spec, spec_problems +from unit_tests.helpers import SmaugLabTestCase, configs_dir, first_output + +#: Every pair the full pool can draw, by name. +ALL_PAIRS = (("mean", "std"), ("mean", "p1"), ("mean", "p99"), ("std", "p1"), ("std", "p99"), ("p1", "p99")) + + +def restored(image, spec, *, seg=None, transform=None): + """Run `image` through a bracket, with one augmentation inside it by default.""" + snapshot, restore = make_bracket(spec) + inner = RandomGammaGPU(p=1.0) if transform is None else transform + keys = ["input", "mask"] if seg is not None else ["input"] + pipeline = AugmentationSequentialCustom(snapshot, inner, restore, data_keys=keys, same_on_batch=True) + out = pipeline(image.clone(), seg.clone()) if seg is not None else pipeline(image.clone()) + return first_output(out) + + +def stats_of(image, spec, channel=0, mask=None): + """The four statistics of one channel, as a `[4, N]` tensor.""" + return measure(image[:, channel], mask, spec.probabilities) + + +class TestTheSolverHitsTheDrawnPair(SmaugLabTestCase): + """Two statistics determine the map, and the solver puts exactly those two back.""" + + def setUp(self): + super().setUp() + # Off-centre and wider than [0, 1], so a map that merely happened to be the + # identity would be caught: every statistic here is far from its default. + self.image = torch.rand(2, 1, 24, 24, 24) * 4.0 - 1.0 + + def test_each_pair_is_restored_exactly(self): + for pair in ALL_PAIRS: + with self.subTest(pair=pair): + spec = parse_spec({"stats": list(pair)}) + out = restored(self.image, spec) + before, after = stats_of(self.image, spec), stats_of(out, spec) + for name in pair: + index = STAT_NAMES.index(name) + torch.testing.assert_close(after[index], before[index], rtol=1e-4, atol=1e-5, msg=f"{name} was not restored") + + def test_the_statistics_outside_the_pair_are_free(self): + """The counterpart: an over-determined restore would be a bug, not a feature.""" + moved = 0 + for pair in ALL_PAIRS: + spec = parse_spec({"stats": list(pair)}) + before, after = stats_of(self.image, spec), stats_of(restored(self.image, spec), spec) + for index, name in enumerate(STAT_NAMES): + if name not in pair: + moved += int(bool((after[index] - before[index]).abs().max() > 1e-3)) + self.assertGreater(moved, 0, "no untargeted statistic moved, so the restore is pinning more than two") + + def test_it_generalizes_the_per_transform_retain_stats(self): + """With the pool fixed to (mean, std) the bracket is `_restore_stats` over the pipeline. + + Compared loosely on purpose: `_restore_stats` divides by the sample standard + deviation (`unbiased=True`) and everything in this module uses the population one, + as nnU-Net and `ZscoreNormalizationGPU` do. On a 24^3 patch that is a factor of + 1.00004, which is the tolerance below and not a bug. + """ + spec = parse_spec({"stats": ["mean", "std"]}) + torch.manual_seed(7) + through_bracket = restored(self.image, spec) + + torch.manual_seed(7) + snapshot, restore = make_bracket(spec) + gamma = RandomGammaGPU(p=1.0) + pipeline = AugmentationSequentialCustom(snapshot, gamma, restore, data_keys=["input"], same_on_batch=True) + pipeline(self.image.clone()) # same draws, so the gamma below is the one it applied + + torch.manual_seed(7) + reference = first_output( + AugmentationSequentialCustom(RandomGammaGPU(p=1.0), data_keys=["input"], same_on_batch=True)(self.image.clone()) + ) + by_hand = _restore_stats(reference[:, 0], _channel_stats(self.image[:, 0])) + torch.testing.assert_close(through_bracket[:, 0], by_hand, rtol=1e-3, atol=1e-3) + + +class TestTheJitter(SmaugLabTestCase): + def setUp(self): + super().setUp() + self.image = torch.rand(1, 1, 24, 24, 24) * 2.0 + + def test_an_identity_range_restores_the_snapshot_exactly(self): + spec = parse_spec({"stats": ["mean", "std"], "std_scale": [1.0, 1.0]}) + before, after = stats_of(self.image, spec), stats_of(restored(self.image, spec), spec) + torch.testing.assert_close(after[:2], before[:2], rtol=1e-4, atol=1e-5) + + def test_the_std_lands_inside_the_configured_range(self): + spec = parse_spec({"stats": ["mean", "std"], "std_scale": [0.8, 1.25]}) + ratios = [] + for seed in range(25): + torch.manual_seed(seed) + ratios.append((stats_of(restored(self.image, spec), spec)[1] / stats_of(self.image, spec)[1]).item()) + self.assertGreaterEqual(min(ratios), 0.8 - 1e-4, f"below the configured range: {min(ratios)}") + self.assertLessEqual(max(ratios), 1.25 + 1e-4, f"above the configured range: {max(ratios)}") + self.assertGreater(max(ratios) - min(ratios), 0.1, "the range is configured but nothing is being drawn from it") + + def test_the_percentile_span_lands_inside_the_configured_range(self): + spec = parse_spec({"stats": ["p1", "p99"], "percentile_scale": [0.8, 1.25]}) + before = stats_of(self.image, spec) + spans = [] + for seed in range(25): + torch.manual_seed(seed) + after = stats_of(restored(self.image, spec), spec) + spans.append(((after[3] - after[2]) / (before[3] - before[2])).item()) + # Each end is drawn independently against its own distance from the mean, so the + # span is a convex combination of the two factors and stays inside their range. + self.assertGreaterEqual(min(spans), 0.8 - 1e-4, f"span below the range: {min(spans)}") + self.assertLessEqual(max(spans), 1.25 + 1e-4, f"span above the range: {max(spans)}") + + def test_the_mean_target_is_never_jittered(self): + spec = parse_spec({"stats": ["mean", "std"], "std_scale": [0.5, 2.0], "percentile_scale": [0.5, 2.0]}) + for seed in range(5): + with self.subTest(seed=seed): + torch.manual_seed(seed) + before, after = stats_of(self.image, spec), stats_of(restored(self.image, spec), spec) + torch.testing.assert_close(after[0], before[0], rtol=1e-4, atol=1e-5) + + def test_p_zero_leaves_the_image_where_the_augmentations_left_it(self): + spec = parse_spec({"stats": ["mean", "std"], "p": 0.0}) + torch.manual_seed(3) + gated = restored(self.image, spec) + torch.manual_seed(3) + ungated = restored(self.image, parse_spec({"stats": ["mean", "std"], "p": 1.0})) + self.assertFalse(bool(torch.allclose(gated, ungated)), "p=0 restored anyway") + before, after = stats_of(self.image, spec), stats_of(gated, spec) + self.assertGreater((after[0] - before[0]).abs().max().item(), 1e-4, "p=0 still put the mean back") + + +class TestTheDrawDiscipline(SmaugLabTestCase): + """Five uniforms per channel, whatever the configuration says.""" + + def _stream_after(self, spec, channels=1): + image = torch.rand(1, channels, 24, 24, 24) + _, restore = make_bracket(spec) + snapshot, _ = make_bracket(spec) + snapshot.slot = restore.slot + snapshot.apply_transform(image, {}, {}) + torch.manual_seed(11) + restore.apply_transform(image, {}, {}) + return torch.rand(3) + + def test_five_draws_per_channel(self): + spec = parse_spec({}) + torch.manual_seed(11) + torch.rand(5) + expected = torch.rand(3) + torch.testing.assert_close(self._stream_after(spec), expected) + + def test_the_count_does_not_depend_on_the_ranges(self): + wide = self._stream_after(parse_spec({"std_scale": [0.5, 2.0], "percentile_scale": [0.5, 2.0]})) + identity = self._stream_after(parse_spec({})) + torch.testing.assert_close(wide, identity) + + def test_the_count_does_not_depend_on_the_size_of_the_pool(self): + pair = self._stream_after(parse_spec({"stats": ["mean", "std"]})) + full = self._stream_after(parse_spec({})) + torch.testing.assert_close(pair, full) + + def test_two_channels_cost_twice(self): + spec = parse_spec({}) + torch.manual_seed(11) + torch.rand(10) + expected = torch.rand(3) + torch.testing.assert_close(self._stream_after(spec, channels=2), expected) + + def test_the_feature_off_emits_nothing_to_draw_with(self): + section = json.loads((configs_dir() / "transform_params_gpu.json").read_text())["GPU"] + built = [type(t).__name__ for t in build_gpu_pipeline(section)] + self.assertNotIn("StatsSnapshotGPU", built) + self.assertNotIn("StatsRestoreGPU", built) + + +class TestDegenerateInputs(SmaugLabTestCase): + def test_a_constant_patch_is_left_alone(self): + """No span, so no map is determined -- and the restore must not invent one.""" + image = self.constant_volume() + for pair in ALL_PAIRS: + with self.subTest(pair=pair): + out = restored(image, parse_spec({"stats": list(pair)}), transform=torch.nn.Identity()) + self.assertIsImageLike(out, image, f"a constant patch through {pair}") + torch.testing.assert_close(out, image) + + def test_an_unlabelled_mask_falls_back_to_the_whole_volume(self): + image = self.tiny_volume() + spec = parse_spec({"stats": ["mean", "std"], "stats_from": "seg"}) + torch.manual_seed(5) + empty = restored(image, spec, seg=self.empty_seg()) + torch.manual_seed(5) + everywhere = restored(image, parse_spec({"stats": ["mean", "std"], "stats_from": "all"}), seg=self.empty_seg()) + torch.testing.assert_close(empty, everywhere) + + def test_the_statistics_can_be_confined_to_the_labelled_anatomy(self): + image = self.tiny_volume() + seg = self.tiny_seg() + image = image + seg * 5.0 # foreground mean far from the whole-volume mean + spec = parse_spec({"stats": ["mean", "std"], "stats_from": "seg"}) + out = restored(image, spec, seg=seg) + + foreground = seg[:, 0] > 0 + before = image[:, 0][foreground].mean() + after = out[:, 0][foreground].mean() + torch.testing.assert_close(after, before, rtol=1e-4, atol=1e-4) + whole = (out[:, 0].mean() - image[:, 0].mean()).abs() + self.assertGreater(whole.item(), 1e-4, "the whole-volume mean was restored, so the mask was ignored") + + def test_a_restore_without_a_snapshot_is_an_error(self): + _, restore = make_bracket(parse_spec({})) + with self.assertRaises(RuntimeError) as caught: + restore.apply_transform(self.tiny_volume(), {}, {}) + self.assertIn("without a snapshot", str(caught.exception)) + + def test_a_changed_batch_size_is_an_error(self): + spec = parse_spec({}) + snapshot, restore = make_bracket(spec) + snapshot.apply_transform(torch.rand(2, 1, 8, 8, 8), {}, {}) + with self.assertRaises(RuntimeError) as caught: + restore.apply_transform(torch.rand(1, 1, 8, 8, 8), {}, {}) + self.assertIn("samples but the snapshot holds", str(caught.exception)) + + def test_the_slot_is_consumed(self): + spec = parse_spec({}) + snapshot, restore = make_bracket(spec) + snapshot.apply_transform(self.tiny_volume(), {}, {}) + restore.apply_transform(self.tiny_volume(), {}, {}) + self.assertIsNone(restore.slot.channels, "a second pass could read the first one's statistics") + + +class TestTensorContract(SmaugLabTestCase): + def test_the_caller_s_batch_is_not_mutated(self): + image = self.tiny_volume() + reference = image.clone() + restored(image, parse_spec({})) + torch.testing.assert_close(image, reference) + + def test_the_mask_comes_back_untouched(self): + seg = self.multi_label_seg() + snapshot, restore = make_bracket(parse_spec({})) + pipeline = AugmentationSequentialCustom(snapshot, RandomGammaGPU(p=1.0), restore, data_keys=["input", "mask"], same_on_batch=True) + _, mask = pipeline(self.tiny_volume(), seg.clone()) + torch.testing.assert_close(mask, seg) + + def test_both_kornia_routes_agree(self): + """An open autocast region sends the leaves down kornia's path instead of the fast one.""" + image = torch.rand(1, 1, 24, 24, 24) * 3.0 + spec = parse_spec({}) + torch.manual_seed(2) + direct = restored(image, spec, transform=torch.nn.Identity()) + torch.manual_seed(2) + with torch.autocast("cpu", dtype=torch.bfloat16): + through_kornia = restored(image, spec, transform=torch.nn.Identity()) + torch.testing.assert_close(direct, through_kornia, rtol=1e-5, atol=1e-5) + + def test_each_channel_is_restored_to_its_own_statistics(self): + image = torch.stack([torch.rand(1, 24, 24, 24), torch.rand(1, 24, 24, 24) * 10.0 - 5.0], dim=1) + spec = parse_spec({"stats": ["mean", "std"]}) + out = restored(image, spec) + for channel in (0, 1): + with self.subTest(channel=channel): + before, after = stats_of(image, spec, channel), stats_of(out, spec, channel) + torch.testing.assert_close(after[:2], before[:2], rtol=1e-4, atol=1e-4) + + def test_apply_to_channel_restricts_the_bracket(self): + image = torch.stack([torch.rand(1, 24, 24, 24), torch.rand(1, 24, 24, 24)], dim=1) + spec = parse_spec({"stats": ["mean", "std"], "apply_to_channel": [1], "std_scale": [2.0, 2.0]}) + snapshot, restore = make_bracket(spec) + snapshot.apply_transform(image, {}, {}) + out = restore.apply_transform(image, {}, {}) + torch.testing.assert_close(out[:, 0], image[:, 0], msg="channel 0 was touched") + self.assertGreater((out[:, 1] - image[:, 1]).abs().max().item(), 1e-3, "channel 1 was not touched") + + def test_an_out_of_range_channel_is_rejected(self): + snapshot, _ = make_bracket(parse_spec({"apply_to_channel": [3]})) + with self.assertRaises(IndexError): + snapshot.apply_transform(self.tiny_volume(), {}, {}) + + +class TestWhereTheBuilderPutsIt(SmaugLabTestCase): + """`[geometry] [snapshot] [augmentations] [restore] [crop] [normalization]`, every mode.""" + + def setUp(self): + super().setUp() + self.section = json.loads((configs_dir() / "transform_params_gpu.json").read_text())["GPU"] + self.spec = parse_spec({}) + self.geometry = {e.name for e in entries(Backend.GPU) if e.group is AugType.GEO} + + def _names(self, mode, **kwargs): + return [type(t).__name__ for t in build_gpu_pipeline(self.section, mode=mode, retain_stats=self.spec, **kwargs)] + + def test_the_snapshot_follows_the_geometry(self): + for mode in PipelineMode: + with self.subTest(mode=mode.value): + names = self._names(mode) + cut = names.index("StatsSnapshotGPU") + self.assertTrue( + set(names[:cut]) <= self.geometry, f"something that is not geometry runs before the snapshot: {names[:cut]}" + ) + self.assertTrue( + set(names[cut + 1 :]).isdisjoint(self.geometry - {"RandomCropTransformGPU"}), + "a geometric transform other than crop runs after the snapshot", + ) + + def test_the_restore_precedes_the_crop_and_the_normalization(self): + for mode in PipelineMode: + with self.subTest(mode=mode.value): + names = self._names(mode) + self.assertEqual( + names[-3:], + ["StatsRestoreGPU", "RandomCropTransformGPU", "ZscoreNormalizationGPU"], + f"unexpected tail: {names[-4:]}", + ) + + def test_the_low_resolution_hoist_stays_inside_the_bracket(self): + """`random_order` hoists it next to the geometry, but it is an enhancement.""" + names = self._names(PipelineMode.RANDOM_ORDER) + self.assertLess(names.index("StatsSnapshotGPU"), names.index("RandomLowResTransformGPU")) + self.assertLess(names.index("RandomLowResTransformGPU"), names.index("StatsRestoreGPU")) + + def test_the_bracket_is_absent_without_the_block(self): + for mode in PipelineMode: + with self.subTest(mode=mode.value): + plain = [type(t).__name__ for t in build_gpu_pipeline(self.section, mode=mode)] + self.assertNotIn("StatsSnapshotGPU", plain) + self.assertNotIn("StatsRestoreGPU", plain) + + def test_the_only_transform_the_bracket_moves_is_the_crop(self): + """Everything else keeps the position it had; crop is staged after the restore.""" + for mode in PipelineMode: + with self.subTest(mode=mode.value): + plain = [type(t).__name__ for t in build_gpu_pipeline(self.section, mode=mode)] + bracketed = self._names(mode) + self.assertEqual(sorted(bracketed), sorted([*plain, "StatsSnapshotGPU", "StatsRestoreGPU"])) + without_crop = [n for n in bracketed if not n.startswith("Stats") and n != "RandomCropTransformGPU"] + self.assertEqual([n for n in plain if n != "RandomCropTransformGPU"], without_crop) + self.assertIn("RandomCropTransformGPU", bracketed) + + def test_it_also_brackets_a_config_ordered_pipeline(self): + from smauglab.config import OrderSource + + names = self._names(PipelineMode.SEQUENTIAL, order_source=OrderSource.CONFIG) + self.assertEqual(names[-3:], ["StatsRestoreGPU", "RandomCropTransformGPU", "ZscoreNormalizationGPU"]) + self.assertEqual(names.index("StatsSnapshotGPU"), 2) + + def test_a_lone_normalization_is_staged_after_the_restore(self): + """Without the bracket a single normalization keeps its position; with it, it trails.""" + section = {"RandomGammaGPU": {"p": 1.0}, "ZscoreNormalizationGPU": {"p": 1.0}} + self.assertEqual( + [type(t).__name__ for t in build_gpu_pipeline(section, retain_stats=self.spec)], + ["StatsSnapshotGPU", "RandomGammaGPU", "StatsRestoreGPU", "ZscoreNormalizationGPU"], + ) + + +class TestTheConfigSurface(SmaugLabTestCase): + def test_the_shipped_config_carries_the_block(self): + payload = json.loads((configs_dir() / "transform_params_gpu_stats-restore.json").read_text()) + spec = SmaugConfig(payload, source="transform_params_gpu_stats-restore.json").retain_stats_spec() + self.assertEqual(spec.std_scale, (0.9, 1.1)) + self.assertEqual(spec.percentile_scale, (0.9, 1.1)) + self.assertEqual(spec.stats, STAT_NAMES) + self.assertFalse( + any(block.get("retain_stats") for block in payload["GPU"].values() if isinstance(block, dict)), + "a per-transform retain_stats is still on, and the two do not compose", + ) + + def test_the_paper_configs_do_not_carry_it(self): + for name in ("transform_params_one-sequence-to-segment-them-all.json", "transform_params_paper-synthseg.json"): + with self.subTest(config=name): + payload = json.loads((configs_dir() / name).read_text()) + self.assertIsNone(SmaugConfig(payload, source=name).retain_stats_spec()) + + def test_a_typo_is_reported_with_the_rest(self): + with self.assertRaises(InvalidConfigError) as caught: + SmaugConfig({"GPU": {"RandomGammaGPU": {"p": 1}}, "pipeline": {"retain_stats": {"std_scaled": [1, 2]}}}) + self.assertTrue(any("unknown key 'std_scaled'" in problem for problem in caught.exception.problems)) + + def test_the_values_are_checked(self): + for block, expected in ( + ({"stats": ["mean"]}, "at least two statistics"), + ({"stats": ["median"]}, "unknown name"), + ({"std_scale": [1.5, 0.5]}, "non-decreasing"), + ({"std_scale": -1.0}, "must lie inside"), + ({"percentiles": [99.0, 1.0]}, "0 <= low < high <= 100"), + ({"stats_from": "bone"}, "stats_from must be one of"), + ({"p": 2.0}, "must be a number in [0, 1]"), + ({"apply_to_channel": []}, "non-empty list"), + ({"nonzero_threshold": 0.0}, "must be positive"), + ): + with self.subTest(block=block): + problems = spec_problems(block) + self.assertTrue(any(expected in problem for problem in problems), f"{block} -> {problems}") + + def test_an_absent_block_is_off_and_an_empty_one_is_the_defaults(self): + self.assertIsNone(parse_spec(None), "an absent block must emit no leaves at all") + self.assertEqual(parse_spec({}), StatsSpec(), "a written-but-empty block is the defaults, not off") + self.assertEqual(spec_problems(None), []) + + def test_the_pairs_follow_the_canonical_order(self): + """The pair table is indexed by a draw, so the order of the keys must not decide it.""" + self.assertEqual(parse_spec({"stats": ["p99", "mean"]}).stats, ("mean", "p99")) + self.assertEqual(StatsSpec().pairs, tuple((i, j) for i in range(4) for j in range(i + 1, 4)))