This repo contains the code for our CVPR 2022 Oral paper "Learning What Not to Segment: A New Perspective on Few-Shot Segmentation" by Chunbo Lang, Gong Cheng, Binfei Tu, and Junwei Han.
Abstract: Recently few-shot segmentation (FSS) has been extensively developed. Most previous works strive to achieve generalization through the meta-learning framework derived from classification tasks; however, the trained models are biased towards the seen classes instead of being ideally class-agnostic, thus hindering the recognition of new concepts. This paper proposes a fresh and straightforward insight to alleviate the problem. Specifically, we apply an additional branch (base learner) to the conventional FSS model (meta learner) to explicitly identify the targets of base classes, i.e., the regions that do not need to be segmented. Then, the coarse results output by these two learners in parallel are adaptively integrated to yield precise segmentation prediction. Considering the sensitivity of meta learner, we further introduce an adjustment factor to estimate the scene differences between the input image pairs for facilitating the model ensemble forecasting. The substantial performance gains on PASCAL-5i and COCO-20i verify the effectiveness, and surprisingly, our versatile scheme sets a new state-of-the-art even with two plain learners. Moreover, in light of the unique nature of the proposed approach, we also extend it to a more realistic but challenging setting, i.e., generalized FSS, where the pixels of both base and novel classes are required to be determined.
[Mar 2, 2022]
- BAM is accepted to CVPR2022.
[Mar 29, 2022]
- Our paper is selected for an oral presentation.
- Python 3.8
- PyTorch 1.7.0
- cuda 11.0
- torchvision 0.8.1
- tensorboardX 2.14
-
COCO-20i: 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.
- 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-50 and 2 ResNet50 based models for COCO-200.
-
Change configuration via the
.yaml
files inBAM/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 comparison with the state-of-the-art approaches (i.e., HSNet and PFENet) in terms of average mIoU across all folds.
-
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) -
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)
- Support different backbones
- Support various annotations for training/testing
- Multi-GPU training
- FSS-1000 dataset
This repo is mainly built based on PFENet, RePRI, and SemSeg. Thanks for their great work!
If you find our work and this repository useful. Please consider giving a star ⭐ and citation 📚.
@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}
}