CEmb-SAM: Segment Anything Model with Condition Embedding for Joint Learning from Heterogeneous Datasets

TL;DR

CEmb-SAM introduces a novel approach to jointly learn ultrasound image segmentation from heterogeneous datasets by incorporating sub-group condition embeddings into the Segment Anything Model, enhancing generalization and performance.

Contribution

The paper proposes CEmb-SAM, a new method that integrates condition embeddings into SAM to adapt to different dataset sub-groups in medical image segmentation.

Findings

CEmb-SAM outperforms baseline methods on ultrasound segmentation tasks.

The approach effectively leverages dataset variability for improved generalization.

Experimental results demonstrate robustness across different anatomical structures.

Abstract

Automated segmentation of ultrasound images can assist medical experts with diagnostic and therapeutic procedures. Although using the common modality of ultrasound, one typically needs separate datasets in order to segment, for example, different anatomical structures or lesions with different levels of malignancy. In this paper, we consider the problem of jointly learning from heterogeneous datasets so that the model can improve generalization abilities by leveraging the inherent variability among datasets. We merge the heterogeneous datasets into one dataset and refer to each component dataset as a subgroup. We propose to train a single segmentation model so that the model can adapt to each sub-group. For robust segmentation, we leverage recently proposed Segment Anything model (SAM) in order to incorporate sub-group information into the model. We propose SAM with Condition Embedding block (CEmb-SAM) which encodes sub-group conditions and combines them with image embeddings from SAM. The conditional embedding block effectively adapts SAM to each image sub-group by incorporating dataset properties through learnable parameters for normalization. Experiments show that CEmb-SAM outperforms the baseline methods on ultrasound image segmentation for peripheral nerves and breast cancer. The experiments highlight the effectiveness of Cemb-SAM in learning from heterogeneous datasets in medical image segmentation tasks.