RobustSAM: Segment Anything Robustly on Degraded Images (CVPR 2024 Highlight)

Official repository for RobustSAM: Segment Anything Robustly on Degraded Images

HF Model Cards

Updates

Aug 2024: ✨ You can refer to the Hugging Face model cards & demo built by @jadechoghari for easier usage via this link.
July 2024: ✨ Training code, data and model checkpoints for different ViT backbones are released!
June 2024: ✨ Inference code has been released!
Feb 2024: ✨ RobustSAM was accepted into CVPR 2024!

Introduction

Segment Anything Model (SAM) has emerged as a transformative approach in image segmentation, acclaimed for its robust zero-shot segmentation capabilities and flexible prompting system. Nonetheless, its performance is challenged by images with degraded quality. Addressing this limitation, we propose the Robust Segment Anything Model (RobustSAM), which enhances SAM's performance on low-quality images while preserving its promptability and zero-shot generalization.

Our method leverages the pre-trained SAM model with only marginal parameter increments and computational requirements. The additional parameters of RobustSAM can be optimized within 30 hours on eight GPUs, demonstrating its feasibility and practicality for typical research laboratories. We also introduce the Robust-Seg dataset, a collection of 688K image-mask pairs with different degradations designed to train and evaluate our model optimally. Extensive experiments across various segmentation tasks and datasets confirm RobustSAM's superior performance, especially under zero-shot conditions, underscoring its potential for extensive real-world application. Additionally, our method has been shown to effectively improve the performance of SAM-based downstream tasks such as single image dehazing and deblurring.

Setup

Create a conda environment and activate it.

conda create --name robustsam python=3.10 -y
conda activate robustsam

Clone and enter into repo directory.

git clone https://github.com/robustsam/RobustSAM
cd RobustSAM

Use command below to check your CUDA version.

nvidia-smi

Replace the CUDA version with yours in command below.

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu[$YOUR_CUDA_VERSION]
# For example: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117 # cu117 = CUDA_version_11.7

Install remaining dependencies

pip install -r requirements.txt

Download pretrained RobustSAM checkpoints of different sizes and place them into current directory.

Data Preparation

Change current directory to "data" directory.

cd data

Download train, val, test and extra COCO & LVIS dataset. (NOTE: Images in train, val and test dataset consist images from LVIS, MSRA10K, ThinObject-5k, NDD20, STREETS and FSS-1000)

bash download.sh

There are only clear images downloaded in the previous step. Use command below to generate corresponding degraded images.

bash gen_data.sh

Training

If you want to train from scratch, use command below.

python -m torch.distributed.launch train_ddp.py  --multiprocessing-distributed --exp_name [$YOUR_EXP_NAME] --model_size [$MODEL_SIZE]
# Example usage: python -m torch.distributed.launch train_ddp.py  --multiprocessing-distributed --exp_name test --model_size l

If you want to train from pretrained checkpoint, use command below.

python -m torch.distributed.launch train_ddp.py  --multiprocessing-distributed --exp_name [$YOUR_EXP_NAME] --model_size [$MODEL_SIZE] --load_model [$CHECKPOINT_PATH]
# Example usage: python -m torch.distributed.launch train_ddp.py  --multiprocessing-distributed --exp_name test --model_size l --load_model robustsam_checkpoint_l.pth

Local Gradio App

python gradio_app.py

Demo

We have prepared some images im demo_images folder for demo purpose. Besides, two prompting modes are available (box prompts and point prompts).

For box prompt:

python eval.py --bbox --model_size l

For point prompt:

python eval.py --model_size l

In default, demo results will be saved to demo_result/[$PROMPT_TYPE].

Comparison of computational requirements

Visual Comparison

Quantitative Comparison

Seen dataset with synthetic degradation

Unseen dataset with synthetic degradation

Unseen dataset with real degradation

Reference

If you find this work useful, please consider citing us!

@inproceedings{chen2024robustsam,
  title={RobustSAM: Segment Anything Robustly on Degraded Images},
  author={Chen, Wei-Ting and Vong, Yu-Jiet and Kuo, Sy-Yen and Ma, Sizhou and Wang, Jian},
  journal={CVPR},
  year={2024}
}

Acknowledgements

We thank the authors of SAM from which our repo is based off of.