ZECO: ZeroFusion Guided 3D MRI Conditional Generation

TL;DR

ZECO addresses the data scarcity challenge in 3D MRI segmentation by generating high-fidelity MRI volumes conditioned on 3D segmentation masks using a latent diffusion framework. It introduces a Spatial Transformation Module to encode MRIs into a compact latent space and a ZeroFusion structure that progressively maps segmentation masks to MRI features without overfitting pretrained diffusion models. Empirical results on BraTS 2016 and 2020 demonstrate superior quantitative and qualitative performance across FLAIR and T1 modalities, with ablations validating the contribution of ZeroFusion and the STM. The approach preserves inter-slice coherence and shows strong potential for scaling conditional MRI synthesis to other medical imaging tasks with sparse labels, thereby aiding data augmentation and diagnostic workflows.

Abstract

Medical image segmentation is crucial for enhancing diagnostic accuracy and treatment planning in Magnetic Resonance Imaging (MRI). However, acquiring precise lesion masks for segmentation model training demands specialized expertise and significant time investment, leading to a small dataset scale in clinical practice. In this paper, we present ZECO, a ZeroFusion guided 3D MRI conditional generation framework that extracts, compresses, and generates high-fidelity MRI images with corresponding 3D segmentation masks to mitigate data scarcity. To effectively capture inter-slice relationships within volumes, we introduce a Spatial Transformation Module that encodes MRI images into a compact latent space for the diffusion process. Moving beyond unconditional generation, our novel ZeroFusion method progressively maps 3D masks to MRI images in latent space, enabling robust training on limited datasets while avoiding overfitting. ZECO outperforms state-of-the-art models in both quantitative and qualitative evaluations on Brain MRI datasets across various modalities, showcasing its exceptional capability in synthesizing high-quality MRI images conditioned on segmentation masks.

Figures (3)

• Figure 1: ZECO generates FLAIR (left) and T1 (right) MRI modalities conditioning on segmentation masks.
• Figure 2: Pipeline. (a) The Spatial Transformation Module encodes MRI images into latent space. (b) The 3D U-Net reconstruct latent images across timesteps in the reverse diffusion process. (c) The ZeroFusion generates Middle-feature and Down-feature conditioned on segmentation masks $\mathcal{C}$ for controllable generation.
• Figure 3: ZECO generates coherent MRI slices, while SegGuidedDiff fails in regions without masks.