The official repo of "SAM-Med3D: Towards General-purpose Segmentation Models for Volumetric Medical Images".
- [Data] We have now released all of our dataset SA-Med3D-140K. Full Dataset Download Link: Huggingface. Labels Download Link: Baidu Netdisk and Google Drive.
- [Challenge] SAM-Med3D is invited as a baseline of CVPR-MedSegFMCompetition and the tutorial is here. We kindly invite you to join the challenge and build better foundation models for 3D medical image segmentation!
- [Examples] SAM-Med3D is now supported in MedIM, you can easily get our model with one-line Python code. Details can be found in
medim_val_single.py. - [Paper] SAM-Med3D is accepted as ECCV BIC 2024 Oral
- [Model] A newer version of finetuned SAM-Med3D named
SAM-Med3D-turbois released now. We fine-tuned it on 44 datasets (list) to improve the performance. Hope this update can help you π. - [Repos] If you are interested in computer vision, we recommend checking out OpenGVLab for more exciting projects like SAM-Med2D!
- π Curated the most extensive volumetric medical dataset to date for training, boasting 143K 3D masks and 245 categories.
- π€ Achieved efficient promptable segmentation, requiring 10 to 100 times fewer prompt points for satisfactory 3D outcomes.
- π Conducted a thorough assessment of SAM-Med3D across 16 frequently used volumetric medical image segmentation datasets.
SAM-Med3D-turbo: Hugging Face | Google Drive | Baidu NetDisk (pwd:l6ol)
Note: Ground-truth labels are required to generate prompt points. If you want to test an image without ground-truth, please generate a fake ground-truth with the target region for prompt annotated.
First, set up your environment with the following commands:
conda create --name sammed3d python=3.10
conda activate sammed3d
pip install uv
uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0
uv pip install torchio opencv-python-headless matplotlib prefetch_generator monai edt surface-distance medim
Then, use medim_val_single.py to test the model:
python medim_val_single.py
You could set your custom data in the code like:
img_path = "./test_data/Seg_Exps/ACDC/ACDC_test_cases/patient101_frame01_0000.nii.gz"
gt_path = "./test_data/Seg_Exps/ACDC/ACDC_test_gts/patient101_frame01.nii.gz"
out_path = "./test_data/Seg_Exps/ACDC/ACDC_test_SAM_Med3d/patient101_frame01.nii.gz"
(we recommend fine-tuning with SAM-Med3D pre-trained weights from link)
To train the SAM-Med3D model on your own data, follow these steps:
Note: You can easily get PyTorch SAM-Med3D model with pre-trained weights from huggingface use
MedIM.ckpt_path = "https://huggingface.co/blueyo0/SAM-Med3D/blob/main/sam_med3d_turbo.pth" model = medim.create_model("SAM-Med3D", pretrained=True, checkpoint_path=ckpt_path)
Download the checkpoint from ckpt section and move the pth file into SAM_Med3D/ckpt/ (We recommand to use SAM-Med3D-turbo.pth).
Ensure that your training data is organized according to the structure shown in the data/medical_preprocessed directories. The target file structures should be like the following:
data/medical_preprocessed
βββ adrenal
β βββ ct_WORD
β β βββ imagesTr
β β β βββ word_0025.nii.gz
β β β βββ ...
β β βββ labelsTr
β β β βββ word_0025.nii.gz
β β β βββ ...
βββ ...
If the original data are in the nnU-Net style, follow these steps:
For a nnU-Net style dataset, the original file structure should be:
Task010_WORD βββ imagesTr β βββ word_0025_0000.nii.gz β βββ ... βββ labelsTr β βββ word_0025.nii.gz β βββ ...Then you should resample and convert the masks into binary. (You can use script for nnU-Net folder)
data/train βββ adrenal β βββ ct_WORD β β βββ imagesTr β β β βββ word_0025.nii.gz β β β βββ ... β β βββ labelsTr β β β βββ word_0025.nii.gz (binary label) β β β βββ ... βββ liver β βββ ct_WORD β β βββ imagesTr β β β βββ word_0025.nii.gz β β β βββ ... β β βββ labelsTr β β β βββ word_0025.nii.gz (binary label) β β β βββ ... βββ ...
Then, modify img_datas in utils/data_paths.py according to your own data.
img_datas = [
"data/train/adrenal/ct_WORD",
"data/train/liver/ct_WORD",
...
]
or
PROJ_DIR = <YOUR PROJ DIR>
img_datas = glob(os.path.join(PROJ_DIR, "data", "train", "*", "*"))
You can refer to train.sh and train_ddp.sh for training.
Hint: Use the --checkpoint to set the pre-trained weight path, the model will be trained from scratch if no ckpt in the path is found!
@misc{wang2024sammed3dgeneralpurposesegmentationmodels,
title={SAM-Med3D: Towards General-purpose Segmentation Models for Volumetric Medical Images},
author={Haoyu Wang and Sizheng Guo and Jin Ye and Zhongying Deng and Junlong Cheng and Tianbin Li and Jianpin Chen and Yanzhou Su and Ziyan Huang and Yiqing Shen and Bin Fu and Shaoting Zhang and Junjun He and Yu Qiao},
year={2024},
eprint={2310.15161},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2310.15161},
}
This project is released under the Apache 2.0 license.
BTW, welcome to follow our Zhihu official account, we will share more information on medical imaging there.
- We thank all medical workers and dataset owners for making public datasets available to the community.
- Thanks to the open-source of the following projects:
- Hiring: We are hiring researchers, engineers, and interns in General Vision Group, Shanghai AI Lab. If you are interested in Medical Foundation Models and General Medical AI, including designing benchmark datasets, general models, evaluation systems, and efficient tools, please contact us.
- Global Collaboration: We're on a mission to redefine medical research, aiming for a more universally adaptable model. Our passionate team is delving into foundational healthcare models, promoting the development of the medical community. Collaborate with us to increase competitiveness, reduce risk, and expand markets.
- Contact: Junjun He(hejunjun@pjlab.org.cn), Jin Ye(yejin@pjlab.org.cn), and Tianbin Li (litianbin@pjlab.org.cn).