Retrieval-augmented Few-shot Medical Image Segmentation with Foundation Models

TL;DR

This work tackles the challenge of few-shot medical image segmentation across diverse modalities by introducing a retrieval-augmented framework that leverages DINOv2 for retrieving similar annotated samples and SAM 2 with a memory-attention mechanism to condition segmentation on retrieved memories, all without retraining on the target domain. The method constructs a memory bank from retrieved images and masks and uses memory attention to enrich input features for accurate segmentation. Evaluations on three medical datasets (ACDC, CMR T1-Map, Fluoroscopy) show state-of-the-art performance under minimal annotations, outperforming SAM 2 prompts and fully supervised baselines in limited-data settings and surpassing prior few-shot methods on Abd-MRI. The approach demonstrates strong cross-modality generalization and has practical potential as an annotation tool in clinical workflows, with extensions to 3D segmentation and real-time considerations discussed as future directions.

Abstract

Medical image segmentation is crucial for clinical decision-making, but the scarcity of annotated data presents significant challenges. Few-shot segmentation (FSS) methods show promise but often require training on the target domain and struggle to generalize across different modalities. Similarly, adapting foundation models like the Segment Anything Model (SAM) for medical imaging has limitations, including the need for finetuning and domain-specific adaptation. To address these issues, we propose a novel method that adapts DINOv2 and Segment Anything Model 2 (SAM 2) for retrieval-augmented few-shot medical image segmentation. Our approach uses DINOv2's feature as query to retrieve similar samples from limited annotated data, which are then encoded as memories and stored in memory bank. With the memory attention mechanism of SAM 2, the model leverages these memories as conditions to generate accurate segmentation of the target image. We evaluated our framework on three medical image segmentation tasks, demonstrating superior performance and generalizability across various modalities without the need for any retraining or finetuning. Overall, this method offers a practical and effective solution for few-shot medical image segmentation and holds significant potential as a valuable annotation tool in clinical applications.

Figures (5)

• Figure 1: (a) Illustration of how human annotators learn to segment medical images: By studying a few annotated examples, the human annotator can effectively apply the learned knowledge to segment a new, unseen case.(b) Representation of our proposed model's process: Retrieving the contextual and anatomical information from similar annotated example to guide the foundation models to perform the segmentation for new case without any retraining or finetuning.
• Figure 2: Overview of Retrieval-Augmented Few-Shot Medical Image Segmentation Workflow: (a) The main segmentation pipeline starts with input images processed through DINOv2 for dino embedding, followed by querying similar images and corresponding masks which are encoded and stored in a memory bank. The memory attention mechanism integrates the information from memory bank to assist the mask decoder in generating the final segmentation mask. (b) The process of indexing limited annotated data using DINOv2 and Faiss, enabling efficient retrieval of relevant images to enhance segmentation accuracy.
• Figure 3: Qualitative comparison of segmentation performance on ACDC, CMR T1-MAP, and Fluoroscopy Image datasets from a randomly selected sample for each dataset, respectively.
• Figure 4: Segmentation performance on the ACDC dataset with different numbers of queried images, measured by Dice similarity coefficient (DSC) for the right ventricle (RV), myocardium (Myo), and left ventricle (LV).
• Figure 5: Segmentation results on challenging basal and apical slices from the ACDC dataset, with corresponding top three similar queried images retrieved using the DINOv2.