[Paper Review] DiffuseExpand: Expanding dataset for 2D medical image segmentation using diffusion models

Shitong Shao, Xiaohan Yuan+4 more|arXiv (Cornell University)|Apr 26, 2023

Generative Adversarial Networks and Image Synthesis8 citations

TL;DR

DiffuseExpand uses diffusion probabilistic models to generate high-quality, paired image-mask samples for 2D medical image segmentation, enabling efficient data expansion and improved segmentation performance. It introduces a four-stage pipeline with mask-first synthesis, classifier guidance, accelerated sampling, and a final quality filter.

ABSTRACT

Dataset expansion can effectively alleviate the problem of data scarcity for medical image segmentation, due to privacy concerns and labeling difficulties. However, existing expansion algorithms still face great challenges due to their inability of guaranteeing the diversity of synthesized images with paired segmentation masks. In recent years, Diffusion Probabilistic Models (DPMs) have shown powerful image synthesis performance, even better than Generative Adversarial Networks. Based on this insight, we propose an approach called DiffuseExpand for expanding datasets for 2D medical image segmentation using DPM, which first samples a variety of masks from Gaussian noise to ensure the diversity, and then synthesizes images to ensure the alignment of images and masks. After that, DiffuseExpand chooses high-quality samples to further enhance the effectiveness of data expansion. Our comparison and ablation experiments on COVID-19 and CGMH Pelvis datasets demonstrate the effectiveness of DiffuseExpand. Our code is released at https://github.com/shaoshitong/DiffuseExpand.

Motivation & Objective

• Address data scarcity in medical image segmentation due to privacy and labeling constraints.
• Develop a diffusion-model-based pipeline to generate paired image and mask samples with high diversity and alignment.
• Improve data expansion quality by stage-wise conditioning, classifier guidance, and sample screening.
• Demonstrate effectiveness on COVID-19 chest X-ray and CGMH Pelvis datasets.
• Provide open-source code and validate improvements through ablations and comparisons.

Proposed method

• Stage I: Fine-tune a pre-trained diffusion probabilistic model (DPM) without classifier guidance to map distributions to medical images.
• Stage II: Train a segmenter to enable classifier guidance and assess image-mask alignment via Dice and BCE losses.
• Stage III: Use DPM Solver++ to synthesize image-mask pairs under multiple conditions, enabling conditional sampling based on classifier guidance.
• Stage IV: Discard poor-quality samples using a trained network and a Dice-threshold to ensure high-quality expanded data.
• Introduce a continuous-time probability flow ODE formulation for conditional sampling and leverage classifier guidance with a temperature parameter tau to control gradient magnitudes.
• Provide ablation studies to compare classifier-free vs. classifier-guided sampling and assess the impact of tau/s on diversity and Dice performance.

Experimental results

Research questions

• RQ1 Can diffusion models generate diverse, well-aligned image-mask pairs for 2D medical image segmentation?
• RQ2 Does staged conditioning and classifier guidance improve the quality and utility of expanded segmentation datasets?
• RQ3 What is the impact of sample screening (Stage IV) on Dice scores and sample diversity in data expansion?
• RQ4 How do diffusion-based expansion methods compare to GAN-based or augmentation-based approaches in the medical imaging domain?

Key findings

• DiffuseExpand yields high-fidelity, diverse Image-Mask pairs that improve Dice scores on tested datasets.
• Classifier guidance with a trained segmenter, combined with Stage IV screening, enhances alignment between images and masks.
• Using the tau-based conditioning for guidance yields higher Dice scores than gradient-scaling value s in comparable settings.
• Diffusion-based expansion outperforms Synthmed and XLsor in Dice-based segmentation performance under the tested setups.
• Stage IV screening reduces sample diversity but increases overall quality, contributing to better downstream segmentation.