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

Yichi Zhang, Shiyao Hu+4 more|arXiv (Cornell University)|Nov 17, 2023

Brain Tumor Detection and Classification8 citations

TL;DR

UR-SAM enhances SAM with prompt augmentation and an uncertainty rectification module to enable reliable auto-prompting medical image segmentation without extra training. It achieves notable Dice improvements on 3D HaN and abdominal datasets.

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.

Motivation & Objective

• Motivate reliable automatic medical image segmentation with prompt-driven models.
• Leverage uncertainty estimation to identify and rectify segmentation regions with high ambiguity.
• Develop a prompt-augmentation and class-specific confidence mechanism to improve SAM performance on medical data.
• Validate the framework on 3D medical datasets covering head-and-neck and abdominal organs.
• Demonstrate superiority over baseline auto-prompting and approach parity with manual prompting without fine-tuning.

Proposed method

• Use a landmark localization module to auto-generate 3D bounding box prompts for SAM.
• Augment initial prompts with perturbations to create multiple prompts and generate a distribution of segmentation outputs.
• Estimate pixel-level uncertainty from multi-prompt outputs using predictive entropy.
• Apply a class-specific confidence-based threshold to select high-uncertainty regions for rectification.
• Rectify segmentation by leveraging intensity relationships within the image to reassign uncertain regions.
• Evaluate UR-SAM with SAM and MedSAM backbones on StructSeg (HaN organs) and FLARE 22 (abdominal organs).

Experimental results

Research questions

• RQ1 Can prompt augmentation and uncertainty rectification improve SAM-based auto-prompting for medical image segmentation without additional training?
• RQ2 How does class-specific uncertainty-based rectification compare to fixed-threshold correction strategies on medical datasets?
• RQ3 What is the impact of UR-SAM on segmentation accuracy across small and large/complex organs in 3D CT data?
• RQ4 Does UR-SAM bring auto-prompting performance closer to manual prompting?
• RQ5 What are the effects of using SAM versus MedSAM as the base model in UR-SAM?

Key findings

• UR-SAM improves segmentation performance over baseline auto-prompting on 3D medical datasets, achieving up to 10.7% Dice gains in head-and-neck segmentation and up to 13.8% in abdominal segmentation.
• Prompt augmentation with perturbations yields robustness gains, though benefits plateau beyond a small number of augmentations and excessive perturbation can hurt performance.
• Class-specific uncertainty-based filtering consistently improves DSC over fixed-threshold approaches and ensemble methods across datasets.
• UR-SAM with SAM and MedSAM backbones outperforms auto-prompting and can approach manual prompting performance in several targets.
• Rectified segmentation results show reduced errors compared with ensembled results, illustrating the effectiveness of uncertainty-based rectification.
• MedSAM may underperform compared to SAM on unseen datasets, reinforcing the value of the UR-SAM uncertainty rectification strategy.