I-MedSAM: Implicit Medical Image Segmentation with Segment Anything

Xiaobao Wei,Jiajun Cao,Yizhu Jin,Ming Lu,Guangyu Wang,Shanghang Zhang

2024-07-11

Abstract:With the development of Deep Neural Networks (DNNs), many efforts have been made to handle medical image segmentation. Traditional methods such as nnUNet train specific segmentation models on the individual datasets. Plenty of recent methods have been proposed to adapt the foundational Segment Anything Model (SAM) to medical image segmentation. However, they still focus on discrete representations to generate pixel-wise predictions, which are spatially inflexible and scale poorly to higher resolution. In contrast, implicit methods learn continuous representations for segmentation, which is crucial for medical image segmentation. In this paper, we propose I-MedSAM, which leverages the benefits of both continuous representations and SAM, to obtain better cross-domain ability and accurate boundary delineation. Since medical image segmentation needs to predict detailed segmentation boundaries, we designed a novel adapter to enhance the SAM features with high-frequency information during Parameter-Efficient Fine-Tuning (PEFT). To convert the SAM features and coordinates into continuous segmentation output, we utilize Implicit Neural Representation (INR) to learn an implicit segmentation decoder. We also propose an uncertainty-guided sampling strategy for efficient learning of INR. Extensive evaluations on 2D medical image segmentation tasks have shown that our proposed method with only 1.6M trainable parameters outperforms existing methods including discrete and implicit methods. The code will be available at: <a class="link-external link-https" href="https://github.com/ucwxb/I-MedSAM" rel="external noopener nofollow">this https URL</a>.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

### Problems Addressed by the Paper

The paper aims to address several key issues in medical image segmentation:

1. **Spatial Flexibility and Resolution Scalability**:
- Existing medical image segmentation methods mainly rely on discrete representations, which have poor spatial flexibility and resolution scalability when dealing with images of different resolutions. Discrete representations tend to introduce discretization artifacts when scaled to higher resolutions, affecting segmentation accuracy.

2. **Accurate Depiction of Boundary Details**:
- Medical image segmentation requires precise depiction of detailed boundaries, which is crucial for distinguishing different tissues or anatomical structures. Existing methods exhibit blurriness and inaccuracy in extracting fine boundary details.

3. **Cross-Domain Generalization Ability**:
- Existing methods have poor generalization ability when faced with different datasets or images of different resolutions. This limits their robustness and reliability in practical applications.

### Solution

To address the above challenges, the authors propose **I-MedSAM** (Implicit Medical Image Segmentation with Segment Anything), which combines the advantages of continuous representation and the Segment Anything Model (SAM). The specific innovations include:

1. **Frequency Adapter**:
- Introducing a Frequency Adapter to extract high-frequency information from the frequency domain, enhancing feature representation and improving the accuracy of boundary depiction.

2. **Implicit Neural Representation Decoder**:
- Designing a two-stage Implicit Neural Representation Decoder, consisting of a "coarse" decoder and a "fine" decoder. Using an Uncertainty-Guided Sampling (UGS) strategy, high-variance feature points are selected for refinement, improving segmentation accuracy.

3. **Parameter-Efficient Fine-Tuning**:
- Utilizing Parameter-Efficient Fine-Tuning (PEFT) technology to fine-tune only a small number of parameters, allowing the model to maintain high performance while having lower computational costs.

### Experimental Results

- **Quantitative Comparison**:
- In binary polyp segmentation and multi-class organ segmentation tasks, I-MedSAM significantly outperforms existing discrete and implicit methods, especially with fewer parameters.

- **Robustness and Generalization Ability**:
- I-MedSAM performs excellently across different resolutions and datasets, demonstrating good cross-domain generalization ability and robustness to data variations.

- **Boundary Quality**:
- Evaluated by Hausdorff distance, I-MedSAM also excels in boundary quality, accurately depicting boundary details in images.

In summary, by combining the advantages of continuous representation and SAM, I-MedSAM addresses the shortcomings of existing medical image segmentation methods in terms of spatial flexibility, boundary detail depiction, and cross-domain generalization ability, providing a new solution for medical image segmentation.