[Paper Review] I-MedSAM: Implicit Medical Image Segmentation with Segment Anything

Xiaobao Wei, Jiajun Cao+4 more|arXiv (Cornell University)|Nov 28, 2023

Radiomics and Machine Learning in Medical Imaging4 citations

TL;DR

I-MedSAM proposes a novel implicit medical image segmentation framework that combines the Segment Anything Model (SAM) with Implicit Neural Representations (INRs) for improved generalization and boundary accuracy. By integrating a frequency-aware adapter and uncertainty-guided sampling in a coarse-to-fine INR decoder, I-MedSAM achieves state-of-the-art performance with only 1.6M trainable parameters, outperforming both discrete and continuous methods on 2D medical segmentation benchmarks.

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: https://github.com/ucwxb/I-MedSAM.

Motivation & Objective

• To address the limitations of discrete representations in medical image segmentation, such as spatial inflexibility and poor boundary delineation.
• To improve cross-domain generalization and resolution scalability in medical image segmentation using continuous representations.
• To enhance SAM’s performance on medical images by incorporating frequency-domain information through a novel adapter.
• To optimize implicit neural representation learning via a coarse-to-fine INR decoder with uncertainty-guided sampling for efficient and accurate inference.

Proposed method

• I-MedSAM uses a frequency adapter to inject high-frequency information from the DFT spectrum into SAM features during parameter-efficient fine-tuning (PEFT), improving boundary sensitivity.
• The model employs a two-stage implicit neural representation (INR) decoder: a coarse INR predicts initial segmentation, followed by a refined INR on selected points.
• Uncertainty-guided sampling selects Top-K variance feature points from the coarse INR output to focus refinement on ambiguous or complex regions.
• The INR decoder maps image features and coordinate grids to continuous segmentation outputs, enabling resolution-invariant inference.
• The entire framework is trained end-to-end with only 1.6M trainable parameters, leveraging LoRA for efficient fine-tuning of SAM.
• The method integrates spatial and frequency-domain features to enhance representation capacity while maintaining generalization under domain shifts.

Figure 2 : The overall pipeline of I-MedSAM. First, given the medical images and a coarse bounding box as a prompt, I-MedSAM utilizes the medical image encoder and the prompt encoder to generate discrete features. For the medical image encoder, we design low-rank adapters and frequency adapters to e

Experimental results

Research questions

• RQ1 Can implicit neural representations improve the generalization and resolution scalability of medical image segmentation models?
• RQ2 How does incorporating frequency-domain information into SAM features affect segmentation boundary accuracy?
• RQ3 Can uncertainty-guided sampling improve the efficiency and accuracy of implicit segmentation in medical images?
• RQ4 Does a coarse-to-fine INR decoder with selective sampling outperform standard INR or discrete prediction heads in medical segmentation?
• RQ5 To what extent can a lightweight, parameter-efficient adapter enhance SAM’s performance on diverse medical imaging domains?

Key findings

• I-MedSAM achieves a Dice score of 86.28% on the BCV dataset and 79.86% on the AMOS dataset, surpassing state-of-the-art discrete and continuous methods.
• The ablation study shows that a LoRA rank of 4 yields optimal performance with minimal parameters, confirming the efficiency of the PEFT strategy.
• The frequency adapter improves performance by 0.61% Dice when using amplitude spectrum (86.80%) compared to no adapter (86.19%), with amplitude being more critical than phase.
• Uncertainty-guided sampling with Top-12.5% of feature points achieves 86.80% Dice, outperforming full sampling and demonstrating that selective sampling is both effective and efficient.
• I-MedSAM generalizes robustly across domain shifts, as shown in qualitative comparisons on Kvasir-Sessile and BCV datasets, with superior boundary delineation.
• The model maintains high performance with only 1.6M trainable parameters, significantly fewer than baseline methods like nnUNet (126.6M) and SAMed (3.9M).

Figure 3 : The illustration of the proposed frequency adapters and LoRA in the image encoder. With these adapters, I-MedSAM can be trained with minorities of parameters while obtaining high-quality segmentation boundaries.