WholeBrain UNesT (133-class T1 brain parcellation) -- WholeBrain UNesT (MONAI wholeBrainSeg_Large_UNEST v0.2.7)

Description

WholeBrain UNesT, ported to JAX / Equinox from MONAI's wholeBrainSeg_Large_UNEST_segmentation bundle. UNesT pairs a Nested Transformer (NesT) 3D encoder -- which partitions each feature map into non-overlapping cubic blocks, applies transformer layers with attention localised within each block, and aggregates blocks between hierarchy levels via a strided conv-pool -- with a UNETR-style convolutional decoder. The network parcellates a T1 brain MRI patch into 133 whole-brain regions.

Intended use

Whole-brain parcellation of a T1-weighted brain MRI into 133 regions. Input is a single-channel 96x96x96 patch, nonzero channel-wise z-score normalised; output is 133-channel raw logits (argmax over the channel axis for the label map). The bundle tiles a full brain volume with sliding-window inference (roi 96^3, overlap 0.7) -- an out-of-model inference concern not vendored here.

Usage

from ilex.models.wholebrain_unest import WholeBrainUNesT
model = WholeBrainUNesT.from_pretrained('ilex-hub/wholebrain_unest.1')

Authors

Yu X., Yang Q., Zhou Y., Cai L. Y., Gao R., Lee H. H., et al.

Citation

Yu X., Yang Q., Zhou Y., Cai L. Y., Gao R., Lee H. H., et al. (2023). UNesT: Local Spatial Representation Learning with Hierarchical Transformer for Efficient Medical Segmentation. Medical Image Analysis. arXiv:2209.14378. NesT encoder: Zhang Z., et al. (2022), AAAI, arXiv:2105.12723.

References

  • Yu X., Yang Q., Zhou Y., et al. (2023). UNesT: Local Spatial Representation Learning with Hierarchical Transformer for Efficient Medical Segmentation. Medical Image Analysis. arXiv:2209.14378.
  • Zhang Z., Zhang H., Zhao L., Chen T., Arik S. O., Pfister T. (2022). Nested Hierarchical Transformer (NesT). AAAI 2022. arXiv:2105.12723.
  • Bundle: https://huggingface.co/MONAI/wholeBrainSeg_Large_UNEST_segmentation

License

HF Hub license tag: apache-2.0

Effective terms: Apache-2.0 (MONAI Consortium / MASI Lab, Vanderbilt) on both the UNesT network code and the wholeBrainSeg_Large_UNEST_segmentation bundle weights. The NesT encoder is adapted from the Nested Transformer (Zhang et al. 2022). No commercial restrictions; no gating required. The ilex JAX / Equinox port code is separately licensed under Apache-2.0 / GPL-3.0.

Upstream license reference: https://huggingface.co/MONAI/wholeBrainSeg_Large_UNEST_segmentation

Copyright

Network architecture and pretrained weights: copyright (c) the MONAI Consortium / MASI Lab (Vanderbilt), released under the Apache-2.0 License. The NesT encoder is adapted from the Nested Transformer (Zhang et al. 2022). JAX / Equinox port: copyright (c) the ilex authors, released under the Apache-2.0 / GPL-3.0 dual license used by ilex itself.

Upstream source

Original weights / reference implementation: https://huggingface.co/MONAI/wholeBrainSeg_Large_UNEST_segmentation

Provenance

This artefact was produced by ilex's save/load pipeline. The architecture is implemented in ilex.models.wholebrain_unest.WholeBrainUNesT and the weights have been converted from their upstream format. See the upstream source above for the canonical reference.

Downloads last month
18
Safetensors
Model size
87.2M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Papers for ilex-hub/wholebrain_unest.1