Automating MedSAM by Learning Prompts with Weak Few-Shot Supervision

TL;DR

This paper introduces a method to automate image segmentation by learning prompts directly from images using weak supervision, reducing the need for manual prompt design and extensive labeled data, especially in medical imaging.

Contribution

It proposes a lightweight module to learn prompt embeddings from images, enabling automatic segmentation with minimal supervision and fine-tuning of foundation models like MedSAM for medical images.

Findings

Effective segmentation with only 10 samples and weak labels

Validated on three medical datasets in MR and ultrasound imaging

Reduces manual prompt engineering and annotation effort

Abstract

Foundation models such as the recently introduced Segment Anything Model (SAM) have achieved remarkable results in image segmentation tasks. However, these models typically require user interaction through handcrafted prompts such as bounding boxes, which limits their deployment to downstream tasks. Adapting these models to a specific task with fully labeled data also demands expensive prior user interaction to obtain ground-truth annotations. This work proposes to replace conditioning on input prompts with a lightweight module that directly learns a prompt embedding from the image embedding, both of which are subsequently used by the foundation model to output a segmentation mask. Our foundation models with learnable prompts can automatically segment any specific region by 1) modifying the input through a prompt embedding predicted by a simple module, and 2) using weak labels (tight bounding boxes) and few-shot supervision (10 samples). Our approach is validated on MedSAM, a version of SAM fine-tuned for medical images, with results on three medical datasets in MR and ultrasound imaging. Our code is available on https://github.com/Minimel/MedSAMWeakFewShotPromptAutomation.