Source-linked AI summary
RARF: Region-Aware Rectified Flows for 3D Brain MRI Inpainting
Tomas Guija-Valiente, Blanca Rodriguez-Gonzalez, Norberto Malpica, Angel Torrado-Carvajal
TL;DR
Automated brain MRI analysis can be unreliable when pathology replaces healthy anatomy, motivating reconstruction of subject-specific healthy tissue. RARF uses region-aware rectified flow to generate the missing region while keeping observed anatomy fixed, producing anatomically consistent reconstructions in preliminary BraTS results. Its broader evaluation highlights a trade-off between distortion metrics and perceptual or anatomical plausibility.
Problem
Lesions can make downstream brain MRI analysis unreliable, while healthy reference images are generally unavailable for neuro-oncology applications.
Method
RARF restricts rectified-flow interpolation and supervision to the masked region while keeping observed anatomy fixed as conditioning context.
Results
RARF produces anatomically consistent reconstructions in the BraTS-specific instantiation, with official validation using K = 50 averaged samples to prioritize distortion-based metrics.
Takeaways & Limitations
RARF provides a flexible region-aware formulation intended for arbitrary inpainting masks beyond the BraTS benchmark.
Takeaways & Limitations
Averaging multiple samples can improve distortion metrics while producing smoother, blurrier reconstructions and reducing fine anatomical detail.
Abstract
from arXiv · showhide
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.
1 Introduction
Brain MRI analysis methods often assume healthy anatomy, but lesions can make downstream tools unreliable and healthy reference images unavailable. RARF addresses this setting by region-aware rectified-flow inpainting that synthesizes healthy tissue while preserving observed anatomy.
- Motivation: Many automated brain MRI analysis methods assume that input images represent healthy anatomy, an assumption problematic in neuro-oncology where healthy images are generally unavailable.Lesions can also make brain extraction, tissue segmentation, and anatomical parcellation unreliable.
- Motivation: Healthy tissue synthesis can provide a subject-specific anatomical proxy to mitigate pathology-induced bias in subsequent analyses.
- Problem setting: The BraTS Inpainting Challenge formulates healthy tissue synthesis as reconstructing tumor-affected regions from partially masked brain MRI with anatomically plausible, tumor-free tissue.
- Prior approaches: Classical inpainting methods are often insufficient for large pathological regions because complete anatomical structures must be synthesized rather than locally propagated.Deep learning methods instead learn data-driven priors for irregular regions, lesion filling, pathology-free reconstruction, and healthy tissue synthesis.
- Proposed approach: RARF restricts rectified-flow interpolation and supervision to the target region while preserving observed anatomy as conditioning context.The proposed BraTS instantiation is intended to generate plausible healthy tissue without modifying surrounding anatomy.
2 Methods
RARF preprocesses brain MRI cases and performs region-aware rectified-flow inpainting while preserving observed anatomy. Training supervises healthy tissue within the inpainting region, and inference restores the masked region from noise using a learned velocity field.
- Data and pre-processing: The BraTS Local Inpainting dataset contains 1,251 training and 219 validation skull-stripped, coregistered T1-weighted MRI volumes at 240 × 240 × 155 voxels with 1mm isotropic spacing.
- Data and pre-processing: Each training input is voided over the complete inpainting mask, while supervision is restricted to the healthy-tissue mask because tumor intensities are invalid reconstruction targets.
- Data and pre-processing: Mask-aware preprocessing crops volumes using inference-available foreground information, normalizes each case to [0, 1], and augments training with five mask variants per case.
- Region-aware rectified flow: RARF evolves only masked voxels along a localized rectified-flow path, while the observed anatomy remains fixed and conditions the generation process.
- Inference: At inference, Gaussian noise fills the mask, the velocity field is integrated from t = 0 to t = 1, and observed voxels are restored after each update.
- Network and training objective: A 3D U-Net receives the current state and complete mask, and training combines masked flow matching with endpoint, voxel-wise, and structural objectives focused on healthy tissue.
- Inference: Averaging multiple samples can improve distortion metrics but may blur fine anatomical structures, creating a distortion–perception trade-off.
3 Results
The results compare inference strategies through official BraTS validation ablations and challenge evaluation. Multi-sample averaging improves distortion-based fidelity but can smooth anatomical detail, while the released challenge submission uses K = 50 averaging.
- Inference ablation: The ablation compares inference strategies quantitatively through Synapse submissions and qualitatively through representative reconstructed volumes, using the same trained checkpoint.All reported ablations use exponential moving average model parameters with decay β = 0.999.
- Inference ablation: Representative-sample strategies improve distortion metrics over single-sample inference while better preserving sample-level detail, but all K = 30 multi-sample strategies require substantially greater inference time.The comparison reflects a distortion–perception trade-off rather than a uniformly superior inference strategy.
- Main challenge results: The official challenge submission uses a separately trained checkpoint and averages K = 50 samples per case to prioritize distortion-based evaluation metrics.The corresponding model weights and source code are publicly released for reproducibility.
- Inference ablation: Averaging K = 30 samples generally improves distortion-based fidelity but may smooth fine anatomical structures relative to individual samples.Figure 2 restricts absolute-error maps to the synthesized region and uses a common color scale.
4 Discussion
RARF is presented as a general region-aware inpainting framework with flexible configurations and arbitrary masks, while the discussion emphasizes that distortion metrics do not fully capture perceptual or anatomical quality. The authors therefore motivate complementary perceptual, anatomy-aware, and downstream-task evaluation.
- Scope and generality: RARF is designed as a general region-aware inpainting framework rather than one tailored exclusively to BraTS, supporting different training and inference configurations.The formulation is described as applicable to arbitrary inpainting masks.
- Evaluation context: Medical inpainting is highly constrained because masked regions correspond to specific underlying anatomy rather than many equally valid completions.This makes reference-based distortion metrics more informative than in less constrained natural-image settings.
- Evaluation limitations: Lower distortion does not necessarily imply better perceptual quality or anatomical plausibility, since averaging improved MSE and PSNR while producing smoother, blurrier reconstructions.The discussion motivates complementing voxel-wise metrics with perceptual, anatomy-aware, and downstream-task evaluation.
- Disclosure: The study reports no conflict of interests relevant to the research.The acknowledgment states that the study was funded by the MAGERIT-CM project.