Zero-Shot Brain MRI Inpainting with 2.5D Unconditional Flow Priors

Arnela Hadzic; Franz Thaler; Simon Johannes Joham; Martin Urschler

cs.CVOct 6, 2026

Zero-Shot Brain MRI Inpainting with 2.5D Unconditional Flow Priors

Authors: Arnela Hadzic, Franz Thaler, Simon Johannes Joham, Martin Urschler

Organizations: Institute for Medical Informatics, Statistics and Documentation, Medical University of Graz, Graz, Austria

Abstract

Generative inpainting of brain MRI volumes is essential for synthesizing healthy tissue in pathological regions, improving the accuracy and reliability of automated downstream brain analysis applications such as image registration, brain extraction, and segmentation. However, standard 3D approaches are computationally prohibitive, while efficient 2D slice-wise methods suffer from severe inter-slice discontinuities. Furthermore, traditional models rely on conditional training, requiring task-specific learning of masked inputs. We propose a zero-shot brain MRI inpainting framework utilizing 2.5D unconditional flow priors to capture spatial context along the superior-inferior axis without the overhead of full 3D convolutions. During training, our flow matching model learns the joint distribution of adjacent axial slice triplets, modeling the manifold of healthy brain anatomy while explicitly excluding pathological regions from the loss function. At inference, the model processes the input triplets autoregressively along the depth axis. We employ the Restora-Flow solver to constrain the unconditional prior using the input mask, achieving accurate zero-shot inpainting. Evaluations show our 2.5D strategy resolves the structural discontinuities of 2D baselines, synthesizing plausible healthy tissue while maintaining volumetric consistency across the axial, sagittal, and coronal planes. As a final step, we generate and average an ensemble of multiple stochastic reconstructions to form the final prediction. Quantitative results benchmarked on the official BraTS 2026 Inpainting Challenge validation set demonstrate the effectiveness of our proposed approach, yielding an SSIM of 0.816 $\pm$ 0.112, MSE of 0.007 $\pm$ 0.005, and PSNR of 22.923 $\pm$ 4.343. Code is available at https://github.com/imigraz/brats2026-inpainting.

Figures & tables

Panel of Figure 1 — Figure 1: Impact of the number of ODE steps on reconstruction quality, measured via SSIM (left) and PSNR (right) on the internal validation set using a single reconstruction per image obtained from a model trained for 2,000 epochs.

Panel of Figure 2 — Figure 2: Qualitative comparison of volumetric reconstructions between the 2D baseline and the proposed 2.5D strategy. Rows denote the axial, sagittal, and coronal anatomical planes. Columns from left to right display the masked input ( t1n-voided ), the ground truth ( t1n ), the 2D reconstruction, and the 2.5D reconstruction. The boxes indicate the unhealthy (red) and healthy (cyan) masked region. While the 2D approach introduces severe inter-slice discontinuities visible in the sagittal and coronal planes, the 2.5D strategy successfully preserves spatial consistency across all planes. The reconstructions are single-pass predictions of models trained for 2,000 epochs using 32 ODE steps.

Strategy	SSIM $\uparrow$	MSE $\downarrow$	PSNR $\uparrow$
2D	0.727 $\pm$ 0.153	0.014 $\pm$ 0.010	18.899 $\pm$ 2.804
2.5D	0.789 $\pm$ 0.131	0.010 $\pm$ 0.008	20.556 $\pm$ 3.004

Table 1: Quantitative comparison of reconstruction quality between the standard 2D and the proposed 2.5D strategies on the internal validation set. Scores were computed from a single reconstruction using models trained for 2,000 epochs and 32 ODE steps.

Panel of Figure 3 — Figure 3: Impact of ensemble size ( $N$ ) and ODE steps on reconstruction quality, measured via SSIM (left), MSE (center), and PSNR (right) on a subset of 20 cases from the internal validation set. Scores were obtained using a model trained for 2,000 epochs.

Model	SSIM $\uparrow$	MSE $\downarrow$	PSNR $\uparrow$
$N=1$	0.793 $\pm$ 0.122	0.009 $\pm$ 0.007	21.888 $\pm$ 4.328
$N=5$	0.813 $\pm$ 0.113	0.007 $\pm$ 0.006	22.793 $\pm$ 4.335
$N=10$	0.816 $\pm$ 0.112	0.007 $\pm$ 0.005	22.923 $\pm$ 4.343

Table 2: Quantitative results on the challenge validation set. We compare single reconstructions (

N=1

) against ensembles (

N=5,10

) obtained via voxel-wise averaging. We used 32 ODE steps per prediction and a model trained for 10,000 epochs.

Panel of Figure 4 — Figure 4: Qualitative inpainting results of the proposed 2.5D approach on the official validation set where the ground truth is unavailable. The figure highlights three distinct patient cases for which we show the axial, sagittal, and coronal planes for the voided input (col 1), three distinct single reconstructions (col 2-4), and the final ensemble mean with $N=10$ (col 5). The single reconstructions were obtained from a model trained for 10,000 epochs using 32 ODE steps. Yellow boxes indicate the masked healthy and unhealthy region for which a distinction is not possible. The results demonstrate the effectiveness of our approach, producing anatomically plausible reconstructions that maintain structural consistency across all three volumetric planes.

Explore similar work

Sep 3, 2026cs.CV

RARF: Region-Aware Rectified Flows for 3D Brain MRI Inpainting

Medical image inpainting has the potential to improve automated brain MRI analysis by reconstructing healthy tissue within pathological regions. We introduce RARF, a task-agnostic region-aware rectified flow framework for masked data generation. We instantiate the framework for 3D brain MRI inpainting as our submission to the BraTS Inpainting Challenge 2026. RARF restricts the stochastic interpolation process to the inpainting region, while the observed voxels remain fixed and provide patient-specific anatomical context. A three-dimensional neural network receives the partially voided image, with Gaussian noise filling the missing region, together with the inpainting mask and the corresponding timestep. The model is trained using masked flow-matching and reconstruction-consistency objectives, combined with mask-aware preprocessing and data augmentation. During inference, the learned velocity field transports the initial noise toward a plausible reconstruction of the missing tissue, which is then combined with the unchanged observed anatomy. Experiments under the BraTS evaluation protocol show that the proposed approach produces competitive reconstructions while maintaining anatomical consistency. Source code is available at: https://github.com/TomasGuija/rarf.

Tomas Guija-Valiente, Blanca Rodriguez-Gonzalez, Norberto Malpica +1

Medical Image Analysis and Biometry Lab, Universidad Rey Juan Carlos, Madrid, Spain

Inpainting Functional Magnetic Resonance Imaging Rectified Flow 3D Generation

Ask PDF

Jul 30, 2026cs.CV

Now You Have My Healthy Attention: A U-DiT for Brain-MRI Inpainting

The ASNR-MICCAI BraTS Local Synthesis (Inpainting) task asks for the anatomically plausible completion of healthy brain tissue within a masked region of a T1-weighted MRI, providing a tumor-free anatomical reference for downstream analysis. As the task is scored by distortion metrics (SSIM, PSNR, MSE), we build a deterministic regression model and focus on giving it inductive biases tailored to inpainting. Our network follows the U-DiT principle of performing self-attention on a downsampled token grid: a volumetric encoder-decoder imports long-range context through a downsampled global self-attention block with three-dimensional rotary position embeddings, while convolutions and skip connections preserve high-frequency detail. Two ideas drive our results. First, we constrain the attention so that occluded ("void") tokens attend only to known-healthy tokens of the same volume, with a learned bias toward each query's contralateral homologue, forcing the completion to be inferred from observed anatomy rather than from other unknown regions. Second, we add a contralateral-symmetry input that supplies the mirrored healthy hemisphere as a patient-specific prior; since the brain is approximately bilaterally symmetric and lesions are typically unilateral, this prior improves the distortion metrics at matched structural similarity. On the official BraTS-2026 validation leaderboard our submission reaches a mean healthy-region SSIM of

0.864

, PSNR of

24.7

,dB and MSE of

4.6{\times}10^{-3}

over

219

cases. We further analyse the residual smoothness inherent to distortion-optimal regression and discuss its implications for anatomical realism.

Danilo Danese, Angela Lombardi, Tommaso Di Noia

Politecnico di Bari, Bari, Italy

Brain Tumor Segmentation Inpainting Anatomy Distortion Diet

Ask PDF

Oct 29, 2025eess.IV

MedForj: An open, large-scale foundational generative prior for high-resolution 3D brain MRI

This work introduces MedForj, a suite of 3D foundational generative priors based on diffusion models. The MedForj models were trained on

72{,}659

1~mm isotropic 3D

T_1

-weighted MRI human brain image volumes from

38{,}174

subjects, drawn from a curated corpus of

80{,}675

volumes from

42{,}506

subjects spanning

38

publicly available datasets. These training images were manually inspected to exclude those with poor quality and excessive pathology, and otherwise were minimally processed. The models include six different diffusion training strategies: rectified flow, latent diffusion rectified flow, flow matching, velocity prediction, clean prediction, and noise prediction. Image samples produced by each of these models were compared to each other and against real, ground truth data under downstream segmentation distributions, FID, five inverse problems, and blind human inspection in an observer study. Flow matching was the strongest strategy overall, achieving the best inverse problem solving results at

28.80

~dB PSNR and

0.874

SSIM averaged over the five forward problems, the highest rate of reconstructions judged real by blind human raters at

72.6\%

, and the closest per-structure match to real segmented anatomy in a permutation test. It was not best everywhere: rectified flow produced the most convincing unconditional samples in the observer study and the best FID, and the latent rectified-flow model achieved the smallest joint distributional distance to real anatomy. No other strategy, however, performed consistently well across all four evaluations. We therefore recommend flow matching as the default MedForj prior, while releasing every strategy so that the choice can be revisited per application. All model weights and corresponding code are publicly available at https://github.com/piksl-research/medforj.

Samuel W. Remedios, Aaron Carass, Jerry L. Prince +1

Image Analysis and Communications Laboratory, Johns Hopkins University · Department of Neurology, Johns Hopkins Hospital

Medical Image Generation Diffusion Priors

Ask PDF