PRIMED weights

Weights for PRIMED (PRior-Informed Missing-modality synthesis by Endpoint Displacement), from the WACV 2027 paper

When Iteration Cannot Help: Single-Pass Missing-Modality MRI Synthesis and What Actually Limits It
Atharva Atul Rege, Adinath Madhavrao Dukre, Imran Razzak

Code: github.com/Atharva-Rege/PRIMED

PRIMED synthesises a missing MRI contrast from the acquired ones in a single network evaluation, with no noise schedule and no solver. It works in a fixed 3D Haar wavelet basis and starts from the mean of the acquired contrasts (the informed prior). A 3D U-Net predicts the displacement from that prior to the target. A learned target-contrast embedding lets one set of weights synthesise every contrast.

Files

File	Training data	Synthesises	Released weights
primed_brats.safetensors	BraTS 2023 GLI, 1032 subjects	any of T1n, T1c, T2w, T2f from the other three	best checkpoint, step 15 000 of up to 20 000
primed_brats_avail.safetensors	as primed_brats, plus random input availability	any of the four from one, two or three acquired contrasts	final, 8 000 further steps
primed_brats_robust.safetensors	as primed_brats, plus misregistration and thick-slice augmentation	any of the four from three possibly misregistered or thick-slice contrasts	final, 8 000 further steps
primed_ixi.safetensors	IXI, 477 subjects	any of T1, T2, PD from the other two	best checkpoint, step 28 000 of up to 30 000

Each file holds the network weights only, 97.7M parameters in float32 (about 391 MB). The two fine-tunes start from primed_brats.safetensors and are released with their final raw weights, as reported in the paper.

Usage

git clone https://github.com/Atharva-Rege/PRIMED.git
cd PRIMED
pip install -e .
hf download adidukre/PRIMED --local-dir /path/to/weights

python synthesize.py \
  --weights /path/to/weights/primed_brats.safetensors \
  --target t1c \
  --inputs t1n=/path/to/subject-t1n.nii.gz t2w=/path/to/subject-t2w.nii.gz t2f=/path/to/subject-t2f.nii.gz \
  --output /path/to/subject-t1c-synthesized.nii.gz

import torch
from primed import load_model
from primed.data import load_volume

device = "cuda" if torch.cuda.is_available() else "cpu"
model = load_model("/path/to/weights/primed_brats.safetensors", device)
inputs = torch.stack(
    [load_volume(f"/path/to/subject-{m}.nii.gz") for m in ("t1n", "t2w", "t2f")]
).unsqueeze(0).to(device)
prior = inputs.mean(dim=1)
t1c = model.synthesize(prior, inputs, torch.tensor([1], device=device))

Inputs. Co-registered volumes on the 1 mm BraTS grid (240 x 240 x 155). The BraTS models expect skull-stripped volumes; primed_ixi expects volumes that are not skull-stripped, as IXI provides them. Each volume is clipped to its 0.1 and 99.9 percentiles, rescaled to [0, 1], padded to 160 slices and cropped in plane to 224 x 224; load_volume and synthesize.py do this automatically. The contrasts are, in class order:

• BraTS: T1n, T1c, T2w, T2f;
• IXI: T1, T2, PD.

The acquired contrasts are passed in this order with the target left out. Output: the target contrast in the same normalised [0, 1] range.

Results

On 216 held-out BraTS 2023 GLI subjects, primed_brats reaches:

• whole-volume PSNR of 29.372 / 28.137 / 28.953 / 27.778 dB for T1 / T1c / T2 / FLAIR;
• 28.560 dB on average, with 0.9439 SSIM;
• 0.152 s per volume on one NVIDIA RTX PRO 6000 Blackwell.

On cohorts not seen in training, primed_brats reaches:

• 26.555 dB on 60 BraTS-MEN training subjects;
• 26.413 dB on the 141-subject BraTS-MEN validation release;
• 26.398 dB on 99 BraTS-PED subjects.

primed_ixi reaches 31.436 dB on 96 held-out IXI subjects.

With fewer acquired contrasts, primed_brats_avail reaches 28.36 / 27.35 / 25.97 dB with three / two / one inputs. The zero-shot primed_brats reaches 28.56 / 25.40 / 22.37 dB.

primed_brats_robust gives up 0.12 dB on registered inputs and reduces the loss under every perturbation tested, for example:

Inputs	primed_brats	primed_brats_robust
one input misregistered by 2 mm, 2°	26.60 dB	27.28 dB
inputs from 5 mm slices	24.66 dB	26.82 dB

Training

The training recipe:

• 3D Haar wavelet basis, with the low-pass band divided by 3;
• five-level MONAI DiffusionModelUNet with channels (64, 128, 128, 256, 256), three residual blocks per level and one self-attention block in the bottleneck;
• mean squared error on the displacement at the source endpoint;
• AdamW with learning-rate warmup and cosine decay, gradient clipping and bf16 mixed precision;
• EMA decay 0.9999 and random flips;
• checkpoint selection on 12 held-out cases.

The BraTS model was fine-tuned from a network of the same architecture trained along the full bridge path, which is not released. The IXI model was trained from random initialisation. The configs and subject lists are in the code repository.

Intended use and limitations

These are research models, not medical devices, and are not intended for clinical decisions. A synthesised T1c is a conditional mean: it cannot recover contrast enhancement that the acquired contrasts do not reveal, and it may show enhancement that is not there. The models assume co-registered inputs on the BraTS grid. Performance drops on misregistered or thick-slice inputs, which primed_brats_robust partly addresses. The models were trained on BraTS and IXI data only.

License

CC BY-NC 4.0, for non-commercial use. The models are trained on BraTS 2023 and IXI data; please respect their terms of use.