PRIMED / README.md

Upload README.md

d0b8f3e verified 9 days ago

5.99 kB

---
license: cc-by-nc-4.0
library_name: pytorch
pipeline_tag: image-to-image
tags:
  - medical-imaging
  - mri
  - brain
  - brats
  - missing-modality
  - image-synthesis
  - 3d
  - wacv2027
---
# 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**<br>
Atharva Atul Rege, Adinath Madhavrao Dukre, Imran Razzak

Code: [github.com/Atharva-Rege/PRIMED](https://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

```bash
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
```

```python
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](https://www.synapse.org/Synapse:syn51156910/wiki/621282) and
[IXI](https://brain-development.org/ixi-dataset/) data; please respect their terms of use.

	---
	license: cc-by-nc-4.0
	library_name: pytorch
	pipeline_tag: image-to-image
	tags:
	- medical-imaging
	- mri
	- brain
	- brats
	- missing-modality
	- image-synthesis
	- 3d
	- wacv2027
	---

	# 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<br>
	Atharva Atul Rege, Adinath Madhavrao Dukre, Imran Razzak

	Code: [github.com/Atharva-Rege/PRIMED](https://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

	```bash
	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
	```

	```python
	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](https://www.synapse.org/Synapse:syn51156910/wiki/621282) and
	[IXI](https://brain-development.org/ixi-dataset/) data; please respect their terms of use.