diff --git a/.gitignore b/.gitignore index 83fbfeb..3fdaf94 100644 --- a/.gitignore +++ b/.gitignore @@ -46,6 +46,9 @@ MANIFEST # migration table, the test fixtures and the migrator itself. !unit_tests/fixtures/**/*.json !migration/*.json +# Generated by `smauglab template --write` and checked by CI; it is source, +# not per-experiment output. +!smauglab/configs/all_augmentations.json *.yaml *.nii *.nii.gz diff --git a/README.md b/README.md index e2e1853..ffb35e7 100644 --- a/README.md +++ b/README.md @@ -168,3 +168,47 @@ If you use SmaugLab, please make sure to cite the following paper: year={2026} } ``` + +## Available augmentations + +Which augmentations exist, and which backends implement each one. An empty cell +means no implementation on that backend yet. Regenerate with `smauglab matrix --write`. + + +| Augmentation | Group | GPU | CPU | MONAI | +| --- | --- | --- | --- | --- | +| flip | GEO | `RandomFlipTransformGPU` | — | — | +| affine | GEO | `RandomAffineGPU` | — | — | +| crop | GEO | `RandomCropTransformGPU` | — | — | +| spatial | GEO | — | `SpatialTransform` | — | +| gaussian_noise | GE | `RandomGaussianNoiseGPU` | `GaussianNoiseTransform` | — | +| gaussian_blur | GE | `RandomGaussianBlurGPU` | `GaussianBlurTransform` | — | +| brightness | GE | `RandomBrightnessGPU` | `MultiplicativeBrightnessTransform` | — | +| contrast | GE | `RandomContrastGPU` | `ContrastTransform` | — | +| gamma | GE | `RandomGammaGPU` | `GammaTransform` | — | +| inv_gamma | GE | `RandomInvGammaGPU` | `InvertedGammaTransform` | — | +| clamp | GE | `RandomClampGPU` | — | — | +| low_res | GE | `RandomLowResTransformGPU` | `SimulateLowResolutionTransform` | — | +| acq | GE | `RandomAcqTransformGPU` | — | — | +| zscore | GE | `ZscoreNormalizationGPU` | `ZscoreNormalization` | — | +| mirror | GEO | — | `MirrorTransform` | — | +| scharr | TA | `RandomScharrGPU` | `ScharrConvTransform` | — | +| laplace | TA | `RandomLaplaceGPU` | `LaplaceConvTransform` | — | +| unsharp_mask | TA | `RandomUnsharpMaskGPU` | — | — | +| rand_conv | TA | `RandomRandConvGPU` | — | — | +| bias_field | TA | `RandomBiasFieldGPU` | — | — | +| inverse | TA | `RandomInverseGPU` | — | — | +| histogram_equal | TA | `RandomHistogramEqualizationGPU` | `HistogramEqualTransform` | — | +| redistribute_seg | TA | `RandomRedistributeSegGPU` | `RedistributeTransform` | — | +| palette | TA | `RandomPaletteGPU` | — | — | +| domain_transfer | TA | `RandomDomainTransferGPU` | — | — | +| synthseg | TA | `RandomSynthSegGPU` | — | — | +| artifact | TA | — | `ArtifactTransform` | — | +| spatial_custom | GEO | — | `SpatialCustomTransform` | — | +| shape | GE | — | `ShapeTransform` | — | +| func_log1p | TA | `RandomLog1pGPU` | `Log1pTransform` | — | +| func_sqrt | TA | `RandomSqrtGPU` | `SqrtTransform` | — | +| func_sin | TA | `RandomSinGPU` | `SinTransform` | — | +| func_exp | TA | `RandomExpGPU` | `ExpTransform` | — | +| func_sigmoid | TA | `RandomSigmoidGPU` | `SigmoidTransform` | — | + diff --git a/migration/migrate.py b/migration/migrate.py index bda700a..27b99bf 100644 --- a/migration/migrate.py +++ b/migration/migrate.py @@ -389,6 +389,26 @@ def migrate(payload: dict, source: str = "") -> tuple[dict, list[str]]: target["SpatialTransform"] = _migrate_block("SpatialTransform", spatial, Backend.CPU, source, legacy_key="SpatialTransform") notes.append("moved nnUNetSpatialTransform -> CPU.SpatialTransform") + if backend is Backend.CPU and "SpatialTransform" not in migrated and "SpatialTransform" not in out.get("CPU", {}): + # nnU-Net's own SpatialTransform used to be appended by the trainer with + # these values hardcoded, so no config named it. The trainer builds its CPU + # pipeline from the config now, so it has to be in the file or it silently + # stops running. Values are verbatim from get_training_transforms; + # patch_size and rotation stay out because nnU-Net supplies them at runtime + # (they are context_params on the registry entry). + migrated["SpatialTransform"] = { + "patch_center_dist_from_border": 0, + "random_crop": False, + "p_elastic_deform": 0, + "p_rotation": 0, + "p_scaling": 0, + "scaling": [0.7, 1.4], + "p_synchronize_scaling_across_axes": 1, + "bg_style_seg_sampling": False, + "mode_seg": "nearest", + } + notes.append("added CPU.SpatialTransform, which the trainer used to hardcode") + if backend is Backend.CPU: axes = section.get("mirror_axes") if axes: @@ -398,6 +418,11 @@ def migrate(payload: dict, source: str = "") -> tuple[dict, list[str]]: choose = section.get("RandomChooseXTransforms") if isinstance(choose, dict): out.setdefault("pipeline", {})["random_choose"] = choose + # Which pipeline ran used to be the *trainer class*: a config carrying this + # block was only ever used with the list trainer. That choice is + # pipeline.mode now, so it has to be written down. + out["pipeline"].setdefault("mode", "random_order") + notes.append("moved RandomChooseXTransforms -> pipeline.random_choose, mode=random_order") notes.append("moved RandomChooseXTransforms -> pipeline.random_choose") if migrated: diff --git a/pyproject.toml b/pyproject.toml index a00694a..44d122f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,6 +94,7 @@ all = ["monai", "tqdm", "wandb"] dev = ["build", "coverage", "mypy", "pre-commit", "pytest", "pytest-cov", "ruff", "twine"] [tool.poetry.scripts] +smauglab = "smauglab.cli:main" smauglab_add_nnunettrainer = "smauglab.add_trainer:main" [build-system] diff --git a/smauglab/cli.py b/smauglab/cli.py new file mode 100644 index 0000000..c2cc236 --- /dev/null +++ b/smauglab/cli.py @@ -0,0 +1,279 @@ +"""The `smauglab` command line: what augmentations exist, and are my configs valid? + +Before the registry, both questions could only be answered by reading four +hand-written dispatch ladders side by side. These subcommands read the registry, so +they cannot go out of date: + + smauglab list --backend gpu what a GPU config can name, in pipeline order + smauglab matrix which backends implement each augmentation + smauglab show RandomScharrGPU one augmentation's parameters and defaults + smauglab validate config.json strict check, every problem at once + smauglab template --backend gpu a config naming everything, at defaults + smauglab hash config.json content-addressed config identity + +Bringing a pre-registry config forward is a one-time job and is not a subcommand: the +migrator lives in `migration/` in the repository, not in the wheel. + +`hash` is the only one that does not need the registry; the rest import torch and +kornia to populate it, which takes a few seconds on first use. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +REPO = Path(__file__).resolve().parent.parent +README = REPO / "README.md" +TEMPLATE_PATH = Path(__file__).resolve().parent / "configs" / "all_augmentations.json" +BEGIN = "" +END = "" + + +def _backend(name: str): + from smauglab.registry import Backend + + return Backend[name.upper()] + + +# --- generated artefacts ---------------------------------------------------------- + + +def matrix_block() -> str: + from smauglab import registry + + return f"{BEGIN}\n{registry.render_matrix('md')}\n{END}" + + +def readme_with_matrix(text: str) -> str: + """Replace the marked block, or append the section if it is not there yet.""" + if BEGIN in text and END in text: + head, rest = text.split(BEGIN, 1) + _, tail = rest.split(END, 1) + return f"{head}{matrix_block()}{tail}" + section = ( + "\n## Available augmentations\n\n" + "Which augmentations exist, and which backends implement each one. An empty cell\n" + "means no implementation on that backend yet. Regenerate with `smauglab matrix --write`.\n\n" + matrix_block() + "\n" + ) + return text.rstrip("\n") + "\n" + section + + +def template_json() -> str: + from smauglab import registry + from smauglab.registry import Backend + + payload = {b.value: registry.render_template(b) for b in (Backend.GPU, Backend.CPU)} + return json.dumps(payload, indent=4) + "\n" + + +def _sync(targets: dict[Path, str], *, write: bool, check: bool) -> int: + stale = [path for path, content in targets.items() if not path.is_file() or path.read_text() != content] + if check: + for path in stale: + print(f"out of date: {path.relative_to(REPO)}") + if stale: + print("Run `smauglab matrix --write` / `smauglab template --write` and commit the result.") + return 1 if stale else 0 + if write: + for path in stale: + path.write_text(targets[path]) + print(f"wrote {path.relative_to(REPO)}") + if not stale: + print("already up to date") + return 0 + + +# --- subcommands ------------------------------------------------------------------ + + +def cmd_list(args) -> int: + from smauglab import registry + from smauglab.registry import AugType + + backend = _backend(args.backend) if args.backend else None + group = AugType[args.group.upper()] if args.group else None + entries = registry.entries(backend=backend, group=group) + + if args.json: + print( + json.dumps( + [ + { + "name": e.name, + "backend": e.backend.value, + "aug_id": e.aug_id.value, + "group": e.group.value, + "position": registry.pipeline_position(e), + "summary": e.summary, + } + for e in entries + ], + indent=2, + ) + ) + return 0 + + if not entries: + print("no augmentations match") + return 0 + width = max(len(e.name) for e in entries) + current = None + for entry in entries: + if entry.backend is not current: + current = entry.backend + print(f"\n{current.value} ({len([e for e in entries if e.backend is current])}), in pipeline order:") + print(f" {registry.pipeline_position(entry):>4} {entry.name:<{width}} {entry.group.value:<4} {entry.summary}") + return 0 + + +def cmd_matrix(args) -> int: + from smauglab import registry + + if args.write or args.check: + return _sync( + {README: readme_with_matrix(README.read_text()), TEMPLATE_PATH: template_json()}, + write=args.write, + check=args.check, + ) + if args.format == "json": + table = registry.matrix() + print( + json.dumps( + {aug_id.value: {b.value: (e.name if e else None) for b, e in row.items()} for aug_id, row in table.items()}, + indent=2, + ) + ) + return 0 + print(registry.render_matrix(args.format)) + return 0 + + +def cmd_show(args) -> int: + from smauglab import registry + + try: + entry = registry.get(args.name) + except registry.UnknownAugmentationError as exc: + print(str(exc), file=sys.stderr) + return 1 + + print( + f"{entry.name} ({entry.backend.value}, group {entry.group.value}, position {registry.pipeline_position(entry)}, aug_id {entry.aug_id.value})" + ) + if entry.summary: + print(f" {entry.summary}") + if entry.forwards_to is not None: + print(f" forwards extra parameters to {entry.forwards_to.__name__}") + if entry.external_asset: + print(f" needs an external asset; set ${entry.external_asset}") + print(f"\n module: {entry.cls.__module__}") + + params = registry.accepted_params(entry) + required = registry.required_params(entry) + print(f"\n parameters ({len(params)}):") + for name in sorted(params): + default = params[name].default + shown = "REQUIRED" if name in required else repr(default) + print(f" {name:<34} {shown}") + if entry.context_params: + print(f"\n supplied by the trainer, not the config: {', '.join(entry.context_params)}") + return 0 + + +def cmd_validate(args) -> int: + from smauglab.config import validate_file + + failed = 0 + for path in args.configs: + problems = validate_file(path) + if problems: + failed += 1 + print(f"{path}: {len(problems)} problem(s)") + for problem in problems: + print(f" - {problem}") + elif not args.quiet: + print(f"{path}: ok") + return 1 if failed else 0 + + +def cmd_template(args) -> int: + from smauglab import registry + + if args.write or args.check: + return _sync({TEMPLATE_PATH: template_json()}, write=args.write, check=args.check) + section = registry.render_template(_backend(args.backend)) + text = json.dumps({args.backend.upper(): section}, indent=4) + "\n" + if args.output: + Path(args.output).write_text(text) + print(f"wrote {args.output}") + else: + print(text, end="") + return 0 + + +def cmd_hash(args) -> int: + from smauglab.config import config_hash + + for path in args.configs: + payload = json.loads(Path(path).read_text()) + print(f"{config_hash(payload)[:8]} {path}") + return 0 + + +# --- wiring ----------------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="smauglab", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = parser.add_subparsers(dest="command", required=True) + + p = sub.add_parser("list", help="registered augmentations, in pipeline order") + p.add_argument("--backend", choices=["gpu", "cpu", "monai"]) + p.add_argument("--group", choices=["geo", "ge", "ta"]) + p.add_argument("--json", action="store_true") + p.set_defaults(func=cmd_list) + + p = sub.add_parser("matrix", help="which backends implement each augmentation") + p.add_argument("--format", choices=["md", "table", "json"], default="table") + p.add_argument("--write", action="store_true", help="update README.md and the template config") + p.add_argument("--check", action="store_true", help="exit 1 if either is out of date") + p.set_defaults(func=cmd_matrix) + + p = sub.add_parser("show", help="one augmentation's parameters and defaults") + p.add_argument("name") + p.set_defaults(func=cmd_show) + + p = sub.add_parser("validate", help="strict config check, reporting every problem") + p.add_argument("configs", nargs="+") + p.add_argument("-q", "--quiet", action="store_true", help="only report failures") + p.set_defaults(func=cmd_validate) + + p = sub.add_parser("template", help="a config naming every augmentation at its defaults") + p.add_argument("--backend", choices=["gpu", "cpu"], default="gpu") + p.add_argument("-o", "--output") + p.add_argument("--write", action="store_true", help="update the shipped template config") + p.add_argument("--check", action="store_true", help="exit 1 if it is out of date") + p.set_defaults(func=cmd_template) + + p = sub.add_parser("hash", help="content-addressed config identity") + p.add_argument("configs", nargs="+") + p.set_defaults(func=cmd_hash) + + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + if getattr(args, "output", None) and len(getattr(args, "configs", [])) > 1: + print("-o takes a single input config", file=sys.stderr) + return 2 + result: Any = args.func(args) + return int(result or 0) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json new file mode 100644 index 0000000..b25b005 --- /dev/null +++ b/smauglab/configs/all_augmentations.json @@ -0,0 +1,667 @@ +{ + "GPU": { + "RandomFlipTransformGPU": { + "flip_axis": [ + 0 + ], + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false + }, + "RandomAffineGPU": { + "align_corners": true, + "degrees": 10, + "keepdim": true, + "p": 0.5, + "p_batch": 1.0, + "resample": "BILINEAR", + "same_on_batch": false, + "scale": [ + 0.9, + 1.1 + ], + "shears": [ + -10, + 10, + -10, + 10, + -10, + 10 + ], + "translate": [ + 0.1, + 0.1, + 0.1 + ] + }, + "RandomSynthSegGPU": { + "apply_bias_field": true, + "apply_intensity_augmentation": true, + "apply_resolution": true, + "apply_to_channel": null, + "atlas_res": 1.0, + "bias_field_std": 0.7, + "bias_scale": 0.025, + "blur_range": 1.03, + "clip": 300.0, + "data_res": null, + "em_background_clusters_range": [ + 3, + 10 + ], + "em_background_label": 0, + "em_label_completion": false, + "em_max_fit_voxels": 100000, + "em_n_foreground_clusters": 2, + "em_n_iters": 20, + "em_same_on_batch": false, + "flip_axis": 2, + "gamma_std": 0.5, + "generation_classes": null, + "generation_labels": null, + "keepdim": true, + "max_res_aniso": 8.0, + "max_res_iso": 4.0, + "n_channels": 1, + "n_neutral_labels": null, + "nonlin_scale": 0.04, + "nonlin_std": 4.0, + "normalise": true, + "output_labels": null, + "p": 0.5, + "p_batch": 1.0, + "prior_distributions": "uniform", + "prior_means": null, + "prior_stds": null, + "randomise_res": true, + "rotation_bounds": 15.0, + "same_on_batch": false, + "scaling_bounds": 0.2, + "shearing_bounds": 0.012, + "svf_integration_steps": 7, + "thickness": null, + "translation_bounds": false + }, + "RandomPaletteGPU": { + "alpha_magnitude_range": [ + 0.5, + 2.0 + ], + "blur_sigmas": [ + 0.0, + 0.0, + 0.0, + 0.3, + 0.5, + 0.8 + ], + "c_choices": [ + 2, + 3, + 4, + 5, + 6 + ], + "dark_threshold": 0.01, + "keepdim": false, + "label_classes": null, + "label_remap_prob": 0.5, + "min_label_voxels": 4, + "n_kmeans_subsample": 10000, + "p": 1.0, + "p_batch": 1.0, + "s_choices": [ + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10 + ], + "same_on_batch": false, + "skip_parcellation_prob": 0.1, + "skip_sub_parc_prob": 0.4 + }, + "RandomDomainTransferGPU": { + "any_source": false, + "apply_to_channel": null, + "bank_path": null, + "bias_field_std": 0.0, + "bias_scale": 0.03, + "blend_concentration": 1.0, + "blend_targets": 1, + "include_self": true, + "keepdim": true, + "p": 0.2, + "p_batch": 1.0, + "p_class_mix": 0.0, + "p_spatial_mix": 0.0, + "pct": 1.0, + "same_on_batch": false, + "sigma": 2.0, + "source_label": null, + "spatial_mix_gain": 3.0, + "spatial_mix_scale": 0.03, + "targets": null, + "zscore_io": "auto" + }, + "RandomInverseGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomHistogramEqualizationGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomRedistributeSegGPU": { + "apply_to_channel": [ + 0 + ], + "dilation_iterations_range": [ + 1, + 3 + ], + "in_seg": 0.2, + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false, + "std_noise_range": [ + 0.1, + 0.3 + ] + }, + "RandomScharrGPU": { + "absolute": true, + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": true, + "same_on_batch": false + }, + "RandomLaplaceGPU": { + "absolute": false, + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomUnsharpMaskGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false, + "sigma": 1.0, + "unsharp_amount": 1.5 + }, + "RandomRandConvGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "kernel_sizes": [ + 1, + 3, + 5, + 7 + ], + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomClampGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "max_clamp_amount": 0.0, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomGaussianNoiseGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mean": 0.0, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false, + "std": 1.0 + }, + "RandomGaussianBlurGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false, + "sigma": 1.0 + }, + "RandomBrightnessGPU": { + "apply_to_channel": [ + 0 + ], + "brightness_range": [ + 0.5, + 1.5 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false + }, + "RandomGammaGPU": { + "apply_to_channel": [ + 0 + ], + "gamma_range": [ + 0.7, + 1.5 + ], + "in_seg": 0.0, + "keepdim": false, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomInvGammaGPU": { + "apply_to_channel": [ + 0 + ], + "gamma_range": [ + 0.7, + 1.5 + ], + "in_seg": 0.0, + "keepdim": false, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomContrastGPU": { + "apply_to_channel": [ + 0 + ], + "contrast_range": [ + 0.75, + 1.25 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomLog1pGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomSqrtGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomSinGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomExpGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomSigmoidGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomLowResTransformGPU": { + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false, + "scale": [ + 0.3, + 1.0 + ] + }, + "RandomAcqTransformGPU": { + "apply_to_channel": [ + 0 + ], + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false, + "scale": [ + 0.3, + 1.0 + ] + }, + "RandomCropTransformGPU": { + "crop": [ + 1.0, + 1.0 + ], + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "pos": [ + 0.0, + 1.0 + ], + "same_on_batch": false + }, + "RandomBiasFieldGPU": { + "apply_to_channel": [ + 0 + ], + "coefficients": 0.5, + "in_seg": 0.0, + "invert": false, + "keepdim": true, + "mix_in_out": false, + "order": 3, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "ZscoreNormalizationGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0 + } + }, + "CPU": { + "LaplaceConvTransform": { + "absolute": false, + "p": 1.0, + "retain_stats": false + }, + "ScharrConvTransform": { + "absolute": true, + "p": 1.0, + "retain_stats": false + }, + "Log1pTransform": { + "p": 1.0, + "retain_stats": false + }, + "SqrtTransform": { + "p": 1.0, + "retain_stats": false + }, + "SinTransform": { + "p": 1.0, + "retain_stats": false + }, + "ExpTransform": { + "p": 1.0, + "retain_stats": false + }, + "SigmoidTransform": { + "p": 1.0, + "retain_stats": false + }, + "HistogramEqualTransform": { + "p": 1.0, + "retain_stats": false + }, + "RedistributeTransform": { + "classes": null, + "in_seg": 0.2, + "p": 1.0, + "retain_stats": false + }, + "ShapeTransform": { + "ignore_axes": [], + "p": 1.0, + "shape_min": 1 + }, + "ArtifactTransform": { + "bias_field": false, + "blur": false, + "ghosting": false, + "motion": false, + "noise": false, + "p": 1.0, + "random_pick": false, + "spike": false, + "swap": false + }, + "SpatialCustomTransform": { + "affine": false, + "anisotropy": false, + "elastic": false, + "flip": false, + "p": 1.0, + "random_pick": false + }, + "SpatialTransform": { + "align_corners": false, + "bg_style_seg_sampling": true, + "border_mode_seg": "zeros", + "center_deformation": true, + "elastic_deform_magnitude": [ + 0, + 0.2 + ], + "elastic_deform_scale": [ + 0, + 0.2 + ], + "mode_image": "bilinear", + "mode_seg": "bilinear", + "p_elastic_deform": 0, + "p_rot_per_axis": 1, + "p_rotation": 0, + "p_scaling": 0, + "p_synchronize_def_scale_across_axes": 0, + "p_synchronize_scaling_across_axes": 0, + "padding_mode_image": "zeros", + "padding_value_image": 0, + "padding_value_seg": 0, + "patch_center_dist_from_border": null, + "random_crop": null, + "scaling": [ + 0.7, + 1.3 + ] + }, + "GaussianNoiseTransform": { + "noise_variance": [ + 0, + 0.1 + ], + "p": 1.0, + "p_per_channel": 1.0, + "synchronize_channels": false + }, + "GaussianBlurTransform": { + "benchmark": false, + "blur_sigma": [ + 1, + 5 + ], + "p": 1.0, + "p_per_channel": 1, + "synchronize_axes": false, + "synchronize_channels": false + }, + "MultiplicativeBrightnessTransform": { + "multiplier_range": null, + "p": 1.0, + "p_per_channel": 1, + "synchronize_channels": null + }, + "ContrastTransform": { + "contrast_range": null, + "p": 1.0, + "p_per_channel": 1.0, + "preserve_range": null, + "synchronize_channels": null + }, + "SimulateLowResolutionTransform": { + "allowed_channels": null, + "ignore_axes": null, + "p": 1.0, + "p_per_channel": 1, + "scale": null, + "synchronize_axes": null, + "synchronize_channels": null + }, + "InvertedGammaTransform": { + "gamma": [ + 0.7, + 1.5 + ], + "p": 1.0, + "p_per_channel": 1, + "p_retain_stats": 1, + "synchronize_channels": false + }, + "GammaTransform": { + "gamma": null, + "p": 1.0, + "p_invert_image": null, + "p_per_channel": null, + "p_retain_stats": null, + "synchronize_channels": null + }, + "MirrorTransform": { + "allowed_axes": null + }, + "ZscoreNormalization": { + "p": 1.0 + } + } +} diff --git a/smauglab/configs/transform_params.json b/smauglab/configs/transform_params.json index 621c5ff..bbaaeaf 100644 --- a/smauglab/configs/transform_params.json +++ b/smauglab/configs/transform_params.json @@ -133,6 +133,20 @@ "p_retain_stats": 1, "p": 0.3 }, + "SpatialTransform": { + "patch_center_dist_from_border": 0, + "random_crop": false, + "p_elastic_deform": 0, + "p_rotation": 0, + "p_scaling": 0, + "scaling": [ + 0.7, + 1.4 + ], + "p_synchronize_scaling_across_axes": 1, + "bg_style_seg_sampling": false, + "mode_seg": "nearest" + }, "MirrorTransform": { "allowed_axes": [ 0, diff --git a/smauglab/trainers/nnUNetTrainerDAExt.py b/smauglab/trainers/nnUNetTrainerDAExt.py index 47098ab..f1ae88e 100644 --- a/smauglab/trainers/nnUNetTrainerDAExt.py +++ b/smauglab/trainers/nnUNetTrainerDAExt.py @@ -1,7 +1,25 @@ +"""The SmaugLab nnU-Net trainer. + +One class, because the config already says everything the three previous trainers +encoded between them. `nnUNetTrainerDAExt` built the CPU pipeline, `...GPU` built the +GPU one plus nnU-Net's SpatialTransform, and `...Hybrid` built both -- but a config is +sectioned into "CPU" and "GPU", so which sections are populated decides that on its +own: + + transform_params.json CPU: 19 GPU: 0 -> CPU-only, as ...DAExt did + transform_params_gpu.json CPU: 1 GPU: 26 -> GPU-only, as ...DAExtGPU did + transform_params_hybrid.json CPU: 19 GPU: 24 -> both, as ...DAExtHybrid did + +The class keeps the name `nnUNetTrainerDAExtGPU` whatever the config contains. That is +not cosmetic: nnU-Net writes the trainer class name into every checkpoint +(`checkpoint['trainer_name']`) and resolves the class from it at inference, so +renaming it would make several hundred trained models unloadable. +""" + import importlib -import json import os import shutil +import warnings from typing import Union import numpy as np @@ -9,10 +27,8 @@ from batchgeneratorsv2.helpers.scalar_type import RandomScalar from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform from batchgeneratorsv2.transforms.nnunet.seg_to_onehot import MoveSegAsOneHotToDataTransform -from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform from batchgeneratorsv2.transforms.utils.compose import ComposeTransforms from batchgeneratorsv2.transforms.utils.deep_supervision_downsampling import DownsampleSegForDSTransform -from batchgeneratorsv2.transforms.utils.pseudo2d import Convert2DTo3DTransform, Convert3DTo2DTransform from batchgeneratorsv2.transforms.utils.remove_label import RemoveLabelTansform from batchgeneratorsv2.transforms.utils.seg_to_regions import ConvertSegmentationToRegionsTransform from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer @@ -20,114 +36,79 @@ from torch import autocast from smauglab import configs +from smauglab.config import load_config +from smauglab.registry import Backend from smauglab.trainers.utils import DownsampleSegForDSTransformCustom, nnunet_tail_transforms -from smauglab.transforms.cpu.transforms import AugTransforms -from smauglab.transforms.gpu.transforms import AugTransformsGPU - - -class nnUNetTrainerDAExt(nnUNetTrainer): - def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): - super().__init__(plans, configuration, fold, dataset_json, device) - - def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): - rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = ( - super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() +from smauglab.transforms.build import build_cpu_pipeline, build_gpu_pipeline +from smauglab.transforms.gpu.base import AugmentationSequentialCustom + +#: Env var naming the config to train with. +CONFIG_ENV = "SMAUGLAB_PARAMS_JSON" + +#: Previous name, still honoured so existing sweep scripts keep working. The three +#: old variables (_CPU_JSON, _GPU_JSON, _HYBRID_JSON) picked a trainer as much as a +#: file; only this one ever had an external caller. +LEGACY_CONFIG_ENV = "SMAUGLAB_PARAMS_GPU_JSON" + +DEFAULT_CONFIG = "transform_params_gpu.json" + + +def resolve_config_path() -> str: + """Locate the config: new env var, then the deprecated one, then the default.""" + path = os.environ.get(CONFIG_ENV) + if path: + return path + legacy = os.environ.get(LEGACY_CONFIG_ENV) + if legacy: + warnings.warn( + f"{LEGACY_CONFIG_ENV} is deprecated; use {CONFIG_ENV}. The config's CPU and GPU " + "sections now decide which augmentations run, so the name no longer implies a backend.", + DeprecationWarning, + stacklevel=2, ) - # Remove mirroring - mirror_axes = None - self.inference_allowed_mirroring_axes = None - return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes + return legacy + return str(importlib.resources.files(configs) / DEFAULT_CONFIG) - @staticmethod - def get_training_transforms( - patch_size: Union[np.ndarray, tuple[int, ...]], - rotation_for_DA: RandomScalar, - deep_supervision_scales: Union[list, tuple, None], - mirror_axes: tuple[int, ...], - do_dummy_2d_data_aug: bool, - use_mask_for_norm: list[bool] | None = None, - is_cascaded: bool = False, - foreground_labels: Union[tuple[int, ...], list[int]] | None = None, - regions: list[Union[list[int], tuple[int, ...], int]] | None = None, - ignore_label: int | None = None, - retain_stats: bool = False, - ) -> BasicTransform: - transforms = [] - ### Adds transforms - # Load transform parameters from json file - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_CPU_JSON", str(configs_path / "transform_params.json")) - transforms.append( - AugTransforms( - json_path=json_path, - do_dummy_2d_data_aug=do_dummy_2d_data_aug, - patch_size=patch_size, - rotation_for_DA=rotation_for_DA, - mirror_axes=mirror_axes, - ) - ) +def _has_gpu_augmentations(config) -> bool: + """Whether this config asks for anything on the GPU side.""" + return bool(config.names(Backend.GPU)) - if do_dummy_2d_data_aug: - transforms.append(Convert3DTo2DTransform()) - patch_size_spatial = patch_size[1:] - else: - patch_size_spatial = patch_size - transforms.append( - SpatialTransform( - patch_size_spatial, - patch_center_dist_from_border=0, - random_crop=False, - p_elastic_deform=0, - p_rotation=0, - rotation=rotation_for_DA, - p_scaling=0, - scaling=(0.7, 1.4), - p_synchronize_scaling_across_axes=1, - bg_style_seg_sampling=False, - mode_seg="nearest", - ) - ) - - if do_dummy_2d_data_aug: - transforms.append(Convert2DTo3DTransform()) - transforms.extend( - nnunet_tail_transforms( - use_mask_for_norm=use_mask_for_norm, - deep_supervision_scales=deep_supervision_scales, - is_cascaded=is_cascaded, - foreground_labels=foreground_labels, - regions=regions, - ignore_label=ignore_label, - ) - ) - return ComposeTransforms(transforms) +class nnUNetTrainerDAExtGPU(nnUNetTrainer): + """nnU-Net trainer driven entirely by a SmaugLab config. + CPU-section augmentations run in the dataloader worker; GPU-section ones run on + the batch in `train_step`. + """ -class nnUNetTrainerDAExtGPU(nnUNetTrainer): def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): super().__init__(plans, configuration, fold, dataset_json, device) - self.num_epochs = 1000 - - # Load transform parameters from json file - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_GPU_JSON", str(configs_path / "transform_params_gpu.json")) - self.transforms = AugTransformsGPU(json_path=json_path).to(self.device) - print(f"Using SmaugLab GPU transforms with parameters from: {json_path}") - - # Copy json transfrom parameters to output folder - shutil.copy(json_path, os.path.join(self.output_folder, "transform_params_gpu_used_for_training.json")) - - def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): - rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = ( - super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() + json_path = resolve_config_path() + config = load_config(json_path) + + # Built only when the config actually asks for GPU augmentations, so a + # CPU-only config costs nothing per step and train_step stays a no-op. + self.transforms: AugmentationSequentialCustom | None = None + if _has_gpu_augmentations(config): + self.transforms = AugmentationSequentialCustom( + *build_gpu_pipeline( + config.section(Backend.GPU), + mode=config.pipeline_mode(), + options=config.pipeline_options("random_choose"), + source=config.source, + ), + data_keys=["input", "mask"], + same_on_batch=True, + ).to(self.device) + + print(f"Using SmaugLab transforms from: {json_path}") + print( + f" CPU: {len(config.names(Backend.CPU))} augmentations, GPU: {len(config.names(Backend.GPU))}, mode: {config.pipeline_mode().value}" ) - # Remove mirroring - mirror_axes = None - self.inference_allowed_mirroring_axes = None - return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes + + shutil.copy(json_path, os.path.join(self.output_folder, "transform_params_used_for_training.json")) @staticmethod def get_training_transforms( @@ -141,53 +122,41 @@ def get_training_transforms( foreground_labels: Union[tuple[int, ...], list[int]] | None = None, regions: list[Union[list[int], tuple[int, ...], int]] | None = None, ignore_label: int | None = None, - retain_stats: bool = False, ) -> BasicTransform: - transforms = [] - - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_GPU_JSON", str(configs_path / "transform_params_gpu.json")) - with open(json_path) as f: - config = json.load(f) + """Dataloader-side augmentations: whatever the config's CPU section names. - ### Keep some nnunet transforms - if do_dummy_2d_data_aug: - transforms.append(Convert3DTo2DTransform()) - patch_size_spatial = patch_size[1:] - else: - patch_size_spatial = patch_size - - spatial_params = config.get("nnUNetSpatialTransform", {}) + A staticmethod because that is nnU-Net's contract, so it cannot reach the + instance's parsed config and resolves the path itself. `load_config` is + cached, so the file is still read and validated once. + """ + transforms = [] - transforms.append( - SpatialTransform( - patch_size_spatial, - patch_center_dist_from_border=spatial_params.get("patch_center_dist_from_border", 0), - random_crop=spatial_params.get("random_crop", False), - p_elastic_deform=spatial_params.get("p_elastic_deform", 0), - p_rotation=spatial_params.get("p_rotation", 0), + config = load_config(resolve_config_path()) + transforms.extend( + build_cpu_pipeline( + config.section(Backend.CPU), + do_dummy_2d_data_aug=do_dummy_2d_data_aug, + patch_size=patch_size, rotation=rotation_for_DA, - p_scaling=spatial_params.get("p_scaling", 0), - scaling=spatial_params.get("scaling", (0.7, 1.4)), - p_synchronize_scaling_across_axes=spatial_params.get("p_synchronize_scaling_across_axes", 1), - bg_style_seg_sampling=False, - mode_seg="nearest", + source=config.source, ) ) - if do_dummy_2d_data_aug: - transforms.append(Convert2DTo3DTransform()) - + # Deep supervision has to come after whatever last deformed the mask. With GPU + # augmentations that is train_step, so the downsampling happens there; without + # them nothing touches the mask after this point and it belongs here, which is + # where nnU-Net puts it. Passing None is how the tail is told to skip it. transforms.extend( nnunet_tail_transforms( use_mask_for_norm=use_mask_for_norm, - deep_supervision_scales=None, + deep_supervision_scales=None if _has_gpu_augmentations(config) else deep_supervision_scales, is_cascaded=is_cascaded, foreground_labels=foreground_labels, regions=regions, ignore_label=ignore_label, ) ) + return ComposeTransforms(transforms) @staticmethod @@ -220,63 +189,6 @@ def get_validation_transforms( transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales)) return ComposeTransforms(transforms) - def train_step(self, batch: dict) -> dict: - data = batch["data"] - target = batch["target"] - - data = data.to(self.device, non_blocking=True) - # Now target should be a single tensor, not a list - target = target.to(self.device, non_blocking=True) - # if isinstance(target, list): - # target = [i.to(self.device, non_blocking=True) for i in target] - # else: - # target = target.to(self.device, non_blocking=True) - - self.optimizer.zero_grad(set_to_none=True) - # Autocast can be annoying - # If the device_type is 'cpu' then it's slow as heck and needs to be disabled. - # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False) - # So autocast will only be active if we have a cuda device. - with autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): - # Apply GPU augmentations to full-resolution data/target - data, target = self.transforms(data, target) - - # Create multi-scale targets for deep supervision after augmentation - deep_supervision_scales = self._get_deep_supervision_scales() - if deep_supervision_scales is not None: - ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) - target = ds_transform(target) - - output = self.network(data) - # del data - l = self.loss(output, target) - - if self.grad_scaler is not None: - self.grad_scaler.scale(l).backward() - self.grad_scaler.unscale_(self.optimizer) - torch.nn.utils.clip_grad_norm_(self.network.parameters(), 12) - self.grad_scaler.step(self.optimizer) - self.grad_scaler.update() - else: - l.backward() - torch.nn.utils.clip_grad_norm_(self.network.parameters(), 12) - self.optimizer.step() - return {"loss": l.detach().cpu().numpy()} - - -class nnUNetTrainerDAExtHybrid(nnUNetTrainer): - def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): - super().__init__(plans, configuration, fold, dataset_json, device) - - # Load transform parameters from json file - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_HYBRID_JSON", str(configs_path / "transform_params_hybrid.json")) - self.transforms = AugTransformsGPU(json_path=json_path).to(self.device) - print(f"Using SmaugLab hybrid transforms with parameters from: {json_path}") - - # Copy json transfrom parameters to output folder - shutil.copy(json_path, os.path.join(self.output_folder, "transform_params_hybrid_used_for_training.json")) - def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = ( super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() @@ -286,48 +198,6 @@ def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): self.inference_allowed_mirroring_axes = None return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes - @staticmethod - def get_training_transforms( - patch_size: Union[np.ndarray, tuple[int, ...]], - rotation_for_DA: RandomScalar, - deep_supervision_scales: Union[list, tuple, None], - mirror_axes: tuple[int, ...], - do_dummy_2d_data_aug: bool, - use_mask_for_norm: list[bool] | None = None, - is_cascaded: bool = False, - foreground_labels: Union[tuple[int, ...], list[int]] | None = None, - regions: list[Union[list[int], tuple[int, ...], int]] | None = None, - ignore_label: int | None = None, - retain_stats: bool = False, - ) -> BasicTransform: - transforms = [] - - ### Adds transforms - # Load transform parameters from json file - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_HYBRID_JSON", str(configs_path / "transform_params_hybrid.json")) - transforms.append( - AugTransforms( - json_path=json_path, - do_dummy_2d_data_aug=do_dummy_2d_data_aug, - patch_size=patch_size, - rotation_for_DA=rotation_for_DA, - mirror_axes=mirror_axes, - ) - ) - - transforms.extend( - nnunet_tail_transforms( - use_mask_for_norm=use_mask_for_norm, - deep_supervision_scales=None, - is_cascaded=is_cascaded, - foreground_labels=foreground_labels, - regions=regions, - ignore_label=ignore_label, - ) - ) - return ComposeTransforms(transforms) - def train_step(self, batch: dict) -> dict: data = batch["data"] target = batch["target"] @@ -346,14 +216,17 @@ def train_step(self, batch: dict) -> dict: # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False) # So autocast will only be active if we have a cuda device. with autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): - # Apply GPU augmentations to full-resolution data/target - data, target = self.transforms(data, target) - - # Create multi-scale targets for deep supervision after augmentation - deep_supervision_scales = self._get_deep_supervision_scales() - if deep_supervision_scales is not None: - ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) - target = ds_transform(target) + # Apply GPU augmentations to full-resolution data/target, then build the + # deep-supervision targets from the *augmented* mask. A CPU-only config + # builds no GPU pipeline; nothing has touched the mask since the + # dataloader, which already produced those targets. + if self.transforms is not None: + data, target = self.transforms(data, target) + + deep_supervision_scales = self._get_deep_supervision_scales() + if deep_supervision_scales is not None: + ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) + target = ds_transform(target) output = self.network(data) # del data diff --git a/smauglab/trainers/nnUNetTrainerTest.py b/smauglab/trainers/nnUNetTrainerTest.py index a51f32f..f1fbebd 100644 --- a/smauglab/trainers/nnUNetTrainerTest.py +++ b/smauglab/trainers/nnUNetTrainerTest.py @@ -76,6 +76,7 @@ def get_training_transforms( ignore_label=ignore_label, ) ) + return ComposeTransforms(transforms) @@ -129,6 +130,9 @@ def get_training_transforms( if do_dummy_2d_data_aug: transforms.append(Convert2DTo3DTransform()) + # deep_supervision_scales=None: this trainer always runs GPU augmentations, so + # the mask is still being deformed after this point and train_step builds the + # multi-scale targets from the augmented mask instead. transforms.extend( nnunet_tail_transforms( use_mask_for_norm=use_mask_for_norm, @@ -139,6 +143,7 @@ def get_training_transforms( ignore_label=ignore_label, ) ) + return ComposeTransforms(transforms) def train_step(self, batch: dict) -> dict: diff --git a/smauglab/trainers/utils.py b/smauglab/trainers/utils.py index e7a511e..96d952f 100644 --- a/smauglab/trainers/utils.py +++ b/smauglab/trainers/utils.py @@ -16,18 +16,14 @@ def nnunet_tail_transforms( ) -> list[Any]: """The nnU-Net transforms that follow SmaugLab's augmentations, in order. - Five `get_training_transforms` methods across two modules ended with a - character-identical copy of this -- intensity masking, the -1 label removal, the - two cascade transforms, region conversion and deep-supervision downsampling. - - `deep_supervision_scales=None` skips the downsampling, which is how the GPU - trainers had it: they carry the block commented out, because with GPU - augmentations the mask is still being deformed after this point and the - multi-scale targets have to be built from the augmented mask in `train_step`. - - `get_validation_transforms` deliberately does not use this. Its cascade branch - adds only MoveSegAsOneHotToDataTransform, without the two RandomTransform - wrappers, so it is a different sequence rather than another copy of this one. + Every `get_training_transforms` in this package ended with a character-identical + copy of this -- three of them -- covering intensity masking, the -1 label removal, + the two cascade transforms, region conversion and deep-supervision downsampling. + + `deep_supervision_scales=None` skips the downsampling, which is how + `nnUNetTrainerDAExtGPU` defers it to `train_step`: with GPU augmentations the mask + is still being deformed after this point, so the multi-scale targets have to be + built from the augmented mask instead. """ from batchgeneratorsv2.transforms.nnunet.random_binary_operator import ApplyRandomBinaryOperatorTransform from batchgeneratorsv2.transforms.nnunet.remove_connected_components import ( diff --git a/unit_tests/test_cli.py b/unit_tests/test_cli.py new file mode 100644 index 0000000..320977b --- /dev/null +++ b/unit_tests/test_cli.py @@ -0,0 +1,141 @@ +"""The `smauglab` command line. + +Driven through `cli.main` with captured stdout rather than as a subprocess: the +registry import costs a few seconds, and paying it once per process keeps the suite +fast enough to gate every pull request. +""" + +from __future__ import annotations + +import contextlib +import io +import json +import tempfile +import unittest +from pathlib import Path + +from smauglab import cli + +REPO = Path(__file__).resolve().parent.parent +# A *tracked* config. Most of smauglab/configs is gitignored (per-experiment +# sweeps), so naming one of those makes the test pass locally and fail in CI. +DEFAULT_GPU = REPO / "smauglab" / "configs" / "transform_params_gpu.json" +LEGACY = REPO / "unit_tests" / "fixtures" / "legacy_configs" / "configs" / "transform_params_gpu.json" + + +def run(*argv: str) -> tuple[int, str]: + out = io.StringIO() + with contextlib.redirect_stdout(out), contextlib.redirect_stderr(out): + code = cli.main(list(argv)) + return code, out.getvalue() + + +class TestList(unittest.TestCase): + def test_lists_gpu_augmentations_in_pipeline_order(self): + code, out = run("list", "--backend", "gpu") + self.assertEqual(code, 0) + self.assertIn("RandomFlipTransformGPU", out) + # order column ascends + orders = [int(line.split()[0]) for line in out.splitlines() if line.startswith(" ") and line.split()[0].isdigit()] + self.assertEqual(orders, sorted(orders)) + + def test_group_filter(self): + _, out = run("list", "--backend", "gpu", "--group", "geo") + self.assertIn("RandomFlipTransformGPU", out) + self.assertNotIn("RandomScharrGPU", out) + + def test_json_output_is_machine_readable(self): + _, out = run("list", "--backend", "cpu", "--json") + payload = json.loads(out) + self.assertTrue(all({"name", "backend", "aug_id", "group", "position"} <= set(e) for e in payload)) + + +class TestMatrix(unittest.TestCase): + def test_markdown_shows_the_monai_column_as_empty(self): + code, out = run("matrix", "--format", "md") + self.assertEqual(code, 0) + self.assertIn("| Augmentation | Group | GPU | CPU | MONAI |", out) + + def test_json_form_reports_missing_backends_as_null(self): + _, out = run("matrix", "--format", "json") + payload = json.loads(out) + self.assertIsNone(payload["palette"]["CPU"]) + self.assertEqual(payload["palette"]["GPU"], "RandomPaletteGPU") + self.assertTrue(all(row["MONAI"] is None for row in payload.values())) + + def test_check_passes_on_a_clean_tree(self): + code, _ = run("matrix", "--check") + self.assertEqual(code, 0) + + +class TestShow(unittest.TestCase): + def test_reports_parameters_and_defaults(self): + code, out = run("show", "RandomScharrGPU") + self.assertEqual(code, 0) + self.assertIn("group TA", out) + self.assertIn("absolute", out) + + def test_marks_required_parameters(self): + _, out = run("show", "MirrorTransform") + self.assertIn("REQUIRED", out) + + def test_flags_trainer_supplied_parameters(self): + _, out = run("show", "SpatialTransform") + self.assertIn("supplied by the trainer", out) + + def test_unknown_name_exits_nonzero_with_a_suggestion(self): + code, out = run("show", "ScharrTransform") + self.assertEqual(code, 1) + self.assertIn("RandomScharrGPU", out) + + +class TestValidate(unittest.TestCase): + def test_shipped_config_is_valid(self): + code, out = run("validate", str(DEFAULT_GPU)) + self.assertEqual(code, 0) + self.assertIn("ok", out) + + def test_a_broken_config_exits_nonzero_and_reports_every_problem(self): + payload = {"GPU": {"ScharrTransform": {"p": 0.1}, "RandomGaussianNoiseGPU": {"probability": 0.2}}} + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "broken.json" + path.write_text(json.dumps(payload)) + code, out = run("validate", str(path)) + self.assertEqual(code, 1) + self.assertIn("2 problem(s)", out) + self.assertIn("RandomScharrGPU", out) + self.assertIn("'probability' -> p", out) + + def test_a_legacy_config_is_rejected_and_points_at_migrate(self): + code, out = run("validate", str(LEGACY)) + self.assertEqual(code, 1) + self.assertIn("migration/", out) + + +class TestTemplateAndHash(unittest.TestCase): + def test_template_names_every_gpu_augmentation(self): + from smauglab import registry + from smauglab.registry import Backend + + _, out = run("template", "--backend", "gpu") + self.assertEqual(set(json.loads(out)["GPU"]), set(registry.names(Backend.GPU))) + + def test_template_round_trips_through_validation(self): + """A template that cannot be loaded would be worse than no template.""" + from smauglab.config import validate_file + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "template.json" + run("template", "--backend", "gpu", "-o", str(path)) + self.assertEqual(validate_file(path), []) + + def test_hash_is_stable_and_content_addressed(self): + code, first = run("hash", str(DEFAULT_GPU)) + self.assertEqual(code, 0) + _, second = run("hash", str(DEFAULT_GPU)) + self.assertEqual(first, second) + self.assertEqual(len(first.split()[0]), 8) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_registered_augmentations.py b/unit_tests/test_registered_augmentations.py index 0a4bf9f..f2774cb 100644 --- a/unit_tests/test_registered_augmentations.py +++ b/unit_tests/test_registered_augmentations.py @@ -11,6 +11,10 @@ from __future__ import annotations +import inspect +import json +import subprocess +import sys import unittest from pathlib import Path @@ -113,3 +117,64 @@ def test_constructible_with_declared_defaults(self): ) continue entry.cls(**dict(entry.smoke_kwargs)) + + +class TestGeneratedArtifactsAreCurrent(unittest.TestCase): + """`smauglab matrix --check` is the anti-staleness guarantee. + + Run as a subprocess rather than re-deriving the comparison here, so the test + exercises the same code path a developer and CI run. + """ + + def test_readme_matrix_and_template_are_up_to_date(self): + result = subprocess.run( + [sys.executable, "-m", "smauglab.cli", "matrix", "--check"], + capture_output=True, + text=True, + cwd=REPO, + check=False, # a non-zero exit is the assertion below, not an error here + ) + self.assertEqual( + result.returncode, + 0, + f"generated artifacts are stale:\n{result.stdout}{result.stderr}", + ) + + def test_template_covers_exactly_the_registered_augmentations(self): + payload = json.loads(TEMPLATE.read_text()) + for backend in (Backend.GPU, Backend.CPU): + with self.subTest(backend=backend.value): + self.assertEqual(set(payload[backend.value]), set(registry.names(backend))) + + def test_template_parameters_are_all_accepted(self): + """Every key in the template must survive the validation stage 6 will apply.""" + payload = json.loads(TEMPLATE.read_text()) + for backend in (Backend.GPU, Backend.CPU): + for name, params in payload[backend.value].items(): + entry = registry.get(name, backend) + accepted = set(registry.accepted_params(entry)) + with self.subTest(entry=f"{backend.value}.{name}"): + self.assertEqual(set(params) - accepted, set()) + + def test_forced_parameters_are_not_offered(self): + """RandomSynthSegGPU overrides these internally; offering them would lie.""" + payload = json.loads(TEMPLATE.read_text()) + synthseg = payload["GPU"]["RandomSynthSegGPU"] + for forced in ("apply_affine", "apply_nonlinear", "flipping", "output_shape"): + with self.subTest(param=forced): + self.assertNotIn(forced, synthseg) + + +class TestForwardedSignatures(unittest.TestCase): + def test_synthseg_accepts_its_generator_parameters(self): + """forwards_to unions the two signatures, replacing a hand-listed allowlist.""" + from smauglab.transforms.synthseg.generator import SynthSegGenerator + + entry = registry.get("RandomSynthSegGPU", Backend.GPU) + accepted = set(registry.accepted_params(entry)) + generator = set(inspect.signature(SynthSegGenerator).parameters) - set(entry.context_params) + self.assertEqual(generator - accepted, set(), "generator parameters missing from the accepted set") + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_trainers.py b/unit_tests/test_trainers.py new file mode 100644 index 0000000..7f67e5d --- /dev/null +++ b/unit_tests/test_trainers.py @@ -0,0 +1,198 @@ +"""The nnU-Net trainer, which is now one class driven entirely by the config. + +Three trainers used to encode the CPU/GPU split in their class names. The config +already carries it -- which sections are populated says what runs -- so there is one +class, and these tests pin that it reproduces what each of the three used to build. + +`get_training_transforms` is a staticmethod (nnU-Net's contract), so it can be driven +directly: no plans.json, no dataset, no GPU. +""" + +from __future__ import annotations + +import importlib.util +import os +import unittest +import warnings +from pathlib import Path + +from batchgeneratorsv2.transforms.utils.compose import ComposeTransforms + +from smauglab.config import load_config +from smauglab.registry import Backend +from smauglab.transforms.build import PipelineMode, build_cpu_pipeline + +REPO = Path(__file__).resolve().parent.parent +CONFIGS = REPO / "smauglab" / "configs" +PATCH = (24, 24, 24) +ROTATION = (-10, 10) +DS_SCALES = [[1, 1, 1], [0.5, 0.5, 0.5]] + +# Raised at import so it works under pytest and `python -m unittest` alike; a +# `pytestmark` would only be understood by one of them. +if importlib.util.find_spec("nnunetv2") is None: + raise unittest.SkipTest("the trainer needs the nnunetv2 extra") + + +def flatten(transform) -> list[str]: + """Class names, descending into Compose and unwrapping RandomTransform.""" + if isinstance(transform, ComposeTransforms): + return [name for child in transform.transforms for name in flatten(child)] + wrapped = getattr(transform, "transform", None) + if type(transform).__name__ == "RandomTransform" and wrapped is not None: + return [type(wrapped).__name__] + return [type(transform).__name__] + + +def training_transforms(config_name: str, **overrides) -> list[str]: + from smauglab.trainers.nnUNetTrainerDAExt import nnUNetTrainerDAExtGPU + + os.environ["SMAUGLAB_PARAMS_JSON"] = str(CONFIGS / config_name) + kwargs = { + "patch_size": PATCH, + "rotation_for_DA": ROTATION, + "deep_supervision_scales": DS_SCALES, + "mirror_axes": (0, 1, 2), + "do_dummy_2d_data_aug": False, + "use_mask_for_norm": None, + "is_cascaded": False, + "foreground_labels": None, + "regions": None, + "ignore_label": None, + } + kwargs.update(overrides) + return flatten(nnUNetTrainerDAExtGPU.get_training_transforms(**kwargs)) + + +def cpu_block(config_name: str) -> list[str]: + """The CPU pipeline a config asks for, independent of the trainer.""" + config = load_config(str(CONFIGS / config_name)) + built = build_cpu_pipeline(config.section(Backend.CPU), do_dummy_2d_data_aug=False, patch_size=PATCH, rotation=ROTATION) + return [name for transform in built for name in flatten(transform)] + + +class TestOnlyOneTrainerRemains(unittest.TestCase): + def test_the_load_bearing_name_survives(self): + """nnU-Net writes the class name into every checkpoint and resolves the class + from it at inference, and several hundred trained runs record this one.""" + from smauglab.trainers import nnUNetTrainerDAExt + + self.assertTrue(hasattr(nnUNetTrainerDAExt, "nnUNetTrainerDAExtGPU")) + + def test_the_backend_specific_trainers_are_gone(self): + """They differed only in which config they defaulted to, which the config + itself now says. Neither had a single run on disk.""" + from smauglab.trainers import nnUNetTrainerDAExt + + for gone in ("nnUNetTrainerDAExtHybrid", "nnUNetTrainerDAExt"): + with self.subTest(trainer=gone): + self.assertFalse(hasattr(nnUNetTrainerDAExt, gone)) + + +class TestCompositionMatchesTheOldTrainers(unittest.TestCase): + """Each config must build what its dedicated trainer used to build.""" + + def test_gpu_config_matches_the_old_gpu_trainer(self): + self.assertEqual(training_transforms("transform_params_gpu.json"), ["SpatialTransform", "RemoveLabelTransform"]) + + def test_hybrid_config_matches_the_old_hybrid_trainer(self): + expected = [*cpu_block("transform_params_hybrid.json"), "RemoveLabelTransform"] + self.assertEqual(training_transforms("transform_params_hybrid.json"), expected) + + def test_cpu_config_builds_the_whole_cpu_pipeline(self): + got = training_transforms("transform_params.json") + self.assertEqual(got[: len(cpu_block("transform_params.json"))], cpu_block("transform_params.json")) + self.assertIn("SpatialTransform", got) + + def test_the_cpu_config_carries_the_spatial_transform_the_trainer_used_to_hardcode(self): + """The old CPU trainer appended a SpatialTransform with every probability at + 0 -- a no-op that only enforces the patch size. The merged trainer builds only + what the config names, so the config has to say it.""" + section = load_config(str(CONFIGS / "transform_params.json")).section(Backend.CPU) + self.assertIn("SpatialTransform", section) + spatial = section["SpatialTransform"] + self.assertEqual(spatial["p_rotation"], 0) + self.assertEqual(spatial["p_scaling"], 0) + self.assertEqual(spatial["p_elastic_deform"], 0) + self.assertEqual(spatial["mode_seg"], "nearest") + + +class TestDeepSupervisionPlacement(unittest.TestCase): + """Downsampling must follow whatever last deformed the mask. + + With GPU augmentations that is `train_step`, so it happens there; without them + nothing touches the mask after the dataloader and it belongs there. Getting this + backwards would train against targets that no longer match the image. + """ + + def test_a_gpu_config_leaves_downsampling_to_train_step(self): + for config in ("transform_params_gpu.json", "transform_params_hybrid.json"): + with self.subTest(config=config): + self.assertNotIn("DownsampleSegForDSTransform", training_transforms(config)) + + def test_a_cpu_only_config_downsamples_in_the_dataloader(self): + self.assertIn("DownsampleSegForDSTransform", training_transforms("transform_params.json")) + + def test_no_downsampling_when_no_scales_are_requested(self): + got = training_transforms("transform_params.json", deep_supervision_scales=None) + self.assertNotIn("DownsampleSegForDSTransform", got) + + +class TestDummy2D(unittest.TestCase): + def test_the_converters_bracket_the_spatial_transform(self): + got = training_transforms("transform_params_gpu.json", do_dummy_2d_data_aug=True) + spatial = got.index("SpatialTransform") + self.assertEqual(got[spatial - 1], "Convert3DTo2DTransform") + self.assertEqual(got[spatial + 1], "Convert2DTo3DTransform") + + +class TestConfigResolution(unittest.TestCase): + def setUp(self): + self._saved = {k: os.environ.get(k) for k in ("SMAUGLAB_PARAMS_JSON", "SMAUGLAB_PARAMS_GPU_JSON")} + for key in self._saved: + os.environ.pop(key, None) + + def tearDown(self): + for key, value in self._saved.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + def test_the_packaged_default_is_used_when_nothing_is_set(self): + from smauglab.trainers.nnUNetTrainerDAExt import resolve_config_path + + self.assertTrue(resolve_config_path().endswith("transform_params_gpu.json")) + + def test_the_new_variable_wins(self): + from smauglab.trainers.nnUNetTrainerDAExt import resolve_config_path + + os.environ["SMAUGLAB_PARAMS_JSON"] = "/new.json" + os.environ["SMAUGLAB_PARAMS_GPU_JSON"] = "/old.json" + self.assertEqual(resolve_config_path(), "/new.json") + + def test_the_old_variable_still_works_but_warns(self): + """segtransferaug/run_trainings.py sets it, and it drives every historical run.""" + from smauglab.trainers.nnUNetTrainerDAExt import resolve_config_path + + os.environ["SMAUGLAB_PARAMS_GPU_JSON"] = "/old.json" + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + self.assertEqual(resolve_config_path(), "/old.json") + self.assertTrue(any(issubclass(w.category, DeprecationWarning) for w in caught)) + + +class TestPipelineMode(unittest.TestCase): + def test_configs_default_to_the_sequential_pipeline(self): + self.assertIs(load_config(str(CONFIGS / "transform_params_gpu.json")).pipeline_mode(), PipelineMode.SEQUENTIAL) + + def test_the_list_configs_ask_for_the_random_order_pipeline(self): + """They carry a random_choose block, so they were written for the ChooseX + trainers; the config says that now instead of the class name.""" + for config in sorted(CONFIGS.glob("*-List*.json")): + with self.subTest(config=config.name): + self.assertIs(load_config(str(config)).pipeline_mode(), PipelineMode.RANDOM_ORDER) + + +if __name__ == "__main__": + unittest.main()