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