Skip to content

MrGiovanni/SuPreM

Repository files navigation

We developed a suite of pre-trained 3D models, named SuPreM, that combined the best of large-scale datasets and per-voxel annotations, showing the transferability across a range of 3D medical imaging tasks.

Paper

How Well Do Supervised 3D Models Transfer to Medical Imaging Tasks?
Wenxuan Li, Alan Yuille, and Zongwei Zhou*
Johns Hopkins University
International Conference on Learning Representations (ICLR) 2024 (oral; top 1.2%)
YouTube

Transitioning to Fully-Supervised Pre-Training with Large-Scale Radiology ImageNet for Improved AI Transferability in Three-Dimensional Medical Segmentation
Wenxuan Li1, Junfei Xiao1, Jie Liu2, Yucheng Tang3, Alan Yuille1, and Zongwei Zhou1,*
1Johns Hopkins University
2City University of Hong Kong
3NVIDIA
Radiological Society of North America (RSNA) 2023

★ We have maintained a document for Frequently Asked Questions.

★ We have maintained a paper list for Awesome Medical Pre-Training.

★ We have maintained a paper list for Awesome Medical Segmentation Backbones.

An Extensive Dataset: AbdomenAtlas 1.1

The release of AbdomenAtlas 1.0 can be found at https://huggingface.co/datasets/AbdomenAtlas/AbdomenAtlas1.0Mini

AbdomenAtlas 1.1 is an extensive dataset of 9,262 CT volumes with per-voxel annotation of 25 organs and pseudo annotations for seven types of tumors, enabling us to finally perform supervised pre-training of AI models at scale. Based on AbdomenAtlas 1.1, we also provide a suite of pre-trained models comprising several widely recognized AI models.

Prelimianry benchmark showed that supervised pre-training strikes as a preferred choice in terms of performance and efficiency compared with self-supervised pre-training.

We anticipate that the release of large, annotated datasets (AbdomenAtlas 1.1) and the suite of pre-trained models (SuPreM) will bolster collaborative endeavors in establishing Foundation Datasets and Foundation Models for the broader applications of 3D volumetric medical image analysis.

The AbdomenAtlas 1.1 dataset is organized as

AbdomenAtlas1.1
    ├── BDMAP_00000001
    │   ├── ct.nii.gz
    │   └── segmentations
    │       ├── aorta.nii.gz
    │       ├── gall_bladder.nii.gz
    │       ├── kidney_left.nii.gz
    │       ├── kidney_right.nii.gz
    │       ├── liver.nii.gz
    │       ├── pancreas.nii.gz
    │       ├── postcava.nii.gz
    │       ├── spleen.nii.gz
    │       ├── stomach.nii.gz
    │       └── ...
    ├── BDMAP_00000002
    │   ├── ct.nii.gz
    │   └── segmentations
    │       ├── aorta.nii.gz
    │       ├── gall_bladder.nii.gz
    │       ├── kidney_left.nii.gz
    │       ├── kidney_right.nii.gz
    │       ├── liver.nii.gz
    │       ├── pancreas.nii.gz
    │       ├── postcava.nii.gz
    │       ├── spleen.nii.gz
    │       ├── stomach.nii.gz
    │       └── ...
    ├── BDMAP_00000003
    │   ├── ct.nii.gz
    │   └── segmentations
    │       ├── aorta.nii.gz
    │       ├── gall_bladder.nii.gz
    │       ├── kidney_left.nii.gz
    │       ├── kidney_right.nii.gz
    │       ├── liver.nii.gz
    │       ├── pancreas.nii.gz
    │       ├── postcava.nii.gz
    │       ├── spleen.nii.gz
    │       ├── stomach.nii.gz
    │       └── ...
    ...

A Suite of Pre-trained Models: SuPreM

The following is a list of supported model backbones in our collection. Select the appropriate family of backbones and click to expand the table, download a specific backbone and its pre-trained weights (name and download), and save the weights into ./pretrained_weights/. More backbones will be added along time. Please suggest the backbone in this channel if you want us to pre-train it on AbdomenAtlas 1.1 containing 9,262 annotated CT volumes.

Swin UNETR
name params pre-trained data resources download
Tang et al. 62.19M 5050 CT GitHub stars weights
Jose Valanaras et al. 62.19M 50000 CT/MRI GitHub stars weights
Universal Model 62.19M 2100 CT GitHub stars weights
SuPreM 62.19M 2100 CT ours 🌟 weights
U-Net
name params pre-trained data resources download
Models Genesis 19.08M 623 CT GitHub stars weights
UniMiSS tiny 5022 CT&MRI GitHub stars weights
small 5022 CT&MRI weights
Med3D 85.75M 1638 CT GitHub stars weights
DoDNet 17.29M 920 CT GitHub stars weights
Universal Model 19.08M 2100 CT GitHub stars weights
SuPreM 19.08M 2100 CT ours 🌟 weights
SegResNet
name params pre-trained data resources download
SuPreM 4.70M 2100 CT ours 🌟 weights

Examples of predicting organ masks on unseen CT volumes using our SuPreM: README

Examples of fine-tuning our SuPreM on other downstream medical tasks are provided in this repository.

task dataset document
organ, muscle, vertebrae, cardiac, rib segmentation TotalSegmentator README
pancreas tumor detection JHH README

If you want to re-pre-train SuPreM on AbdomenAtlas 1.1 (not recommended), please refer to our instruction

Estimated cost:

  • Eight A100 GPUs
  • At least seven days
  • 733GB disk space

★ Or simply make a request here: #1

Citation

@inproceedings{li2024well,
  title={How Well Do Supervised Models Transfer to 3D Image Segmentation?},
  author={Li, Wenxuan and Yuille, Alan and Zhou, Zongwei},
  booktitle={The Twelfth International Conference on Learning Representations},
  year={2024}
}

@article{qu2023abdomenatlas,
  title={Abdomenatlas-8k: Annotating 8,000 CT volumes for multi-organ segmentation in three weeks},
  author={Qu, Chongyu and Zhang, Tiezheng and Qiao, Hualin and Tang, Yucheng and Yuille, Alan L and Zhou, Zongwei and others},
  journal={Advances in Neural Information Processing Systems},
  volume={36},
  year={2023}
}

Acknowledgement

This work was supported by the Lustgarten Foundation for Pancreatic Cancer Research and the McGovern Foundation. The codebase is modified from NVIDIA MONAI. Paper content is covered by patents pending.