Generative Medical Segmentation

TL;DR

This paper introduces Generative Medical Segmentation (GMS), a novel generative approach using pre-trained vision models to improve medical image segmentation accuracy and generalization across diverse datasets.

Contribution

GMS leverages a generative framework with pre-trained vision models, reducing overfitting and enhancing cross-dataset generalization in medical segmentation tasks.

Findings

GMS outperforms existing discriminative and generative models on five datasets.

GMS demonstrates strong cross-center generalization within the same modality.

Fewer trainable parameters in GMS reduce overfitting risk.

Abstract

Rapid advancements in medical image segmentation performance have been significantly driven by the development of Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs). These models follow the discriminative pixel-wise classification learning paradigm and often have limited ability to generalize across diverse medical imaging datasets. In this manuscript, we introduce Generative Medical Segmentation (GMS), a novel approach leveraging a generative model to perform image segmentation. Concretely, GMS employs a robust pre-trained vision foundation model to extract latent representations for images and corresponding ground truth masks, followed by a model that learns a mapping function from the image to the mask in the latent space. Once trained, the model generates an estimated segmentation mask using the pre-trained vision foundation model to decode the predicted latent representation back into the image space. The design of GMS leads to fewer trainable parameters in the model which reduces the risk of overfitting and enhances its generalization capability. Our experimental analysis across five public datasets in different medical imaging domains demonstrates GMS outperforms existing discriminative and generative segmentation models. Furthermore, GMS is able to generalize well across datasets from different centers within the same imaging modality. Our experiments suggest GMS offers a scalable and effective solution for medical image segmentation. GMS implementation and trained model weights are available at https://github.com/King-HAW/GMS.