youhuang67/cfr-icl-interactive-segmentation

Cascade-Forward Refinement with Iterative Click Loss for Interactive Image Segmentation

CFR-ICL: Cascade-Forward Refinement with Iterative Click Loss for Interactive Image Segmentation

Environment

Training and evaluation environment: Python 3.9, PyTorch 1.13.1, CUDA 11.0. Run the following command to install required packages.

pip3 install -r requirements.txt

You need to configue the paths to the datasets in config.yml before training or testing. A script download_datasets.sh is prepared to download and organize required datasets.

Demo

An example script to run the demo.

python demo.py --checkpoint=weights/cocolvis_icl_vit_huge.pth --gpu 0

Evaluation

Before evaluation, please download the datasets and models, and then configure the path in config.yml.

Download our model, please download below 3 zipped files and extract before use:

• cocolvis_icl_vit_huge.pth.7z.001
• cocolvis_icl_vit_huge.pth.7z.002
• cocolvis_icl_vit_huge.pth.7z.003

Use the following code to evaluate the huge model.

python scripts/evaluate_model.py NoBRS \
    --gpu=0 \
    --checkpoint=cocolvis_icl_vit_huge.pth \
    --datasets=GrabCut,Berkeley,DAVIS,PascalVOC,SBD \\
    --cf-n=4 \\
    --acf

# cf-n: CFR steps
# acf: adaptive CFR

Training

Before training, please download the MAE pretrained weights (click to download: ViT-Base, ViT-Large, ViT-Huge) and configure the dowloaded path in config.yml

Please also download the pretrained SimpleClick models from here.

Use the following code to train a huge model on C+L:

python train.py models/plainvit_huge448_cocolvis.py \
    --batch-size=32 \
    --ngpus=4

Citation

@article{sun2023cfricl,
      title={CFR-ICL: Cascade-Forward Refinement with Iterative Click Loss for Interactive Image Segmentation}, 
      author={Shoukun Sun and Min Xian and Fei Xu and Tiankai Yao and Luca Capriotti},
      year={2023},
      eprint={2303.05620},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Acknowledgement

Our project is developed based on RITM and SimpleClick