Enhancing the Reliability of Segment Anything Model for Auto-Prompting Medical Image Segmentation with Uncertainty Rectification

TL;DR

UR-SAM is proposed, an uncertainty rectified SAM framework to enhance the reliability for auto-prompting medical image segmentation and incorporates a prompt augmentation module to obtain a series of input prompts for SAM for uncertainty estimation and an uncertainty-based rectification module to improve the segmentation performance.

Abstract

The Segment Anything Model (SAM) has recently emerged as a groundbreaking foundation model for prompt-driven image segmentation tasks. However, both the original SAM and its medical variants require slice-by-slice manual prompting of target structures, which directly increase the burden for applications. Despite attempts of auto-prompting to turn SAM into a fully automatic manner, it still exhibits subpar performance and lacks of reliability especially in the field of medical imaging. In this paper, we propose UR-SAM, an uncertainty rectified SAM framework to enhance the reliability for auto-prompting medical image segmentation. Building upon a localization framework for automatic prompt generation, our method incorporates a prompt augmentation module to obtain a series of input prompts for SAM for uncertainty estimation and an uncertainty-based rectification module to further utilize the distribution of estimated uncertainty to improve the segmentation performance. Extensive experiments on two public 3D medical datasets covering the segmentation of 35 organs demonstrate that without supplementary training or fine-tuning, our method further improves the segmentation performance with up to 10.7 % and 13.8 % in dice similarity coefficient, demonstrating efficiency and broad capabilities for medical image segmentation without manual prompting.

Figures (5)

• Figure 1: Overview of our proposed Uncertainty Rectified Segment Anything Model (UR-SAM) framework. We utilize an landmark localization model for auto-prompting with augmentation to generate a series of different bounding box prompts $B$ for each image $x$. we can obtain the distributions of predictions for uncertainty estimation. The SAM architecture in red-line utilizes the original image and generated prompts to obtain a set of predictions for uncertainty estimation. Then the estimated uncertainty can be utilized for rectification to obtain the final segmentation results.
• Figure 2: Detailed pipeline of prompt generation in our framework. For each target organ, six extreme points in three volumetric directions is localized to generate bounding box prompts. Furthermore, the initial generated bounding box prompt is augmented by random shifting within pre-defined ratios.
• Figure 3: Performance of original and uncertainty rectified segmentation results based on SAM SAM and MedSAM MedSAM for 3D head-and-neck organ segmentation in the StructSeg dataset.
• Figure 4: Performance of original and uncertainty rectified segmentation results based on SAM SAM and MedSAM MedSAM for 3D abdominal organ segmentation in the FLARE 22 dataset.
• Figure 5: Visual comparison of segmentation results before and after uncertainty rectification on StructSeg and FLARE 22 datasets.