/wesam

[CVPR 2024] Code for "Improving the Generalization of Segmentation Foundation Model under Distribution Shift via Weakly Supervised Adaptation"

Primary LanguagePythonMIT LicenseMIT

Improving the Generalization of Segmentation Foundation Model under Distribution Shift via Weakly Supervised Adaptation

🎈 News

  • [2024.2.27] Our work has been accepted to CVPR 2024 🎉
  • [2024.3.1] Training and inference code released

🚀 Introduction

image

Segment Anything Model was pre-trained on a large-scale dataset but exhibits awkward performance on diverse downstream segmentation tasks. We adapt SAM through weak supervision to enhance its generalization capabilities.

📻 Overview

image

The proposed self-training architecture with anchor network regularization and contrastive loss regularization. Red arrows indicates the backpropagation flow.

📆 TODO

  • Release code

🎮 Getting Started

1. Install Environment

see INSTALL.

2. Prepare Dataset and Checkpoints

see PREPARE.

3. Adapt with Weak Supervision

# 1 modify configs/config.py 
# Prompt type: box, point, coarse

# 2 adapt
python adaptation.py

4. Validation

python validate.py --ckpt /path/to/checkpoint

🖼️ Visualization

image

🎫 License

The content of this project itself is licensed under LICENSE.

💡 Acknowledgement

🖊️ Citation

If you find this project useful in your research, please consider cite:

@inproceedings{zhang2024improving,
  title={Improving the generalization of segmentation foundation model under distribution shift via weakly supervised adaptation},
  author={Zhang, Haojie and Su, Yongyi and Xu, Xun and Jia, Kui},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={23385--23395},
  year={2024}
}