DiffuSeg: Domain-Driven Diffusion for Medical Image Segmentation

In recent years, the deployment of supervised machine learning techniques for segmentation tasks has significantly increased. Nonetheless, the annotation process for extensive datasets remains costly, labor-intensive, and error-prone. While acquiring sufficiently large datasets to train deep learning models is feasible, these datasets often experience a distribution shift relative to the actual test data. This problem is particularly critical in the domain of medical imaging, where it adversely affects the efficacy of automatic segmentation models. In this work, we introduce DiffuSeg, a novel conditional diffusion model developed for medical image data, that exploits any labels to synthesize new images in the target domain. This allows a number of new research directions, including the segmentation task that motivates this work. Our method only requires label maps from any existing datasets and unlabelled images from the target domain for image diffusion. To learn the target domain knowledge, a feature factorization variational autoencoder is proposed to provide conditional information for the diffusion model. Consequently, the segmentation network can be trained with the given labels and the synthetic images, thus avoiding human annotations. Initially, we apply our method to the MNIST dataset and subsequently adapt it for use with medical image segmentation datasets, such as retinal fundus images for vessel segmentation and MRI images for heart segmentation. Our approach exhibits significant improvements over relevant baselines in both image generation and segmentation accuracy, especially in scenarios where annotations for the target dataset are unavailable during training. An open-source implementation of our approach can be released after reviewing. .

Authors

• Le Zhang Fuxi AI Lab, Netease
• Fuping Wu
• Kevin Bronik
• Bartlomiej W. Papiez

Keywords

• Image segmentation
• Biomedical imaging
• Annotations
• Training
• Diffusion models
• Bioinformatics
• Image synthesis
• Generative adversarial networks
• Testing
• Feature extraction
• Medical Imaging
• Medical Image Segmentation
• Imaging Data
• Cognitive Domains
• Domain Shift
• Diffusion Model
• Segmentation Model
• Image Generation
• MRI Images
• Segmentation Accuracy
• Retinal Images
• Target Domain
• Segmentation Task
• Synthetic Images
• Variational Autoencoder
• Segmentation Dataset
• MNIST Dataset
• Medical Datasets
• Medical Imaging Data
• Domain Features
• Segmentation Labels
• Target Model
• Domain Adaptation
• Segmentation Branch
• Downstream Segments
• Medical Image Analysis
• Image Reconstruction
• Segmentation Map
• domain transfer
• Humans
• Image Processing, Computer-Assisted
• Algorithms
• Deep Learning
• Databases, Factual
• Brain
• Image Interpretation, Computer-Assisted