Awesome

ASPS: Augmented Segment Anything Model for Polyp Segmentation

News

2024/6/25: 🎉Our method was accepted by MICCAI 2024.

2024/5/21: Add data loader for Skin Lesion Segmentation (ISIC2017).

Requirements

Install the dependencies of SAM.

Install mmcv-full for CNN encoder.

conda create --name ASPS python=3.8
conda activate ASPS
pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116
pip install mmcv-full==1.6.2 -f https://download.openmmlab.com/mmcv/dist/cu116/torch1.12/index.html

pip install tqdm
pip install opencv-python
pip install albumentations==1.3.0

Dataset

We conduct extensive experiments on five polyp segmentation datasets following PraNet.

For skin lesion segmentation: following EGE-UNet, needing to modify from dataset.Segmentation_other to from dataset.Segmentation_isic both in train.py and infer.py.

Training

We used train.py to train our framework.

The --exp_name is the name of the experiment, and --polyp_dir is the path to the training dataset.

python train.py --exp_name '0308_E_L' --polyp_dir "polyp_seg/TrainDataset/"

Evaluating

We used infer.py to evaluate our framework.

The --dataset_name is the name of the dataset, and --test_seg_dir is the path to the testing dataset.

python infer.py --exp_name '0308_E_L' --dataset_name 'CVC-300'  --test_seg_dir "polyp_seg/TestDataset/CVC-300/"   
python infer.py --exp_name '0308_E_L' --dataset_name 'CVC-ClinicDB'  --test_seg_dir "polyp_seg/TestDataset/CVC-ClinicDB/"   
python infer.py --exp_name '0308_E_L' --dataset_name 'CVC-ColonDB'  --test_seg_dir "polyp_seg/TestDataset/CVC-ColonDB/"   
python infer.py --exp_name '0308_E_L' --dataset_name 'ETIS-LaribPolypDB'  --test_seg_dir "polyp_seg/TestDataset/ETIS-LaribPolypDB/"   
python infer.py --exp_name '0308_E_L' --dataset_name 'Kvasir'  --test_seg_dir "polyp_seg/TestDataset/Kvasir/"  

You can directly run the train.sh to train and evaluate our framework.

Note: If using SUN_SEG dataset, the training and evaluating codes are in 'scripts/'.

Visualize and Inference

To inference single image or visualize the results, run vis.py.

raw image	pred mask	GT

Checkpoints

Citation

@article{li2024asps,
  title={ASPS: Augmented Segment Anything Model for Polyp Segmentation},
  author={Li, Huiqian and Zhang, Dingwen and Yao, Jieru and Han, Longfei and Li, Zhongyu and Han, Junwei},
  journal={arXiv preprint arXiv:2407.00718},
  year={2024}
}