Mining Latent Classes for Few-shot Segmentation

Lihe Yang, Wei Zhuo, Lei Qi, Yinghuan Shi, Yang Gao.

This codebase contains baseline of our paper Mining Latent Classes for Few-shot Segmentation, ICCV 2021 Oral.

Several key modifications to the simple yet effective metric learning framework: - Remove the final residual stage in ResNet for stronger generalization - Remove the final ReLU for feature matching - Freeze all the BatchNorms from ImageNet pretrained model

Environment

This codebase was tested with the following environment configurations.

• Ubuntu 18.04
• CUDA 11.2
• Python 3.7.4
• PyTorch 1.6.0
• Pillow, numpy, torchvision, tqdm
• Two NVIDIA V100 GPUs

Getting Started

Data Preparation

Pretrained model: ResNet-50 | ResNet-101

Dataset: Pascal JPEGImages | SegmentationClass | ImageSets

File Organization

├── ./pretrained
    ├── resnet50.pth
    └── resnet101.pth

├── [Your Pascal Path]
    ├── JPEGImages
    │   ├── 2007_000032.jpg
    │   └── ...
    │
    ├── SegmentationClass
    │   ├── 2007_000032.png
    │   └── ...
    │
    └── ImageSets
        ├── train.txt
        └── val.txt

Run the Code

CUDA_VISIBLE_DEVICES=0,1 python -W ignore main.py \
  --dataset pascal --data-root [Your Pascal Path] \
  --backbone resnet50 --fold 0 --shot 1

You may change the backbone from resnet50 to resnet101, change the fold from 0 to 1/2/3, or change the shot from 1 to 5 for other settings.

Performance and Trained Models

Here we report the performance of our modified baseline on Pascal. You can click on the numbers to download corresponding trained models.

The training time is measured on two V100 GPUs. Compared with other works, our method is efficient to train.

Acknowledgement

We thank PANet, PPNet, PFENet and other FSS works for their great contributions.

Citation

If you find this project useful for your research, please consider citing:

@inproceedings{yang2021mining,
  title={Mining Latent Classes for Few-shot Segmentation},
  author={Yang, Lihe and Zhuo, Wei and Qi, Lei and Shi, Yinghuan and Gao, Yang},
  booktitle={ICCV},
  year={2021}
}