Awesome

Rethinking Interactive Image Segmentation with Low Latency, High Quality, and Diverse Prompts

Pytorch implementation for paper Rethinking Interactive Image Segmentation with Low Latency, High Quality, and Diverse Prompts, CVPR 2024. <br> Qin Liu, Jaemin Cho, Mohit Bansal, Marc Niethammer <br> UNC-Chapel Hill

Paper | Project | Demos

Installation

The code is tested with python=3.10, torch=2.2.0, torchvision=0.17.0.

git clone https://github.com/uncbiag/SegNext
cd SegNext

Now, create a new conda environment and install required packages accordingly.

conda create -n segnext python=3.10
conda activate segnext
conda install pytorch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install -r requirements.txt

Getting Started

First, download three model weights: vitb_sax1 (408M), vitb_sax2 (435M), and vitb_sax2_ft (435M). These weights will be automatically saved to the weights folder.

python download.py

Run interactive GUI with the downloaded weights. The assets contains images for demo.

./run_demo.sh

Datasets

We train and test our method on three datasets: DAVIS, COCO+LVIS, and HQSeg-44K.

Don't forget to change the paths to the datasets in config.yml after downloading and unpacking.

(*) To prepare COCO+LVIS, you need to download original LVIS v1.0, then download and unpack pre-processed annotations that are obtained by combining COCO and LVIS dataset into the folder with LVIS v1.0. (The combined annotations are prepared by RITM.)

Evaluation

We provide a script (run_eval.sh) to evaluate our presented models. The following command runs the NoC evaluation on all test datasets.

python ./segnext/scripts/evaluate_model.py --gpus=0 --checkpoint=./weights/vitb_sa2_cocolvis_hq44k_epoch_0.pth --datasets=DAVIS,HQSeg44K

For SAT latency evaluation, please refer to eval_sat_latency.ipynb.

Training

We provide a script (run_train.sh) for training our models on the HQSeg-44K dataset. You can start training with the following commands. By default we use 4 A6000 GPUs for training.

# train vitb-sax1 model on coco+lvis 
MODEL_CONFIG=./segnext/models/default/plainvit_base1024_cocolvis_sax1.py
torchrun --nproc-per-node=4 --master-port 29504 ./segnext/train.py ${MODEL_CONFIG} --batch-size=16 --gpus=0,1,2,3

# train vitb-sax2 model on coco+lvis 
MODEL_CONFIG=./segnext/models/default/plainvit_base1024_cocolvis_sax2.py
torchrun --nproc-per-node=4 --master-port 29505 ./segnext/train.py ${MODEL_CONFIG} --batch-size=16 --gpus=0,1,2,3

# finetune vitb-sax2 model on hqseg-44k 
MODEL_CONFIG=./segnext/models/default/plainvit_base1024_hqseg44k_sax2.py
torchrun --nproc-per-node=4 --master-port 29506 ./segnext/train.py ${MODEL_CONFIG} --batch-size=12 --gpus=0,1,2,3 --weights ./weights/vitb_sa2_cocolvis_epoch_90.pth

Citation

@article{liu2024rethinking,
  title={Rethinking Interactive Image Segmentation with Low Latency, High Quality, and Diverse Prompts},
  author={Liu, Qin and Cho, Jaemin and Bansal, Mohit and Niethammer, Marc},
  journal={arXiv preprint arXiv:2404.00741},
  year={2024}
}