Unleashing the Potential of SAM for Medical Adaptation via Hierarchical Decoding
Zhiheng Cheng, Qingyue Wei, Hongru Zhu, Yan Wang, Liangqiong Qu, Wei Shao, and Yuyin Zhou
CVPR 2024
paper | code
git clone https://github.com/Cccccczh404/H-SAM.gitPlease run the following commands to create an environment and obtain requirements.
conda create -n H-SAM python=3.10
conda activate H-SAM
pip install -r requirements.txt1.1 Please download Synapse multi-organ CT dataset from the original website (https://www.synapse.org/#!Synapse:syn3193805/wiki/217789) or the preprocessed version (https://drive.google.com/file/d/1zuOQRyfo0QYgjcU_uZs0X3LdCnAC2m3G/view?usp=share_link), unzip the [training set] with resolution 224x224, and put it in <Your folder>. We also prepare the [training set] with resolution 512x512 (https://drive.google.com/file/d/1F42WMa80UpH98Pw95oAzYDmxAAO2ApYg/view?usp=share_link). The 224x224 version of training set is downsampled from the 512x512 version.
1.2 Please download and unzip the [testset] with resolution 512x512 (https://drive.google.com/file/d/1RczbNSB37OzPseKJZ1tDxa5OO1IIICzK/view?usp=share_link) and put it in the ./testset folder. Then, unzip and delete this file.
1.3 Please download the pretrained SAM models from the original SAM repository (https://github.com/facebookresearch/segment-anything), and put them in the ./checkpoints folder.
Use the train.py file for training models. An example script of is
CUDA_VISIBLE_DEVICES="0,1" python train.py --root_path <Your folder>/train_npz_new_224/ --split='train' --batch_size=8 --base_lr=0.0026 --img_size=224 --warmup --AdamW --max_epochs=300 --stop_epoch=300 --vit_name='vit_l' --ckpt='checkpoints/sam_vit_l_0b3195.pth'
To use the few-shot training strategy with 10% of the training data, please first run the following command to rename the train.txt file.
mv lists/lists_Synapse/train.txt lists/lists_Synapse/train_full.txt
Use the randomly generated 10% training list.
mv lists/lists_Synapse/train_220.txt lists/lists_Synapse/train.txt
Or generate a new training list by running the following command.
python preprocess/make_list.py
An example script under few-shot setting is
CUDA_VISIBLE_DEVICES="0,1" python train.py --root_path <Your folder>/train_npz/ --split='train' --batch_size=8 --base_lr=0.0025 --img_size=512 --warmup --AdamW --max_epochs=300 --stop_epoch=300 --vit_name='vit_b' --ckpt='checkpoints/sam_vit_b_01ec64.pth'
Use the test.py file for testing models. An example script is
CUDA_VISIBLE_DEVICES="0" python test.py --is_savenii --lora_ckpt outputs/Synapse_224_pretrain_vit_l_epo300_bs8_lr0.0026/epoch_299.pth --vit_name='vit_l' --ckpt=checkpoints/sam_vit_l_0b3195.pth --img_size=224 --stage=3
#### 10% training data: https://drive.google.com/file/d/1A1wqTT_FD6FXzqCj9LNStxT425ZfYpnE/view?usp=sharing
CUDA_VISIBLE_DEVICES="0" python test.py --is_savenii --lora_ckpt 220_epoch_299.pth --vit_name='vit_b' --ckpt=checkpoints/sam_vit_b_01ec64.pth --stage=3 --img_size=512
If you find this code useful in your research, please consider citing:
@article{cheng2024unleashing,
title={Unleashing the Potential of SAM for Medical Adaptation via Hierarchical Decoding},
author={Cheng, Zhiheng and Wei, Qingyue and Zhu, Hongru and Wang, Yan and Qu, Liangqiong and Shao, Wei and Zhou, Yuyin},
journal={arXiv preprint arXiv:2403.18271},
year={2024}
}We appreciate the developers of Segment Anything Model and the provider of the Synapse multi-organ segmentation dataset. Our code of H-SAM is built upon SAMed, and we express our gratitude to these awesome projects.
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent