SAM Fewshot Finetuning for Anatomical Segmentation in Medical Images

TL;DR

The paper tackles the challenge of adapting a large segmentation foundation model, SAM, to anatomical segmentation in medical images with limited labeled data. It proposes a few-shot fine-tuning strategy that freezes the image encoder and retrains only the mask decoder, using few-shot embeddings derived from labeled images as prompts, thereby replacing traditional user prompts. Across four datasets and six anatomical tasks, the method outperforms SAM with point prompts and achieves performance comparable to, or approaching, fully supervised nnU-Net while reducing labeled data by at least an order of magnitude. The approach also offers practical efficiency via embedding caching and demonstrates generality to 2D/3D segmentation beyond medical imaging, suggesting a versatile framework for token-query-based segmentation.

Abstract

We propose a straightforward yet highly effective few-shot fine-tuning strategy for adapting the Segment Anything (SAM) to anatomical segmentation tasks in medical images. Our novel approach revolves around reformulating the mask decoder within SAM, leveraging few-shot embeddings derived from a limited set of labeled images (few-shot collection) as prompts for querying anatomical objects captured in image embeddings. This innovative reformulation greatly reduces the need for time-consuming online user interactions for labeling volumetric images, such as exhaustively marking points and bounding boxes to provide prompts slice by slice. With our method, users can manually segment a few 2D slices offline, and the embeddings of these annotated image regions serve as effective prompts for online segmentation tasks. Our method prioritizes the efficiency of the fine-tuning process by exclusively training the mask decoder through caching mechanisms while keeping the image encoder frozen. Importantly, this approach is not limited to volumetric medical images, but can generically be applied to any 2D/3D segmentation task. To thoroughly evaluate our method, we conducted extensive validation on four datasets, covering six anatomical segmentation tasks across two modalities. Furthermore, we conducted a comparative analysis of different prompting options within SAM and the fully-supervised nnU-Net. The results demonstrate the superior performance of our method compared to SAM employing only point prompts (approximately 50% improvement in IoU) and performs on-par with fully supervised methods whilst reducing the requirement of labeled data by at least an order of magnitude.

Figures (3)

• Figure 1: Two cases using SAM point prompting. The points are marked in yellow. Segmentation is colored red. The top row is an axial slice from a chest CT scan with a segmented aorta, and the bottom row is a coronal slice of a CT image with a segmented femur. We show three segmentation predictions from SAM for each case, with predicted IoU and stability score above the default thresholds (0.88 and 0.95, respectively).
• Figure 2: Mask decoder (b) in SAM using user prompts (d) and the proposed prompting method based on fewshot target embeddings derived from a set of labeled images (c). Two-way transformer layers allow both image embeddings and prompt embeddings to attend to each other's information (a). To maintain clarity, we have omitted the process of dense mask embeddings and positional encodings. Green dash lines indicate the process in the mask decoder that are modified for fewshot finetuning. And the part in the box marked by the red dash line is removed in the modified mask decoder as the proposed method does not use user prompt embeddings.
• Figure 3: Qualitative results of the segmentation methods in comparison. The first column represents the ground truth segmentations, where all anatomies are visualized in red, except for the femur, which is visualized in green to differentiate it from the tibia. The subsequent columns depict the results of different segmentation methods: SAM with bounding box prompt (5th column), the fully supervised nnU-Net (6th column), and the proposed fewshot finetuning with 5, 20, and 50 labeled images (2nd-4th columns). Each row corresponds to a different case.