Uncertainty-aware Fine-tuning of Segmentation Foundation Models (SUM)

Official implementation of Uncertainty-aware Fine-tuning of Segmentation Foundation Models (NeurIPS 2024).

Kangning Liu^1,2, Brian Price², Jason Kuen², Yifei Fan², Zijun Wei², Luis Figueroa², Krzysztof J. Geras¹, Carlos Fernandez-Granda¹

¹ New York University
² Adobe

NeurIPS 2024 Poster

Project Website

Table of Contents

• Status Update
  • Current Progress
  • Next Steps
  • Known Issues
• Prerequisites
• Dataset
• Notebook
• Contact

Status Update

Current Progress

• [NEW] SUM (HQ-SAM arch.) Provide the training code and inference code of SUM implemented with HQ-SAM architecture sum_on_hq-sam

• Main experiments

  • Provided the model building code build_sam.py
  • Provided the key components of uncertainty-aware fine-tuning for the main experiment:
    • Uncertainty-aware loss losses.py
    • Uncertainty-aware prompt sampling interactive_sampling.py

Next Steps

• Main experiments
  • Provide demo Jupyter notebooks
  • Add support for the evaluation dataloader
  • Release model weights trained on the public dataset
  • Provide the full training code

Known Issues

• Some scripts may require additional dependencies not listed in the prerequisites.
• Documentation is still in progress and may lack detailed instructions for some scripts.

Prerequisites

The code requires python>=3.8, as well as pytorch>=1.7 and torchvision>=0.8. Please follow the instructions here to install both PyTorch and TorchVision dependencies. Installing both PyTorch and TorchVision with CUDA support is strongly recommended.

Dataset

TODO

Notebook

TODO

Contact

For any questions or issues, please contact:

• Kangning Liu - kangning.liu@nyu.edu