Generative Enhancement for 3D Medical Images

TL;DR

3D medical image datasets are scarce due to privacy and cost. GEM-3D tackles this by using conditional diffusion with an informed 2D slice as patient prior and segmentation masks, trained in two stages (slice diffusion followed by volumetric tuning) and aided by bi-directional, overlapping sampling to ensure volumetric consistency. The framework enables dataset enhancement, counterfactual synthesis, and dataset-level de-enhancement, demonstrated on BraTS brain MRI and AbdomenCT-1K CT data. The method provides a memory-efficient, flexible approach to generate high-fidelity 3D volumes from existing data, with practical impact for data augmentation and privacy-preserving research in medical imaging.

Abstract

The limited availability of 3D medical image datasets, due to privacy concerns and high collection or annotation costs, poses significant challenges in the field of medical imaging. While a promising alternative is the use of synthesized medical data, there are few solutions for realistic 3D medical image synthesis due to difficulties in backbone design and fewer 3D training samples compared to 2D counterparts. In this paper, we propose GEM-3D, a novel generative approach to the synthesis of 3D medical images and the enhancement of existing datasets using conditional diffusion models. Our method begins with a 2D slice, noted as the informed slice to serve the patient prior, and propagates the generation process using a 3D segmentation mask. By decomposing the 3D medical images into masks and patient prior information, GEM-3D offers a flexible yet effective solution for generating versatile 3D images from existing datasets. GEM-3D can enable dataset enhancement by combining informed slice selection and generation at random positions, along with editable mask volumes to introduce large variations in diffusion sampling. Moreover, as the informed slice contains patient-wise information, GEM-3D can also facilitate counterfactual image synthesis and dataset-level de-enhancement with desired control. Experiments on brain MRI and abdomen CT images demonstrate that GEM-3D is capable of synthesizing high-quality 3D medical images with volumetric consistency, offering a straightforward solution for dataset enhancement during inference. The code is available at https://github.com/HKU-MedAI/GEM-3D.

Figures (8)

• Figure 1: Overview of GEM-3D Framework.(a) During training, our method is built upon volume diffusion and we sample a window of images and their corresponding masks as training samples. The informed slice, sampled from the volume window, combines with mask data as conditions. (b) For inference, we decouple conditional generation into the combination of mask volume and informed slice. Our designs include random starting positions, editable mask volume, and a selective or generative informed slice for increased variations, employing bi-directional propagation in sampling.
• Figure 2: Qualitative comparison on BraTS and AbdomenCT-1K for training and testing samples.(a) For training samples, our method synthesizes new samples by introducing variations through randomly chosen informed slices from the given volumes, even when only provided with the training split. In comparison, the baseline method outputs fitting results but still exhibits volumetric inconsistency. (b) For testing samples, our method leverages the additional information of informed slices in the true data and maintains high fidelity, resulting in superior details and improved volumetric consistency comparing with the baseline method.
• Figure 3: GEM-3D enables counterfactual synthesis. We show synthesis results under different scenarios with respect to informed slices and masks. The proposed method produces high-fidelity generation performances and offers solutions for counterfactual synthesis under different occasions.
• Figure 4: GEM-3D enables generative de-enhancement. We show three volumes samples on BraTS with the last sample exhibiting poorer visualization quality due to the absence of normalization. Using two choices of informed slices for the initial sampling, GEM-3D generates two types of well-normalized entire volumes through de-enhancement. It is important to note that counterfactual samples are generated for dataset-level normalization, and thus, they are not required to maintain all the anatomical details as observed in the true samples (first row).
• Figure 5: Ablation on Overlapped Inpainting. We present two cases for comparison, with each case displaying three consecutive slices in a sample. OI refers to overlapped inpainting.