Home

Awesome

M-SAM [Paper]

<a src="https://img.shields.io/badge/cs.CV-2403.05912-b31b1b?logo=arxiv&logoColor=red" href="https://arxiv.org/pdf/2403.05912"> <img src="https://img.shields.io/badge/cs.CV-2403.05912-b31b1b?logo=arxiv&logoColor=red">

We introduce Mask-Enhanced SAM (M-SAM), an innovative architecture tailored for 3D tumor lesion segmentation. This method is elaborated on the paper Mask-Enhanced Segment Anything Model for Tumor Lesion Semantic Segmentation.

🌟 Highlights

🫱 A Quick Overview

M-SAM consists of multiple stages of iterative refinement, which makes it possible to refine the predicted segmentation masks iteratively, thus obtaining more accurate segmentation boundaries progressively.

<p align="center"><img width="800" alt="image" src="https://github.com/nanase1025/M-SAM/blob/main/assets/architecture.jpg"></p> Our MEA is proposed to aggregate the image embedding with corresponding mask, so that the updated image embedding can perceive position priors of the lesion regions. <p align="center"><img width="800" alt="image" src="https://github.com/nanase1025/M-SAM/blob/main/assets/MEA.jpg"></p>

🫱 Requirement

Install the environment:

pip install -r requirements.txt

Then download SAM checkpoint, and put it at .work_dir/SAM/

🫱 Usage

  1. Prepare your own train dataset and refer to the samples in M-SAM/demo to replace them according to your specific scenario. The target file structures should be like the following:
data/train
      ├── ct_BraTS
      │ ├── imagesTr
      │ │ ├── word_0025.nii.gz
      │ │ ├── ...
      │ ├── labelsTr
      │ │ ├── word_0025.nii.gz
      │ │ ├── ...
      │ ── ct_LiTS
      │ ├── imagesTr
      │ │ ├── word_0025.nii.gz
      │ │ ├── ...
      │ ├── labelsTr
      │ │ ├── word_0025.nii.gz
      │ │ ├── ...
      ├── ...
  1. Start Training!
cd ./M-SAM
python train_msam.py

This will start the training process of the M-SAM model on your prepared data. The key options are listed below:

💡Note: Use the --checkpoint to set the pre-trained weight path, if no checkpoint is found in the specified path, the model will undergo training from scratch.

  1. Evaluation Prepare your own test dataset and refer to the samples in M-SAM/demo to replace them according to your specific scenario. The target file structures should be like the following:
data/test
      ├── adrenal
      │ ├── ct_WORD
      │ │ ├── imagesTr
      │ │ │ ├── word_0027.nii.gz
      │ │ │ ├── ...
      │ │ ├── labelsTr
      │ │ │ ├── word_0027.nii.gz
      │ │ │ ├── ...
      ├── liver
      │ ├── ct_WORD
      │ │ ├── imagesTr
      │ │ │ ├── word_0029.nii.gz
      │ │ │ ├── ...
      │ │ ├── labelsTr
      │ │ │ ├── word_0029.nii.gz
      │ │ │ ├── ...
      ├── ...

And

cd ./M-SAM
python test_msam.py

This will start the testing process of the M-SAM model on your prepared data.

📧 Citation

@inproceedings{shi2024mask,
  title={Mask-enhanced segment anything model for tumor lesion semantic segmentation},
  author={Shi, Hairong and Han, Songhao and Huang, Shaofei and Liao, Yue and Li, Guanbin and Kong, Xiangxing and Zhu, Hua and Wang, Xiaomu and Liu, Si},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={403--413},
  year={2024},
  organization={Springer}
}

🤝 Acknowledgement

📚 License

This project is released under the Apache 2.0 license.