[Paper Review] ASPS: Augmented Segment Anything Model for Polyp Segmentation

Huiqian Li, Dingwen Zhang+4 more|arXiv (Cornell University)|Jun 30, 2024

Vehicle License Plate Recognition4 citations

TL;DR

This paper proposes ASPS, a novel adaptation of the Segment Anything Model (SAM) for polyp segmentation that overcomes domain shift and poor out-of-distribution generalization. By integrating a trainable CNN encoder with SAM’s frozen ViT via Cross-branch Feature Augmentation (CFA) and applying Uncertainty-guided Prediction Regularization (UPR) to improve confidence calibration, ASPS achieves state-of-the-art performance across five polyp datasets without requiring prompts.

ABSTRACT

Polyp segmentation plays a pivotal role in colorectal cancer diagnosis. Recently, the emergence of the Segment Anything Model (SAM) has introduced unprecedented potential for polyp segmentation, leveraging its powerful pre-training capability on large-scale datasets. However, due to the domain gap between natural and endoscopy images, SAM encounters two limitations in achieving effective performance in polyp segmentation. Firstly, its Transformer-based structure prioritizes global and low-frequency information, potentially overlooking local details, and introducing bias into the learned features. Secondly, when applied to endoscopy images, its poor out-of-distribution (OOD) performance results in substandard predictions and biased confidence output. To tackle these challenges, we introduce a novel approach named Augmented SAM for Polyp Segmentation (ASPS), equipped with two modules: Cross-branch Feature Augmentation (CFA) and Uncertainty-guided Prediction Regularization (UPR). CFA integrates a trainable CNN encoder branch with a frozen ViT encoder, enabling the integration of domain-specific knowledge while enhancing local features and high-frequency details. Moreover, UPR ingeniously leverages SAM's IoU score to mitigate uncertainty during the training procedure, thereby improving OOD performance and domain generalization. Extensive experimental results demonstrate the effectiveness and utility of the proposed method in improving SAM's performance in polyp segmentation. Our code is available at https://github.com/HuiqianLi/ASPS.

Motivation & Objective

• To address the domain gap between natural images in SAM’s pre-training and endoscopy polyp images, which limits performance.
• To improve SAM’s ability to capture local and high-frequency features in polyp images, which are often overlooked due to its focus on global, low-frequency patterns.
• To reduce prediction uncertainty and improve confidence calibration for out-of-distribution (OOD) endoscopy images.
• To eliminate dependency on prompts in clinical applications by enabling end-to-end, prompt-free training.
• To enhance domain generalization and robustness through a lightweight, efficient adaptation framework.

Proposed method

• Introduces Cross-branch Feature Augmentation (CFA), which fuses features from a trainable CNN encoder with the frozen Vision Transformer (ViT) encoder to enhance multi-scale and high-frequency feature learning.
• Employs Uncertainty-guided Prediction Regularization (UPR), which uses ground truth as a hint to calibrate confidence scores and reduce uncertainty during training.
• Adapts normalization layers in the neck and transformer blocks to improve adaptation to endoscopy domain distributions.
• Uses Fourier analysis to validate that the CNN branch captures more high-frequency content than the ViT baseline.
• Trains the model end-to-end without prompts, enabling direct inference on endoscopy images.
• Leverages SAM’s IoU prediction as a supervisory signal in UPR to guide uncertainty reduction and improve OOD generalization.

Experimental results

Research questions

• RQ1 Can a hybrid CNN-ViT encoder architecture improve feature representation for polyp segmentation compared to a pure ViT-based SAM?
• RQ2 Does introducing uncertainty regularization during training enhance out-of-distribution generalization and confidence calibration in polyp segmentation?
• RQ3 Can domain-specific feature enhancement be achieved without fine-tuning the original SAM weights or relying on prompts?
• RQ4 How does the integration of multi-level features and position embedding replacement affect segmentation accuracy and robustness?
• RQ5 To what extent does the proposed method outperform existing SAM-based models on diverse polyp datasets?

Key findings

• ASPS achieves a mean Dice score of 0.914 and mean IoU of 0.843 across five polyp datasets, outperforming baseline SAM and prior methods like MedSAM and SurgicalSAM.
• The CFA module alone improves mean Dice by 31.7% and mean IoU by 41.4% compared to the ViT-B baseline.
• The combination of CFA and UPR achieves the best performance, with a Dice score of 0.950 and IoU of 0.905 on the CVC-ClinicDB dataset.
• UPR with neck layer normalization (NN) and hint-based regularization yields the highest performance, demonstrating its effectiveness in uncertainty reduction.
• Fourier analysis confirms that the CNN branch captures significantly more high-frequency content than the ViT encoder, supporting improved local detail learning.
• Qualitative results show that ASPS predictions are closer to ground truth, especially in capturing fine boundaries and texture details of polyps.