zhang-haojie / wesam

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

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

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

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

🖼️ Demo

COCO Dataset

ISIC Dataset

OCID Dataset

CAMO Dataset

COCO-C Dataset

🎫 License

The content of this project itself is licensed under LICENSE.

💡 Acknowledgement

• SAM

• lightning-sam

• SAM-LoRA

🖊️ Citation

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

@article{zhang2023improving,
      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},
      journal={arXiv preprint arXiv:2312.03502},
      year={2023}
}