ProGiDiff: Prompt-Guided Diffusion-Based Medical Image Segmentation

TL;DR

ProGiDiff introduces a prompt-guided diffusion framework for medical image segmentation that leverages pre-trained models, enabling multi-class segmentation, human interaction, and cross-modality adaptation with strong experimental results.

Contribution

The paper presents a novel conditioning mechanism for diffusion models, allowing effective medical image segmentation guided by natural language prompts, adaptable across modalities with minimal training.

Findings

Outperforms previous segmentation methods on CT images

Supports multi-class segmentation via natural language prompts

Enables cross-modality adaptation with few-shot learning

Abstract

Widely adopted medical image segmentation methods, although efficient, are primarily deterministic and remain poorly amenable to natural language prompts. Thus, they lack the capability to estimate multiple proposals, human interaction, and cross-modality adaptation. Recently, text-to-image diffusion models have shown potential to bridge the gap. However, training them from scratch requires a large dataset-a limitation for medical image segmentation. Furthermore, they are often limited to binary segmentation and cannot be conditioned on a natural language prompt. To this end, we propose a novel framework called ProGiDiff that leverages existing image generation models for medical image segmentation purposes. Specifically, we propose a ControlNet-style conditioning mechanism with a custom encoder, suitable for image conditioning, to steer a pre-trained diffusion model to output segmentation masks. It naturally extends to a multi-class setting simply by prompting the target organ. Our experiment on organ segmentation from CT images demonstrates strong performance compared to previous methods and could greatly benefit from an expert-in-the-loop setting to leverage multiple proposals. Importantly, we demonstrate that the learned conditioning mechanism can be easily transferred through low-rank, few-shot adaptation to segment MR images.