FANzhaoxin666 / BAM

Official PyTorch Implementation of Learning What Not to Segment: A New Perspective on Few-Shot Segmentation (CVPR 2022).

Learning What Not to Segment: A New Perspective on Few-Shot Segmentation

This repo contains the code for our CVPR 2022 paper "Learning What Not to Segment: A New Perspective on Few-Shot Segmentation" by Chunbo Lang, Gong Cheng, Binfei Tu, and Junwei Han.

Dependencies

• Python 3.8
• PyTorch 1.7.0
• cuda 11.0
• torchvision 0.8.1
• tensorboardX 2.14

Datasets

• PASCAL-5ⁱ: VOC2012 + SBD

• COCO-20ⁱ: COCO2014

  Download the data lists (.txt files) and put them into the BAM/lists directory.

• Run util/get_mulway_base_data.py to generate base annotations for stage1, or directly use the trained weights.

Models

• Download the pre-trained backbones from here and put them into the BAM/initmodel directory.
• Download our trained base learners from OneDrive and put them under initmodel/PSPNet.
• We provide 4 trained BAM models for performance evaluation: 2 VGG16 based models for PASCAL-5⁰ and 2 ResNet50 based models for COCO-20⁰.

Usage

• Change configuration via the .yaml files in BAM/config, then run the .sh scripts for training and testing.

• Stage1 Pre-training

  Train the base learner within the standard learning paradigm.

  sh train_base.sh
  

• Stage2 Meta-training

  Train the meta learner and ensemble module within the meta-learning paradigm.

  sh train.sh
  

• Stage3 Meta-testing

  Test the proposed model under the standard few-shot setting.

  sh test.sh
  

• Stage4 Generalized testing

  Test the proposed model under the generalized few-shot setting.

  sh test_GFSS.sh
  

Performance

Performance comparison with the state-of-the-art approaches (i.e., HSNet and PFENet) in terms of average mIoU across all folds.

1. PASCAL-5ⁱ

   Backbone	Method	1-shot	5-shot
   VGG16	HSNet	59.70	64.10
   	BAM (ours)	64.41 (+4.71)	68.76 (+4.66)
   ResNet50	HSNet	64.00	69.50
   	BAM (ours)	67.81 (+3.81)	70.91 (+1.41)

2. COCO-20ⁱ

   Backbone	Method	1-shot	5-shot
   VGG16	PFENet	36.30	40.40
   	BAM (ours)	43.50 (+7.20)	49.34 (+8.94)
   ResNet50	HSNet	39.20	46.90
   	BAM (ours)	46.23 (+7.03)	51.16 (+4.26)

Visualization

References

This repo is mainly built based on PFENet, RePRI, and SemSeg. Thanks for their great work!

To-Do List

• Support different backbones
• Support various annotations for training/testing
• Multi-GPU training
• FSS-1000 dataset

Bibtex

Please consider citing our paper if the project helps your research. BibTeX reference is as follows:

@InProceedings{lang2022bam,
    title={Learning What Not to Segment: A New Perspective on Few-Shot Segmentation},
    author={Lang, Chunbo and Cheng, Gong and Tu, Binfei and Han, Junwei},
    booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year={2022}
}