CoSAM: Self-Correcting SAM for Domain Generalization in 2D Medical Image Segmentation

TL;DR

CoSAM is a novel self-correcting framework for 2D medical image segmentation that improves domain generalization by iteratively refining coarse masks without manual prompts, outperforming existing SAM-based methods.

Contribution

We introduce CoSAM, a self-correcting approach that refines segmentation masks iteratively, eliminating the need for prompt generators and enhancing generalization in medical imaging.

Findings

Outperforms state-of-the-art SAM-based methods on medical benchmarks.

Effectively refines segmentation masks through a self-correcting loop.

Demonstrates robustness across multiple scenarios and domain shifts.

Abstract

Medical images often exhibit distribution shifts due to variations in imaging protocols and scanners across different medical centers. Domain Generalization (DG) methods aim to train models on source domains that can generalize to unseen target domains. Recently, the segment anything model (SAM) has demonstrated strong generalization capabilities due to its prompt-based design, and has gained significant attention in image segmentation tasks. Existing SAM-based approaches attempt to address the need for manual prompts by introducing prompt generators that automatically generate these prompts. However, we argue that auto-generated prompts may not be sufficiently accurate under distribution shifts, potentially leading to incorrect predictions that still require manual verification and correction by clinicians. To address this challenge, we propose a method for 2D medical image segmentation called Self-Correcting SAM (CoSAM). Our approach begins by generating coarse masks using SAM in a prompt-free manner, providing prior prompts for the subsequent stages, and eliminating the need for prompt generators. To automatically refine these coarse masks, we introduce a generalized error decoder that simulates the correction process typically performed by clinicians. Furthermore, we generate diverse prompts as feedback based on the corrected masks, which are used to iteratively refine the predictions within a self-correcting loop, enhancing the generalization performance of our model. Extensive experiments on two medical image segmentation benchmarks across multiple scenarios demonstrate the superiority of CoSAM over state-of-the-art SAM-based methods.