diff --git a/nobrainer/cli/main.py b/nobrainer/cli/main.py index 3b6ea348..55d11626 100644 --- a/nobrainer/cli/main.py +++ b/nobrainer/cli/main.py @@ -668,6 +668,140 @@ def zarr_suggest_shards(n_volumes, volume_shape, dtype, n_input_files, levels): click.echo(json.dumps(result, indent=2)) +# --------------------------------------------------------------------------- +# export subcommands +# --------------------------------------------------------------------------- + + +@cli.group() +def export(): + """Export trained models to interoperable formats.""" + + +@export.command("bundle") +@click.argument("model_dir", type=click.Path(exists=True)) +@click.argument("output", type=click.Path()) +@click.option( + "--no-torchscript", + is_flag=True, + help="Skip TorchScript export (models/model.ts).", +) +@click.option( + "--trace", + is_flag=True, + help=( + "Force torch.jit.trace instead of torch.jit.script. Bakes in a " + "fixed input shape; not the automatic fallback for a script failure." + ), +) +@click.option( + "--allow-stochastic", + is_flag=True, + help=( + "Allow exporting a model whose output differs across two identical " + "forward passes (Bayesian/MC models)." + ), +) +@click.option( + "--spatial-shape", + default=None, + help="Override spatial patch shape as D,H,W (default: provenance block_shape).", +) +@click.option( + "--version", + "bundle_version", + default="0.0.1", + help="Bundle version string.", + **_option_kwds, +) +@click.option("--name", default=None, help="Bundle display name.") +@click.option("--task", default=None, help="Task description.") +@click.option("--description", default=None, help="Longer description.") +@click.option( + "--authors", + default="nobrainer contributors", + help="Author string.", + **_option_kwds, +) +@click.option( + "--copyright", + "copyright_", + default="Copyright (c) nobrainer contributors", + help="Copyright string.", + **_option_kwds, +) +@click.option( + "--labels", + default=None, + help="Comma-separated class label names, index 0 (background) first.", +) +@click.option( + "--reference", + "references", + multiple=True, + help="Reference citation (repeatable).", +) +@click.option( + "--no-verify", + is_flag=True, + help="Skip `python -m monai.bundle verify_metadata` after writing.", +) +def export_bundle_cmd( + *, + model_dir, + output, + no_torchscript, + trace, + allow_stochastic, + spatial_shape, + bundle_version, + name, + task, + description, + authors, + copyright_, + labels, + references, + no_verify, +): + """Export a trained nobrainer model as a MONAI bundle. + + MODEL_DIR is a directory written by Segmentation.save() (model.pth + + croissant.json). OUTPUT is the bundle directory to create; it must not + already exist. + """ + from ..export.bundle import BundleExportError, export_bundle + + shape = None + if spatial_shape: + shape = tuple(int(x) for x in spatial_shape.split(",")) + label_list = labels.split(",") if labels else None + + try: + out = export_bundle( + model_dir, + output, + torchscript=not no_torchscript, + trace=trace, + allow_stochastic=allow_stochastic, + spatial_shape=shape, + version=bundle_version, + name=name, + task=task, + description=description, + authors=authors, + copyright_=copyright_, + labels=label_list, + references=list(references) or None, + verify=not no_verify, + ) + except BundleExportError as exc: + click.echo(click.style(f"ERROR: {exc}", fg="red")) + sys.exit(1) + + click.echo(click.style(f"Bundle exported: {out}", fg="green")) + + # For debugging only. if __name__ == "__main__": cli() diff --git a/nobrainer/export/__init__.py b/nobrainer/export/__init__.py new file mode 100644 index 00000000..355bfca6 --- /dev/null +++ b/nobrainer/export/__init__.py @@ -0,0 +1,5 @@ +"""Export trained nobrainer models to interoperable formats.""" + +from .bundle import export_bundle + +__all__ = ["export_bundle"] diff --git a/nobrainer/export/bundle.py b/nobrainer/export/bundle.py new file mode 100644 index 00000000..de9cddc9 --- /dev/null +++ b/nobrainer/export/bundle.py @@ -0,0 +1,789 @@ +"""Export a trained nobrainer estimator as a MONAI model-zoo bundle. + +See ``https://docs.monai.io/en/stable/mb_specification.html`` for the bundle +directory layout and ``configs/metadata.json`` schema this module targets. +""" + +from __future__ import annotations + +import importlib.metadata +import inspect +import json +from pathlib import Path +import shutil +import subprocess +import sys +import tempfile +from typing import Any +import warnings + +import monai +import numpy as np +import torch + +import nobrainer +from nobrainer.models import get as get_model + +MONAI_META_SCHEMA_URL = ( + "https://github.com/Project-MONAI/MONAI-extra-test-data/" + "releases/download/0.8.1/meta_schema_20240725.json" +) + +# Fully-convolutional / segmentation-contract architectures this exporter +# understands. Excludes autoencoder, simsiam (multi-input forward), +# dcgan/progressivegan (Lightning modules with no single-tensor forward) -- +# none of these describe a (B, n_classes, D, H, W) segmentation contract. +SUPPORTED_ARCHITECTURES = frozenset( + { + "unet", + "vnet", + "attention_unet", + "unetr", + "meshnet", + "highresnet", + "swin_unetr", + "segresnet", + "segformer3d", + "bayesian_meshnet", + "bayesian_vnet", + "kwyk_meshnet", + } +) + +NOT_FOR_CLINICAL_USE = ( + "**NOT FOR CLINICAL USE.** This model is exported for research purposes " + "only. It has not been evaluated, reviewed, or approved for clinical " + "diagnosis, treatment planning, or any other clinical use." +) + + +class BundleExportError(RuntimeError): + """Raised when a trained model cannot be exported as a MONAI bundle.""" + + +def _package_version(name: str) -> str: + """Return the installed version of a distribution. + + Parameters + ---------- + name : str + Distribution name as registered with ``importlib.metadata``. + + Returns + ------- + str + Installed version string, read live -- never hardcoded. + """ + return importlib.metadata.version(name) + + +def _extra_required_packages(base_model: str) -> dict[str, str]: + """Return extra required packages (beyond nobrainer/nibabel) for an architecture. + + Parameters + ---------- + base_model : str + Registry name of the exported architecture. + + Returns + ------- + dict of str to str + Package name mapped to installed version. + """ + extra: dict[str, str] = {} + if base_model == "segformer3d": + extra["einops"] = _package_version("einops") + if base_model in {"bayesian_meshnet", "bayesian_vnet", "kwyk_meshnet"}: + try: + extra["pyro-ppl"] = _package_version("pyro-ppl") + except importlib.metadata.PackageNotFoundError: + pass + return extra + + +def _resolve_in_channels(base_model: str, model_args: dict[str, Any]) -> int: + """Resolve ``in_channels`` for an architecture. + + Croissant provenance does not currently record ``in_channels`` + (see ``nobrainer/processing/croissant.py``), so it is read from + ``model_args`` if present, else from the factory's declared default. + + Parameters + ---------- + base_model : str + Registry name of the architecture. + model_args : dict + Stored ``model_args`` from the estimator's provenance. + + Returns + ------- + int + Number of input channels. + + Raises + ------ + BundleExportError + If ``in_channels`` cannot be resolved. + """ + if "in_channels" in model_args: + return int(model_args["in_channels"]) + factory = get_model(base_model) + default = inspect.signature(factory).parameters["in_channels"].default + if default is inspect.Parameter.empty: + raise BundleExportError( + f"Cannot resolve in_channels for '{base_model}': not present in " + "model_args and the factory has no default." + ) + return int(default) + + +def _resolve_spatial_shape( + base_model: str, + model_args: dict[str, Any], + block_shape: tuple[int, ...], + spatial_shape_override: tuple[int, int, int] | None, +) -> tuple[int, int, int]: + """Resolve the 3-D spatial patch shape used for the bundle contract. + + Uses the literal ``block_shape`` recorded at training time rather than a + symbolic divisibility expression. MONAI's ``spatial_shape`` grammar in + ``metadata.json`` only binds the variables ``p`` and ``n`` + (``monai.bundle.scripts._get_fake_spatial_shape``), which cannot express + "multiple of 32 and >= 64" -- the actual constraint measured for + ``swin_unetr``. ``block_shape`` is also exactly the value used for the + bundle's ``SlidingWindowInferer.roi_size``, so the two cannot drift. + + Parameters + ---------- + base_model : str + Registry name of the architecture. + model_args : dict + Stored ``model_args`` from provenance. + block_shape : tuple of int + ``block_shape`` recorded in the estimator's provenance. + spatial_shape_override : tuple of int, or None + User-supplied override, if any. + + Returns + ------- + tuple of int + Three-element spatial shape. + + Raises + ------ + BundleExportError + If no 3-D shape can be resolved, or (for ``unetr``) the resolved + shape does not match the ``img_size`` baked into the weights. + """ + if spatial_shape_override is not None: + shape = tuple(int(v) for v in spatial_shape_override) + elif len(block_shape) == 3: + shape = tuple(int(v) for v in block_shape) + else: + raise BundleExportError( + "block_shape in the model's provenance is missing or not 3-D " + f"(got {block_shape!r}). Pass an explicit spatial_shape override." + ) + + if base_model == "unetr": + img_size = model_args.get("img_size") + if img_size is not None and tuple(int(v) for v in img_size) != shape: + raise BundleExportError( + f"unetr weights were trained with img_size={tuple(img_size)}, " + f"which does not match the resolved spatial_shape {shape}. " + "UNETR bakes img_size into its weights; the two must match." + ) + return shape + + +def build_metadata( + *, + base_model: str, + in_channels: int, + n_classes: int, + spatial_shape: tuple[int, int, int], + provenance: dict[str, Any], + version: str, + name: str | None, + task: str | None, + description: str | None, + authors: str, + copyright_: str, + labels: list[str] | None, + references: list[str] | None, + stochastic: bool, +) -> dict[str, Any]: + """Build a MONAI bundle ``configs/metadata.json`` dict. + + Emits all 11 schema-required top-level keys plus the optional keys the + MONAI model-zoo bundles ship. Version fields are read from the live + installed packages, never hardcoded. + + Parameters + ---------- + base_model : str + Registry name of the exported architecture. + in_channels : int + Number of input channels. + n_classes : int + Number of output classes. + spatial_shape : tuple of int + Expected input/output spatial patch shape. + provenance : dict + The ``nobrainer:provenance`` block from ``croissant.json``. + version : str + Bundle version string. + name : str or None + Human-readable bundle name; a default is generated if None. + task : str or None + Task description; a default is generated if None. + description : str or None + Longer description; a default is generated if None. + authors : str + Author string. + copyright_ : str + Copyright string. + labels : list of str, or None + Class label names, index-ordered starting at background=0. Must have + length ``n_classes`` if given. + references : list of str, or None + Reference citations. + stochastic : bool + Whether the network was found to be non-deterministic across two + identical forward passes (Bayesian/MC models). + + Returns + ------- + dict + The ``metadata.json`` content. + + Raises + ------ + BundleExportError + If ``labels`` is given but its length does not match ``n_classes``. + """ + pytorch_version = torch.__version__.split("+")[0] + required_packages = { + "nobrainer": nobrainer.__version__, + "nibabel": _package_version("nibabel"), + **_extra_required_packages(base_model), + } + + if labels is None: + channel_def = {"0": "background"} + channel_def.update({str(i): f"class_{i}" for i in range(1, n_classes)}) + else: + if len(labels) != n_classes: + raise BundleExportError( + f"labels has {len(labels)} entries but n_classes={n_classes}." + ) + channel_def = {str(i): label for i, label in enumerate(labels)} + + intended_use = ( + "Research use only; not a substitute for expert diagnosis. " + + NOT_FOR_CLINICAL_USE + ) + if stochastic: + intended_use += ( + " This network is stochastic (Bayesian): each forward pass " + "returns one posterior draw, not a deterministic prediction." + ) + + best_loss = provenance.get("best_loss") + data_source = ", ".join( + d.get("path", "") for d in provenance.get("source_datasets", []) if d + ) + + return { + "schema": MONAI_META_SCHEMA_URL, + "version": version, + "changelog": {version: f"Exported from nobrainer {nobrainer.__version__}"}, + "monai_version": monai.__version__, + "pytorch_version": pytorch_version, + "numpy_version": np.__version__, + "required_packages_version": required_packages, + "name": name or f"Nobrainer {base_model} segmentation", + "task": task or "3D brain MRI segmentation", + "description": description + or ( + f"3-D brain MRI segmentation ({base_model}) exported from " + f"nobrainer {nobrainer.__version__}." + ), + "authors": authors, + "copyright": copyright_, + "data_source": data_source, + "data_type": "nibabel", + "image_classes": f"{in_channels}-channel MRI, intensity scaled to [0, 1]", + "label_classes": f"{n_classes} classes, one-hot", + "pred_classes": f"{n_classes} channels OneHot data", + "eval_metrics": {"best_loss": best_loss} if best_loss is not None else {}, + "intended_use": intended_use, + "references": references or [], + "network_data_format": { + "inputs": { + "image": { + "type": "image", + "format": "magnitude", + "modality": "MR", + "num_channels": in_channels, + "spatial_shape": list(spatial_shape), + "dtype": "float32", + "value_range": [0, 1], + "is_patch_data": True, + "channel_def": {"0": "image"}, + } + }, + "outputs": { + "pred": { + "type": "image", + # The MONAI spec reserves "labels" for N one-hot channels + # and "segmentation" for single-channel categorical + # output, but both real model-zoo bundles (spleen_ct, + # wholeBrainSeg) use "segmentation" for one-hot output. + # Matched here since downstream consumers target the zoo. + "format": "segmentation", + "num_channels": n_classes, + "spatial_shape": list(spatial_shape), + "dtype": "float32", + "value_range": [0, 1], + "is_patch_data": True, + "channel_def": channel_def, + } + }, + }, + } + + +def build_inference_config( + *, + base_model: str, + model_args: dict[str, Any], + in_channels: int, + n_classes: int, + spatial_shape: tuple[int, int, int], +) -> dict[str, Any]: + """Build a runnable MONAI bundle ``configs/inference.json`` dict. + + ``network_def._target_`` is a fully-qualified dotted path into + ``nobrainer.models`` (nobrainer factories are not in MONAI's + ``ComponentLocator`` namespace); ``monai.bundle.ConfigParser`` resolves + dotted paths via ``pydoc.locate`` (verified at plan time). + + Parameters + ---------- + base_model : str + Registry name of the architecture. + model_args : dict + Stored ``model_args`` from provenance (channels, strides, etc.). + in_channels : int + Number of input channels. + n_classes : int + Number of output classes. + spatial_shape : tuple of int + Expected spatial patch shape; used as the sliding-window ROI size. + + Returns + ------- + dict + The ``inference.json`` content. + """ + factory = get_model(base_model) + target = f"{factory.__module__}.{factory.__name__}" + + network_kwargs = { + k: v for k, v in model_args.items() if k not in ("n_classes", "in_channels") + } + network_kwargs["n_classes"] = n_classes + network_kwargs["in_channels"] = in_channels + + return { + "imports": ["$import glob"], + "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", + "ckpt_path": "$@bundle_root + '/models/model.pt'", + "dataset_dir": "/workspace/data", + "datalist": "$list(sorted(glob.glob(@dataset_dir + '/*.nii.gz')))", + "network_def": {"_target_": target, **network_kwargs}, + "network": "$@network_def.to(@device)", + "preprocessing": { + "_target_": "Compose", + "transforms": [ + {"_target_": "LoadImaged", "keys": "image"}, + {"_target_": "EnsureChannelFirstd", "keys": "image"}, + {"_target_": "ScaleIntensityd", "keys": "image"}, + {"_target_": "EnsureTyped", "keys": "image", "device": "@device"}, + ], + }, + "dataset": { + "_target_": "Dataset", + "data": "$[{'image': i} for i in @datalist]", + "transform": "@preprocessing", + }, + "dataloader": { + "_target_": "DataLoader", + "dataset": "@dataset", + "batch_size": 1, + "shuffle": False, + "num_workers": 0, + }, + "inferer": { + "_target_": "SlidingWindowInferer", + "roi_size": list(spatial_shape), + "sw_batch_size": 1, + "overlap": 0.25, + }, + "postprocessing": { + "_target_": "Compose", + "transforms": [ + {"_target_": "Activationsd", "keys": "pred", "softmax": True}, + {"_target_": "AsDiscreted", "keys": "pred", "argmax": True}, + { + "_target_": "SaveImaged", + "keys": "pred", + "meta_keys": "image_meta_dict", + "output_dir": "$@bundle_root + '/eval'", + }, + ], + }, + "handlers": [ + { + "_target_": "CheckpointLoader", + "load_path": "@ckpt_path", + "load_dict": {"model": "@network"}, + } + ], + "evaluator": { + "_target_": "SupervisedEvaluator", + "device": "@device", + "val_data_loader": "@dataloader", + "network": "@network", + "inferer": "@inferer", + "postprocessing": "@postprocessing", + "val_handlers": "@handlers", + }, + "evaluating": ["$@evaluator.run()"], + } + + +def _write_license(output_dir: Path) -> None: + """Write ``LICENSE`` into the bundle root. + + Copies nobrainer's own repository LICENSE file (Apache-2.0, per + ``pyproject.toml``) when it can be located relative to the installed + package; otherwise writes a short pointer to it. + + Parameters + ---------- + output_dir : Path + Bundle root directory. + """ + src = Path(nobrainer.__file__).resolve().parent.parent / "LICENSE" + dest = output_dir / "LICENSE" + if src.exists(): + shutil.copyfile(src, dest) + else: + dest.write_text( + "Apache License 2.0. See " + "https://github.com/neuronets/nobrainer/blob/main/LICENSE\n" + ) + + +def _write_readme( + docs_dir: Path, + *, + base_model: str, + n_classes: int, + in_channels: int, + spatial_shape: tuple[int, int, int], + stochastic: bool, +) -> None: + """Write ``docs/README.md`` with run instructions and a use disclaimer. + + Parameters + ---------- + docs_dir : Path + Bundle ``docs/`` directory. + base_model : str + Registry name of the architecture. + n_classes : int + Number of output classes. + in_channels : int + Number of input channels. + spatial_shape : tuple of int + Expected spatial patch shape. + stochastic : bool + Whether the network is non-deterministic across identical inputs. + """ + stochastic_note = ( + "\n**Note:** this network is stochastic (Bayesian); repeated runs " + "on the same input yield different outputs.\n" + if stochastic + else "" + ) + readme = f"""# Nobrainer {base_model} bundle + +{NOT_FOR_CLINICAL_USE} +{stochastic_note} +Architecture: `{base_model}` ({in_channels} input channel(s), {n_classes} \ +output classes). +Expected patch shape: {list(spatial_shape)}. + +## Run inference + +``` +python -m monai.bundle run \\ + --meta_file configs/metadata.json \\ + --config_file configs/inference.json \\ + --dataset_dir ./input \\ + --bundle_root . +``` +""" + (docs_dir / "README.md").write_text(readme) + + +def _verify_metadata_subprocess(meta_file: Path) -> None: + """Validate ``metadata.json`` against the MONAI bundle schema. + + Shells out to ``python -m monai.bundle verify_metadata`` (downloads the + schema referenced by the ``schema`` key and validates against it) so the + written bundle is checked by MONAI's own validator, not a reimplementation + of it. Requires the optional ``nobrainer[bundle]`` extra (``fire``, + ``jsonschema``) -- MONAI's own CLI entry point and ``verify_metadata`` + depend on them; nobrainer's core dependencies do not. + + The downloaded schema is cached at a stable path under the system temp + directory so repeated calls (e.g. across a test session) do not + re-download it. + + Parameters + ---------- + meta_file : Path + Path to the written ``configs/metadata.json``. + + Raises + ------ + BundleExportError + If ``verify_metadata`` reports an actual schema violation. + """ + schema_cache = ( + Path(tempfile.gettempdir()) / "nobrainer_monai_bundle_meta_schema.json" + ) + result = subprocess.run( + [ + sys.executable, + "-m", + "monai.bundle", + "verify_metadata", + "--meta_file", + str(meta_file), + "--filepath", + str(schema_cache), + ], + capture_output=True, + text=True, + ) + if result.returncode == 0: + return + + combined = result.stdout + result.stderr + if "OptionalImportError" in combined or "ModuleNotFoundError" in combined: + warnings.warn( + "Skipped MONAI schema validation: `python -m monai.bundle " + "verify_metadata` requires the optional 'fire' and 'jsonschema' " + "packages (install with `uv pip install -e '.[bundle]'`). The " + f"bundle was still written to disk.\n{combined}", + stacklevel=2, + ) + return + + raise BundleExportError( + "monai.bundle verify_metadata rejected the exported bundle:\n" + f"{result.stdout}\n{result.stderr}" + ) + + +def export_bundle( + model_dir: str | Path, + output_dir: str | Path, + *, + torchscript: bool = True, + trace: bool = False, + allow_stochastic: bool = False, + spatial_shape: tuple[int, int, int] | None = None, + version: str = "0.0.1", + name: str | None = None, + task: str | None = None, + description: str | None = None, + authors: str = "nobrainer contributors", + copyright_: str = "Copyright (c) nobrainer contributors", + labels: list[str] | None = None, + references: list[str] | None = None, + verify: bool = True, +) -> Path: + """Export a saved nobrainer estimator directory as a MONAI bundle. + + Parameters + ---------- + model_dir : str or Path + Directory written by ``Segmentation.save()`` (``model.pth`` + + ``croissant.json``). + output_dir : str or Path + Bundle directory to create. Must not already exist. + torchscript : bool + Attempt a TorchScript export (``models/model.ts``). On failure a + warning is emitted and ``model.ts`` is omitted -- it is optional per + the bundle spec. + trace : bool + Force ``torch.jit.trace`` instead of ``torch.jit.script``. Bakes in + a fixed input shape; not the automatic fallback for a script + failure. + allow_stochastic : bool + Required to export a network whose output is non-deterministic + across two identical forward passes (Bayesian/MC models). + spatial_shape : tuple of int, or None + Override for the patch shape; defaults to the provenance + ``block_shape``. + version, name, task, description, authors, copyright_, labels, references + Passed through to :func:`build_metadata`. + verify : bool + Run ``monai.bundle verify_metadata`` on the written bundle and raise + on failure. + + Returns + ------- + Path + Path to the written bundle directory (``output_dir``). + + Raises + ------ + BundleExportError + If the output directory exists, the architecture is unsupported, + the spatial shape cannot be resolved, the model is stochastic + without ``allow_stochastic``, or ``verify_metadata`` rejects the + result. + """ + from monai.networks.utils import convert_to_torchscript, save_state + + from nobrainer.processing.segmentation import Segmentation + + model_dir = Path(model_dir) + output_dir = Path(output_dir) + if output_dir.exists(): + raise BundleExportError(f"Output directory already exists: {output_dir}") + + estimator = Segmentation.load(model_dir) + base_model = estimator.base_model + if base_model not in SUPPORTED_ARCHITECTURES: + raise BundleExportError( + f"Architecture '{base_model}' is not supported for bundle " + f"export. Supported: {sorted(SUPPORTED_ARCHITECTURES)}." + ) + + n_classes = estimator.n_classes_ + if not n_classes: + raise BundleExportError( + "n_classes missing from the model's provenance; cannot export." + ) + model_args = dict(estimator.model_args) + in_channels = _resolve_in_channels(base_model, model_args) + shape = _resolve_spatial_shape( + base_model, model_args, tuple(estimator.block_shape_ or ()), spatial_shape + ) + + net = estimator.model_ + net.eval() + probe = torch.rand(1, in_channels, *shape) + with torch.no_grad(): + out_a = net(probe) + out_b = net(probe) + if tuple(out_a.shape) != (1, n_classes, *shape): + raise BundleExportError( + f"Self-verification failed: expected output shape " + f"(1, {n_classes}, {tuple(shape)}), got {tuple(out_a.shape)}." + ) + is_stochastic = not torch.allclose(out_a, out_b) + if is_stochastic and not allow_stochastic: + raise BundleExportError( + f"'{base_model}' produced different output on two identical " + "forward passes (stochastic/Bayesian network). Pass " + "allow_stochastic=True to export it anyway; the resulting " + "bundle documents that inference yields one posterior draw." + ) + + provenance = json.loads((model_dir / "croissant.json").read_text()).get( + "nobrainer:provenance", {} + ) + + metadata = build_metadata( + base_model=base_model, + in_channels=in_channels, + n_classes=n_classes, + spatial_shape=shape, + provenance=provenance, + version=version, + name=name, + task=task, + description=description, + authors=authors, + copyright_=copyright_, + labels=labels, + references=references, + stochastic=is_stochastic, + ) + inference_config = build_inference_config( + base_model=base_model, + model_args=model_args, + in_channels=in_channels, + n_classes=n_classes, + spatial_shape=shape, + ) + + configs_dir = output_dir / "configs" + models_dir = output_dir / "models" + docs_dir = output_dir / "docs" + configs_dir.mkdir(parents=True) + models_dir.mkdir(parents=True) + docs_dir.mkdir(parents=True) + + (configs_dir / "metadata.json").write_text(json.dumps(metadata, indent=4)) + (configs_dir / "inference.json").write_text(json.dumps(inference_config, indent=4)) + + save_state(net, str(models_dir / "model.pt")) + + if torchscript: + try: + if trace: + convert_to_torchscript( + model=net, + filename_or_obj=str(models_dir / "model.ts"), + inputs=[probe], + use_trace=True, + ) + else: + convert_to_torchscript( + model=net, + filename_or_obj=str(models_dir / "model.ts"), + ) + except Exception as exc: # noqa: BLE001 - any scripting/tracing failure + warnings.warn( + f"torch.jit.{'trace' if trace else 'script'} failed for " + f"'{base_model}': {type(exc).__name__}: {exc}. Skipping " + "models/model.ts; the bundle remains spec-valid (model.ts " + "is optional).", + stacklevel=2, + ) + (models_dir / "model.ts").unlink(missing_ok=True) + + _write_license(output_dir) + _write_readme( + docs_dir, + base_model=base_model, + n_classes=n_classes, + in_channels=in_channels, + spatial_shape=shape, + stochastic=is_stochastic, + ) + + if verify: + _verify_metadata_subprocess(configs_dir / "metadata.json") + + return output_dir diff --git a/nobrainer/tests/unit/test_bundle_export.py b/nobrainer/tests/unit/test_bundle_export.py new file mode 100644 index 00000000..02153ab4 --- /dev/null +++ b/nobrainer/tests/unit/test_bundle_export.py @@ -0,0 +1,399 @@ +"""Tests for `nobrainer export bundle` (nobrainer/export/bundle.py).""" + +from __future__ import annotations + +import json +from pathlib import Path +import subprocess +import sys + +import monai +from monai.bundle import ConfigParser, verify_net_in_out +import numpy as np +import pytest +import torch + +import nobrainer +from nobrainer.export.bundle import ( + SUPPORTED_ARCHITECTURES, + BundleExportError, + build_metadata, + export_bundle, +) +from nobrainer.models import get as get_model +from nobrainer.processing.segmentation import Segmentation + +REQUIRED_METADATA_KEYS = { + "schema", + "version", + "monai_version", + "pytorch_version", + "numpy_version", + "required_packages_version", + "task", + "description", + "authors", + "copyright", + "network_data_format", +} +REQUIRED_TENSOR_KEYS = { + "type", + "format", + "num_channels", + "spatial_shape", + "dtype", + "value_range", +} + + +def _save_estimator( + tmp_path: Path, + *, + base_model: str = "unet", + model_args: dict | None = None, + n_classes: int = 3, + block_shape: tuple[int, int, int] | list = (16, 16, 16), + dirname: str = "my_model", +) -> Path: + """Build a tiny trained-looking estimator and save it like ``fit()`` would.""" + model_args = ( + model_args + if model_args is not None + else { + "channels": (4, 8), + "strides": (2,), + } + ) + est = Segmentation(base_model, model_args=model_args) + est.model_ = get_model(base_model)(n_classes=n_classes, **model_args) + est.n_classes_ = n_classes + est.block_shape_ = tuple(block_shape) if block_shape else block_shape + est.volume_shape_ = tuple(block_shape) if block_shape else block_shape + est._training_result = {"history": [{"loss": 0.5}, {"loss": 0.3}]} + est._dataset = None + save_dir = tmp_path / dirname + est.save(save_dir) + return save_dir + + +@pytest.fixture +def unet_model_dir(tmp_path: Path) -> Path: + return _save_estimator(tmp_path) + + +@pytest.fixture +def exported_bundle(tmp_path: Path, unet_model_dir: Path) -> Path: + return export_bundle(unet_model_dir, tmp_path / "MyBundle", verify=False) + + +class TestLayout: + def test_required_files_present(self, exported_bundle: Path) -> None: + assert (exported_bundle / "LICENSE").is_file() + assert (exported_bundle / "configs" / "metadata.json").is_file() + assert (exported_bundle / "configs" / "inference.json").is_file() + assert (exported_bundle / "models" / "model.pt").is_file() + assert (exported_bundle / "docs" / "README.md").is_file() + + def test_output_dir_must_not_exist( + self, tmp_path: Path, unet_model_dir: Path + ) -> None: + out = tmp_path / "AlreadyThere" + out.mkdir() + with pytest.raises(BundleExportError): + export_bundle(unet_model_dir, out, verify=False) + + def test_readme_has_not_for_clinical_use_disclaimer( + self, exported_bundle: Path + ) -> None: + readme = (exported_bundle / "docs" / "README.md").read_text() + assert "NOT FOR CLINICAL USE" in readme + + +class TestMetadataSchema: + def test_required_top_level_keys_present(self, exported_bundle: Path) -> None: + metadata = json.loads( + (exported_bundle / "configs" / "metadata.json").read_text() + ) + missing = REQUIRED_METADATA_KEYS - metadata.keys() + assert not missing, f"missing required metadata keys: {missing}" + + def test_required_tensor_keys_present(self, exported_bundle: Path) -> None: + metadata = json.loads( + (exported_bundle / "configs" / "metadata.json").read_text() + ) + ndf = metadata["network_data_format"] + for block in (ndf["inputs"]["image"], ndf["outputs"]["pred"]): + missing = REQUIRED_TENSOR_KEYS - block.keys() + assert not missing, f"missing required tensor keys: {missing}" + + def test_versions_are_read_live_not_hardcoded(self, exported_bundle: Path) -> None: + metadata = json.loads( + (exported_bundle / "configs" / "metadata.json").read_text() + ) + assert metadata["monai_version"] == monai.__version__ + assert metadata["numpy_version"] == np.__version__ + assert "+" not in metadata["pytorch_version"] + assert metadata["pytorch_version"] == torch.__version__.split("+")[0] + assert ( + metadata["required_packages_version"]["nobrainer"] == nobrainer.__version__ + ) + + def test_spatial_shape_matches_block_shape(self, exported_bundle: Path) -> None: + metadata = json.loads( + (exported_bundle / "configs" / "metadata.json").read_text() + ) + inference = json.loads( + (exported_bundle / "configs" / "inference.json").read_text() + ) + image_shape = metadata["network_data_format"]["inputs"]["image"][ + "spatial_shape" + ] + pred_shape = metadata["network_data_format"]["outputs"]["pred"]["spatial_shape"] + assert image_shape == [16, 16, 16] + assert pred_shape == [16, 16, 16] + assert inference["inferer"]["roi_size"] == [16, 16, 16] + + def test_channel_def_background_and_count(self, exported_bundle: Path) -> None: + metadata = json.loads( + (exported_bundle / "configs" / "metadata.json").read_text() + ) + channel_def = metadata["network_data_format"]["outputs"]["pred"]["channel_def"] + assert len(channel_def) == 3 + assert channel_def["0"] == "background" + assert set(channel_def) == {"0", "1", "2"} + + def test_channel_def_labels_override(self) -> None: + metadata = build_metadata( + base_model="unet", + in_channels=1, + n_classes=3, + spatial_shape=(16, 16, 16), + provenance={}, + version="0.0.1", + name=None, + task=None, + description=None, + authors="a", + copyright_="c", + labels=["bg", "gray", "white"], + references=None, + stochastic=False, + ) + channel_def = metadata["network_data_format"]["outputs"]["pred"]["channel_def"] + assert channel_def == {"0": "bg", "1": "gray", "2": "white"} + + def test_labels_length_mismatch_raises(self) -> None: + with pytest.raises(BundleExportError): + build_metadata( + base_model="unet", + in_channels=1, + n_classes=3, + spatial_shape=(16, 16, 16), + provenance={}, + version="0.0.1", + name=None, + task=None, + description=None, + authors="a", + copyright_="c", + labels=["only_one"], + references=None, + stochastic=False, + ) + + +class TestRoundTrip: + def test_network_def_resolves_and_loads_state_dict( + self, exported_bundle: Path + ) -> None: + parser = ConfigParser() + parser.read_config(str(exported_bundle / "configs" / "inference.json")) + net = parser.get_parsed_content("network_def") + assert isinstance(net, torch.nn.Module) + state = torch.load(exported_bundle / "models" / "model.pt", weights_only=True) + net.load_state_dict(state, strict=True) + + def test_verify_net_in_out(self, exported_bundle: Path) -> None: + verify_net_in_out( + net_id="network_def", + meta_file=str(exported_bundle / "configs" / "metadata.json"), + config_file=str(exported_bundle / "configs" / "inference.json"), + device="cpu", + ) + + +class TestSpatialShapeResolution: + def test_missing_block_shape_raises(self, tmp_path: Path) -> None: + model_dir = _save_estimator(tmp_path, block_shape=()) + with pytest.raises(BundleExportError): + export_bundle(model_dir, tmp_path / "Bundle", verify=False) + + def test_spatial_shape_override(self, tmp_path: Path) -> None: + model_dir = _save_estimator(tmp_path, block_shape=()) + out = export_bundle( + model_dir, + tmp_path / "Bundle", + spatial_shape=(16, 16, 16), + verify=False, + ) + metadata = json.loads((out / "configs" / "metadata.json").read_text()) + assert metadata["network_data_format"]["inputs"]["image"]["spatial_shape"] == [ + 16, + 16, + 16, + ] + + +class TestArchitectureScope: + def test_rejected_architecture_raises(self, tmp_path: Path) -> None: + model_dir = _save_estimator( + tmp_path, + base_model="autoencoder", + model_args={"input_shape": (16, 16, 16)}, + n_classes=1, + ) + with pytest.raises(BundleExportError, match="autoencoder"): + export_bundle(model_dir, tmp_path / "Bundle", verify=False) + + def test_supported_architectures_excludes_non_segmentation_nets(self) -> None: + assert "autoencoder" not in SUPPORTED_ARCHITECTURES + assert "simsiam" not in SUPPORTED_ARCHITECTURES + assert "dcgan" not in SUPPORTED_ARCHITECTURES + assert "progressivegan" not in SUPPORTED_ARCHITECTURES + + +class TestTorchScript: + def test_script_failure_warns_and_omits_model_ts(self, tmp_path: Path) -> None: + model_dir = _save_estimator( + tmp_path, + base_model="swin_unetr", + model_args={"feature_size": 12}, + n_classes=2, + block_shape=(64, 64, 64), + ) + with pytest.warns(UserWarning, match="torch.jit.script failed"): + out = export_bundle(model_dir, tmp_path / "Bundle", verify=False) + assert not (out / "models" / "model.ts").exists() + assert (out / "models" / "model.pt").exists() + assert (out / "configs" / "metadata.json").exists() + + def test_script_success_writes_model_ts(self, exported_bundle: Path) -> None: + assert (exported_bundle / "models" / "model.ts").exists() + + +class TestStochasticGate: + def test_stochastic_model_requires_allow_stochastic(self, tmp_path: Path) -> None: + model_dir = _save_estimator( + tmp_path, + base_model="bayesian_meshnet", + model_args={"filters": 8, "receptive_field": 37}, + n_classes=2, + block_shape=(8, 8, 8), + ) + with pytest.raises(BundleExportError, match="stochastic|allow_stochastic"): + export_bundle(model_dir, tmp_path / "Bundle", verify=False) + + def test_stochastic_model_exports_with_flag(self, tmp_path: Path) -> None: + model_dir = _save_estimator( + tmp_path, + base_model="bayesian_meshnet", + model_args={"filters": 8, "receptive_field": 37}, + n_classes=2, + block_shape=(8, 8, 8), + ) + out = export_bundle( + model_dir, tmp_path / "Bundle", allow_stochastic=True, verify=False + ) + metadata = json.loads((out / "configs" / "metadata.json").read_text()) + assert "posterior draw" in metadata["intended_use"] + + +class TestCLIContract: + def _help(self, cmd: list[str]) -> str: + result = subprocess.run( + [sys.executable, "-m", "nobrainer.cli.main"] + cmd + ["--help"], + capture_output=True, + text=True, + ) + assert ( + result.returncode == 0 + ), f"'{' '.join(cmd)} --help' exited {result.returncode}:\n{result.stderr}" + return result.stdout + + def test_export_bundle_help_exits_zero(self) -> None: + self._help(["export", "bundle"]) + + @pytest.mark.parametrize( + "option", + [ + "--no-torchscript", + "--trace", + "--allow-stochastic", + "--spatial-shape", + "--version", + "--name", + "--task", + "--description", + "--authors", + "--copyright", + "--labels", + "--reference", + "--no-verify", + ], + ) + def test_export_bundle_has_option(self, option: str) -> None: + out = self._help(["export", "bundle"]) + assert option in out + + def test_export_bundle_cli_end_to_end(self, tmp_path: Path) -> None: + model_dir = _save_estimator(tmp_path) + out_dir = tmp_path / "CliBundle" + result = subprocess.run( + [ + sys.executable, + "-m", + "nobrainer.cli.main", + "export", + "bundle", + str(model_dir), + str(out_dir), + "--no-verify", + ], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert (out_dir / "configs" / "metadata.json").is_file() + + +class TestVerifyMetadataSubprocess: + """Exercises the `python -m monai.bundle verify_metadata` integration. + + Requires the optional `nobrainer[bundle]` extra (fire, jsonschema) and + network access to fetch the schema once; skips gracefully otherwise so + the rest of the suite is not network-dependent. + """ + + def test_verify_metadata_exits_zero_on_fresh_export(self, tmp_path: Path) -> None: + pytest.importorskip("fire") + pytest.importorskip("jsonschema") + model_dir = _save_estimator(tmp_path) + try: + out = export_bundle(model_dir, tmp_path / "Bundle", verify=True) + except BundleExportError as exc: + pytest.fail(f"verify_metadata rejected a fresh export: {exc}") + result = subprocess.run( + [ + sys.executable, + "-m", + "monai.bundle", + "verify_metadata", + "--meta_file", + str(out / "configs" / "metadata.json"), + "--filepath", + str(tmp_path / "schema_cache.json"), + ], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/pyproject.toml b/pyproject.toml index ecac034d..a653fa5c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,9 +64,10 @@ zarr = ["zarr >= 3.0", "nifti-zarr", "ome-zarr >= 0.14.0", "dask[array]", "scipy croissant = ["mlcroissant"] versioning = ["datalad >= 0.19"] tfrecord = ["tfrecord >= 1.14"] +bundle = ["fire", "jsonschema"] dev = ["pre-commit", "pytest", "pytest-cov", "scipy"] all = [ - "nobrainer[bayesian,generative,zarr,croissant,versioning,tfrecord,dev]", + "nobrainer[bayesian,generative,zarr,croissant,versioning,tfrecord,bundle,dev]", ] [tool.hatch.version]