feat(export): add MONAI bundle export - #381
Open
dhritimandas wants to merge 1 commit into
Open
Conversation
Trained nobrainer models could not be consumed outside nobrainer: weights are a bare state_dict and the architecture lives only as a registry name in croissant.json. Export them as MONAI model-zoo bundles instead. - nobrainer/export/bundle.py: export_bundle(), build_metadata(), build_inference_config(). Reads a Segmentation.save() directory and writes LICENSE, configs/metadata.json, configs/inference.json, models/model.pt, docs/README.md. - CLI: `nobrainer export bundle MODEL_DIR OUTPUT` - pyproject: new [bundle] extra (fire, jsonschema) — required by `python -m monai.bundle verify_metadata`, which the exporter shells out to and which fails without them. Design notes: - spatial_shape uses the literal block_shape from provenance, not a symbolic divisibility expression. MONAI's grammar binds only `p` and `n`, which cannot express swin_unetr's real constraint (multiple of 32 AND >= 64). block_shape is also the SlidingWindowInferer roi_size, so metadata and inferer cannot drift. - All version fields read from live packages; nothing hardcoded. - torch.jit.script only. segformer3d/swin_unetr/bayesian nets fail to script (**kwargs in forward); that warns and omits models/model.ts, which is optional per spec. Tracing is opt-in (--trace) since it bakes in a fixed input shape. - Export self-verifies with two forward passes: output shape, and a determinism check. bayesian_meshnet/bayesian_vnet ignore mc=False and are stochastic, so they require --allow-stochastic.
dhritimandas
force-pushed
the
feat/monai-bundle-export
branch
from
August 12, 2026 14:20
e6ba5cb to
573b46d
Compare
Open
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
What: We add
nobrainer export bundle. It wraps a trained model as a MONAI model-zoo bundle. MONAI Label, MONAI Deploy, and the zoo consume this format.Why: Trained models could not leave nobrainer. The weights were a bare state_dict with nobrainer-only metadata. External tools could not load them.
Exporter (
nobrainer/export/bundle.py)Segmentation.save()dir (model.pth+croissant.json), reconstructs via the existingSegmentation.load()strict pathLICENSE,configs/metadata.json,configs/inference.json,models/model.pt,models/model.ts(when scriptable),docs/README.mdmonai.__version__,importlib.metadata); nothing hardcodedspatial_shape= literal provenanceblock_shapewithis_patch_data: true— MONAI's symbolic grammar binds onlyp/nand cannot expressswin_unetr's real constraint (multiple of 32 AND >= 64, measured);block_shapeis also theSlidingWindowInfererroi_size, so metadata and inferer cannot drifttorch.jit.scriptonly; on failure warns and omitsmodel.ts(optional per spec).--traceopt-in —segformer3d.forwardhas a shape-dependent branch a trace would freezebayesian_meshnet/bayesian_vnetignoremc=False(genuinely stochastic) and require--allow-stochasticautoencoder/simsiam/dcgan/progressivegan— no(B, n_classes, D, H, W)contractpython -m monai.bundle verify_metadataafter writing; missing-optional-import degrades to a warningCLI + deps
exportgroup +bundlesubcommand innobrainer/cli/main.py, existing conventionsbundle = ["fire", "jsonschema"]— both required by MONAI's ownverify_metadata; neither was declared anywhereTest plan
ConfigParserround-trip + strict load,verify_net_in_out, script-failure path, stochastic gate, CLI contract)python -m monai.bundle verify_metadataexits 0 on a fresh export