Sam2Rad: A Segmentation Model for Medical Images with Learnable Prompts

Assefa Seyoum Wahd,Banafshe Felfeliyan,Yuyue Zhou,Shrimanti Ghosh,Adam McArthur,Jiechen Zhang,Jacob L. Jaremko,Abhilash Hareendranathan

2024-09-11

Abstract:Foundation models like the segment anything model require high-quality manual prompts for medical image segmentation, which is time-consuming and requires expertise. SAM and its variants often fail to segment structures in ultrasound (US) images due to domain shift. We propose Sam2Rad, a prompt learning approach to adapt SAM and its variants for US bone segmentation without human prompts. It introduces a prompt predictor network (PPN) with a cross-attention module to predict prompt embeddings from image encoder features. PPN outputs bounding box and mask prompts, and 256-dimensional embeddings for regions of interest. The framework allows optional manual prompting and can be trained end-to-end using parameter-efficient fine-tuning (PEFT). Sam2Rad was tested on 3 musculoskeletal US datasets: wrist (3822 images), rotator cuff (1605 images), and hip (4849 images). It improved performance across all datasets without manual prompts, increasing Dice scores by 2-7% for hip/wrist and up to 33% for shoulder data. Sam2Rad can be trained with as few as 10 labeled images and is compatible with any SAM architecture for automatic segmentation.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

### Problems the Paper Attempts to Solve

This paper aims to address the issues faced when using foundational models (such as the Segment Anything Model, SAM) in medical image segmentation. Specifically:

1. **The need for high-quality manual prompts**: Existing medical image segmentation methods require high-quality manual prompts, which are typically generated by individuals with medical expertise, making the process time-consuming and costly.
2. **Poor cross-domain adaptability**: Even with sparse prompts (such as boxes, points, or text) or dense prompts (such as masks), SAM and its variants (like MedSAM) perform poorly in segmenting bones in ultrasound images (US), primarily due to significant domain shift issues.
3. **Dependence on manual prompts**: Current methods still rely on manual prompts for image segmentation, which is often challenging in practical applications, especially in the absence of medical expertise.

To address these issues, the authors propose a new prompt learning method—Sam2Rad, which can automatically generate prompts, thereby completing medical image segmentation tasks without human intervention. Specifically, Sam2Rad introduces a Prompt Predictor Network (PPN) to predict prompt embeddings, generating bounding boxes and mask prompts directly from features extracted by the image encoder. Additionally, this framework supports optional manual prompts, which can be combined with learned prompts and input into the mask decoder. Through Parameter-Efficient Fine-Tuning (PEFT), the PPN and mask decoder can be trained end-to-end. Experimental results show that Sam2Rad significantly improves segmentation performance on multiple ultrasound datasets, particularly on shoulder ultrasound images, where the Dice score increased from 49% to 82%.