CELP-Pytorch

This is a PyTorch implementation of CELP for few-shot segmentation

• The pipeline of CELP

• Display of predictions

Usage

Environment

python==3.6.9
GCC=7.5
torch=1.7.1+cu110
torchvision=0.8.2+cu110
pycocotoolssion=2.0.3
cython
tensorboardX
tqdm
PyYAML
opencv-python

Data Preparation

• Pascal-5i: download PASCAL VOC 2012 and SBD dataset. For SBD dataset, the val images should be excluded from the list of training samples.
• COCO-20i: download COCO2014 dataset, including 2014 train images, 2014 val images, and 2014 Train/Val annotations. Extract the files in the same directory, and run

python prepare_coco_data.py

Train

• Create a fold of 'backbones' at the root directory. Download the ImageNet pretrained backbone, and put them into the 'backbones' directory.
• For the implementation of CELP with baseline CyCTR, the Deformable DETR dependencies should be built first.

cd CELP_CyCTR/model/ops/
bash make.sh

• Execute the command at the root directory:

sh train_*.sh {*dataset*} {*model_config*}

For example

sh train_cyctr.sh pascal split0_resnet50

sh train_pfenet.sh pascal split0_resnet50

Run Demo

We provide pre-trained models on PASCAL-5^i and COCO-20^i for testing. Update the config file by specifying the target split and path of weights. Then execute the command:

sh test_*.sh {*dataset*} {*model_config*}

For example

 sh test_cyctr.sh pascal split0_resnet50

Acknowledgement

The project is built upon PFENet, CyCTR, and Deformable-DETR.

• PFENet: https://github.com/dvlab-research/PFENet
• CyCTR: https://github.com/YanFangCS/CyCTR-Pytorch
• Deformable-DETR: https://github.com/fundamentalvision/Deformable-DETR

Thank for their excellent works.

Citation

If you find our codes or models useful, please consider to give us a star or cite with:

@misc{zhao2022contrastive,
      title={Contrastive Enhancement Using Latent Prototype for Few-Shot Segmentation}, 
      author={Xiaoyu Zhao and Xiaoqian Chen and Zhiqiang Gong and Wen Yao and Yunyang Zhang and Xiaohu Zheng},
      year={2022},
      eprint={2203.04095},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}