BrainNormalizer: Anatomy-Informed Pseudo-Healthy Brain Reconstruction from Tumor MRI via Edge-Guided ControlNet

TL;DR

BrainNormalizer tackles the lack of subject-specific pre-tumor references by introducing an anatomy-informed diffusion framework that reconstructs pseudo-healthy brain MRIs from tumorous scans. It combines a two-stage training pipeline—inpainting-based domain adaptation of Stable Diffusion and edge-guided ControlNet conditioning—to enforce anatomical fidelity via boundary cues, with a deliberate misalignment inference strategy using contralateral edge maps. On BraTS2020, the approach yields superior perceptual realism (lower FID), higher structural coherence (higher SSIM), and greater clinical plausibility (lower false positives) compared with baselines. This work enables clinically useful pseudo-healthy references for treatment planning, longitudinal analysis, and counterfactual modeling, while opening avenues for 3D context, multi-contrast integration, and radiologist-validation studies.

Abstract

Brain tumors are among the most clinically significant neurological diseases and remain a major cause of morbidity and mortality due to their aggressive growth and structural heterogeneity. As tumors expand, they induce substantial anatomical deformation that disrupts both local tissue organization and global brain architecture, complicating diagnosis, treatment planning, and surgical navigation. Yet a subject-specific reference of how the brain would appear without tumor-induced changes is fundamentally unobtainable in clinical practice. We present BrainNormalizer, an anatomy-informed diffusion framework that reconstructs pseudo-healthy MRIs directly from tumorous scans by conditioning the generative process on boundary cues extracted from the subject's own anatomy. This boundary-guided conditioning enables anatomically plausible pseudo-healthy reconstruction without requiring paired non-tumorous and tumorous scans. BrainNormalizer employs a two-stage training strategy. The pretrained diffusion model is first adapted through inpainting-based fine-tuning on tumorous and non-tumorous scans. Next, an edge-map-guided ControlNet branch is trained to inject fine-grained anatomical contours into the frozen decoder while preserving learned priors. During inference, a deliberate misalignment strategy pairs tumorous inputs with non-tumorous prompts and mirrored contralateral edge maps, leveraging hemispheric correspondence to guide reconstruction. On the BraTS2020 dataset, BrainNormalizer achieves strong quantitative performance and qualitatively produces anatomically plausible reconstructions in tumor-affected regions while retaining overall structural coherence. BrainNormalizer provides clinically reliable anatomical references for treatment planning and supports new research directions in counterfactual modeling and tumor-induced deformation analysis.

Figures (6)

• Figure 1: Overall architecture of BrainNormalizer based on the ControlNet framework. (a) Stable Diffusion (SD) backbone that performs latent-space denoising conditioned on text and time embeddings. (b) ControlNet branch that processes structural edge maps and injects boundary-aware features into the frozen SD decoder via zero-convolution layers. Abbreviations: VAE, Variational Autoencoder; SD, Stable Diffusion.
• Figure 2: Step 1: Fine-tuning Stable Diffusion (SD). Both non-tumorous and tumorous MRI slices are provided as input pairs $(i, m, p, t)$, consisting of a 2D image slice, binary mask, text prompt, and diffusion timestep. For non-tumorous inputs, a mask from a tumor slice with the same slice index is reused to maintain spatial consistency and robustness. The masked image is encoded into a latent representation $z_t$ through the variational autoencoder (VAE) encoder, while the text prompt and timestep are embedded via the CLIP text encoder and time encoder, respectively. These embeddings are injected into the U-Net denoising network $\epsilon_\theta(z_t, p, t)$ within the SD framework, which predicts and removes noise to reconstruct the full image through the VAE decoder. Abbreviations: VAE, Variational Autoencoder; SD, Stable Diffusion.
• Figure 3: Step 2: ControlNet training for edge-guided anatomical reconstruction. ControlNet is built upon the fine-tuned SD model from Step 1. It receives the MRI slice, its binary mask, and the corresponding edge map $c$ (either $c_n$ or $c_d$). The ControlNet is initialized with the fine-tuned SD U-Net parameters, and their learned features are transmitted to the frozen SD decoder via zero-convolution layers ($1\times1$ convolutions initialized with zeros). Abbreviations: VAE, Variational Autoencoder; SD, Stable Diffusion.
• Figure 4: Step 3: Inference stage of BrainNormalizer. A deliberate misalignment strategy between tumorous inputs and non-tumorous guidance is applied to synthesize a pseudo-healthy brain. The model receives a tumorous slice $i_d$ with its mask $m_d$, a non-tumorous prompt $p_n$, and a mirrored edge map $c_{\text{mirrored}}$ to guide anatomically consistent reconstruction. Abbreviations: VAE, Variational Autoencoder; SD, Stable Diffusion.
• Figure 5: Visualization of pseudo-healthy brain reconstruction results for three representative subjects. Each pair of rows corresponds to the same subject: (a)–(c) show different cases. Columns show the tumorous MRI, reconstructed results by Wolleb et al. wolleb2022diffusion, SD Inpaint, and the proposed BrainNormalizer. The first row presents the input MRI and the corresponding pseudo-healthy reconstructed images, while the second row displays the tumor-overlaid input and the difference maps between the input and reconstructed images.