Implementation of "Segment Anything Model for Medical Images?" in pytorch --for finetuning the SAM with box prompts.

Arxiv link: https://arxiv.org/pdf/2304.14660.pdf

MIA version link: https://www.sciencedirect.com/science/article/pii/S1361841523003213

Our work has been accepted by Medical Image Analysis (MedIA) 2023!

Usage

1. Prepare data and segmentation information json files corresponding to each test set.

{
   "info": {
       "1": "LeftVentricle",
       "2": "LeftVentricularMyocardium",
       "3": "RightVentricle"
   },
   "color": {
       "1": 85,
       "2": 170,
       "3": 255
   }
}

 train: ../data/train_data/images/
 val: ../data/train_data/images/
 test: ../data/test_data/dataset_name/images/
 
 ├── train_data          
 │   ├── images        
 │   │   ├── 000001.png
 │   │   ├── 000002.png
 │   │   └── 000003.png
 │   └── labels         
 │       ├── 00001.png
 │       ├── 00002.png
 │       └── 00003.png
 └── val_data           
 |   ├── images        
 |   │   ├── 000001.png
 |   │   ├── 000002.png
 |   │   └── 000003.png
 |   └── labels         
 |       ├── 00001.png
 |       ├── 00002.png
 |       └── 00003.png
 └── test_data          
     ├── dataset1        
     │   ├── images
     |   |     ├── 000001.png
     |   │     ├── 000002.png
     |   |     └── 000003.png
     │   └── labels
     |         ├── 000001.png
     |         ├── 000002.png
     |         └── 000003.png
     └── dataset2         
         ├── images
         |     ├── 000001.png
         │     ├── 000002.png
         |     └── 000003.png
         └── labels
               ├── 000001.png
               ├── 000002.png
               └── 000003.png

2. Generate embedding for each single image.

$ python pre_grey_rgb2D.py         --img_path  data/train_data/images    --gt_path data/train_data/labels                       --checkpoint sam_vit_b_01ec64.pth              #for preparing training data (embeddings) with ViT-B  
$ python pre_grey_rgb2D_Huge.py    --img_path  data/test_data            --gt_path data/test_data     --task_name 22_Heart       --checkpoint sam_vit_b_01ec64.pth              #for preparing testing data (embeddings) with ViT-B

3. Finetune SAM with your own data.

$ python train_only_box.py    --tr_npz_path data/precompute_vit_b/train  --val_npz_path data/precompute_vit_b/valid --model_type vit_b # finetune ViT-B
                              --tr_npz_path data/precompute_vit_h/train  --val_npz_path data/precompute_vit_h/valid --model_type vit_h # finetune ViT-H

4. Test on finetuned models and output the Dice results.

$ python test_only_box.py    

5. Calculation of all the indicators (Dice, IOU, HD, etc.).

$ python cal_matric.py       

Our pretrained weights

Checkpoints download path: https://drive.google.com/drive/folders/1jry-07RxGYQnT9cQE8weuCZurDCu58pj?usp=sharing

COSMOS 1050K Dataset

We collected and sorted 53 public datasets to build the large COSMOS 1050K medical image segmentation dataset. Following are the links to the datasets used in our paper.

Ownership and license of the datasets belong to their corresponding original papers, authors, or competition organizers. If you use the datasets, please cite the corresponding paper or links.

Acknowledgments

Our code is based on Segment Anything and MedSAM. We appreciate the authors for their great works. We also sincerely appreciate all the challenge organizers and owners for providing the public medical image segmentation datasets.

Citation

If you find the code useful for your research, please cite our paper.

@article{huang2024segment,
  title={Segment anything model for medical images?},
  author={Huang, Yuhao and Yang, Xin and Liu, Lian and Zhou, Han and Chang, Ao and Zhou, Xinrui and Chen, Rusi and Yu, Junxuan and Chen, Jiongquan and Chen, Chaoyu and others},
  journal={Medical Image Analysis},
  volume={92},
  pages={103061},
  year={2024},
  publisher={Elsevier}
}