Request access to the DGR model weights
These weights accompany Long et al., Radiology Advances 2026 (doi:10.1093/radadv/umag031). They were trained on clinical prostate MRI and are released for non-commercial research use only. Every request is reviewed by hand, so please allow a few business days for a reply.
Tell us what you intend to do with these weights. Requests that leave the intended-use field blank, or that answer it with a single generic phrase such as "research" or "testing", will be declined — we cannot evaluate a request we cannot read. A specific answer (what project, what data, what you expect to produce) is normally approved.
By submitting this form you confirm that you will not use these weights for clinical decision-making, that you will not redistribute them, and that you will cite the paper in any work derived from them.
Log in or Sign Up to review the conditions and access this model content.
DGR — Distortion-Guided Restoration for Prostate Diffusion MRI
Model weights for
Distortion-guided restoration: a physics-informed learning framework to correct prostate diffusion MRI artifacts Ziyang Long, Nader Binesh, Lixia Wang, Archana Vadiraj Malaji, Chia-Chi Yang, Haoran Sun, Rola Saouaf, Timothy Daskivich, Hyung Kim, Yibin Xie, Debiao Li, Hsin-Jung Yang. Radiology Advances 3(4), 2026. doi:10.1093/radadv/umag031
Code: https://github.com/Albertlongzi/DGR
DGR corrects susceptibility-induced geometric distortion in single-shot EPI prostate DWI without a field map and without a reverse phase-encode acquisition. It runs in two stages: a CNN front-end produces the geometric correction, then a conditional diffusion module refines the result on the anatomical manifold using the co-registered T2-weighted image as guidance.
Why access is gated
These weights were trained on clinical prostate MRI. Gating is not about the file size — it lets us keep a record of who is using the model and for what, so that (a) misuse in a clinical or diagnostic setting can be discouraged in advance, and (b) we can reach users if a problem with the weights is later found. Requests are approved manually. Please state a concrete intended use in the request form; that field is what the review actually reads.
Files
| Path | Stage | Params | Size | Source epoch | Val loss |
|---|---|---|---|---|---|
stage1_cnn/stage1_cnn.safetensors |
1 — CNN front-end (PHCE2EMageUltraNet) |
32.1 M | 128 MB | 25 | 0.027988 |
stage2_diffusion/stage2_diffusion.safetensors |
2 — conditional diffusion (DiffusionUNetT2AndCNN) |
299.6 M | 1.20 GB | 92 | 0.027970 |
Each directory also holds a config.json recording the architecture hyperparameters, so the
network can be rebuilt without guessing. Optimizer state has been stripped; these are
inference-only weights.
SHA-256:
7b23def2e2d6e9068ca09afd60e4b5dad50db8b5564c67c74e5774772964898f stage1_cnn/stage1_cnn.safetensors
2992be31efe1f375b4ce087036d5c0dd3476b263d206df6a04a701b2b93586d3 stage2_diffusion/stage2_diffusion.safetensors
Usage
Install the code first — these weights are not loadable by transformers or diffusers alone;
they need the dgr package for the architecture definitions.
git clone https://github.com/Albertlongzi/DGR
cd DGR
pip install -r requirements.txt && pip install -e .
hf auth login # required: the repository is gated
Download:
from huggingface_hub import hf_hub_download
cnn_path = hf_hub_download("Zylong/DGR", "stage1_cnn/stage1_cnn.safetensors")
diff_path = hf_hub_download("Zylong/DGR", "stage2_diffusion/stage2_diffusion.safetensors")
Load:
import json
import torch
from safetensors.torch import load_file
from dgr.models.phc_e2e_mageultra_net import PHCE2EMageUltraNet
from dgr.models.diffusion_unet_diffusers import DiffusionUNetT2AndCNN
cfg = json.load(open("stage1_cnn/config.json"))["arch"]
n_slices = 2 * cfg["radius"] + 1 # 2.5-D stack depth, 5 for radius=2
cnn = PHCE2EMageUltraNet(
dwi_channels=n_slices,
t2_channels=n_slices,
base_channels=cfg["base_channels"],
latent_dim=cfg["latent_dim"],
prompt_k=cfg["prompt_k"],
prompt_temp=cfg["prompt_temp"],
)
cnn.load_state_dict(load_file(cnn_path))
cnn.eval()
diff_cfg = json.load(open("stage2_diffusion/config.json"))
unet = DiffusionUNetT2AndCNN(fusion_channels=diff_cfg["arch"]["t2_cond_channels"])
unet.load_state_dict(load_file(diff_path))
unet.eval()
# diff_cfg["noise_scheduler"] carries prediction_type="sample" — the scheduler MUST be
# built with it, since the diffusers default ("epsilon") silently produces noise.
End to end, the simplest route is the inference script in the code repository:
python scripts/restoration/infer_dgr.py \
--cnn_ckpt stage1_cnn/stage1_cnn.safetensors \
--ckpt stage2_diffusion/stage2_diffusion.safetensors \
--test_root /path/to/preprocessed_test_npz \
--out_dir outputs/dgr_infer \
--steps 100 --strength 0.3 --eta 0.0 --sampler dpmsolver \
--radius 2 --t2_cond_channels 64 --b_low 50 --b_high 1400 \
--slice_mode all --save_npz --save_slices
Inputs and outputs
Input — per-subject NPZ holding a distorted single-shot EPI DWI volume at two b-values
(b = 50 and b = 1400 s/mm²) plus a co-registered T2-weighted volume. The network consumes 2.5-D
slice stacks of 2 * radius + 1 = 5 slices.
Output — the geometry-corrected DWI at both b-values, from which ADC is derived.
Settings behind these weights
| Stage 1 | Stage 2 | |
|---|---|---|
| Epochs | 25 | 100 |
| Batch size (per rank) | 6 | 6 |
| LR / warmup | 3e-4 / 3054 steps, sqrt autoscale | 1e-4 / 1000 steps |
| Hardware | 6 GPUs (torchrun) | 4 × H100 (torchrun) |
| Slice radius | 2 | 2 |
| Conditioning | — | T2w, 64 channels, t2_contrast_mod=none; frozen stage-1 output |
| Objective | L1 + SSIM (w 0.25) + multi-scale (0.2 / 0.05) | MSE on the clean image, prediction_type="sample" |
| Validation sampler | — | DPM-Solver, 50 steps, strength 0.3 |
Inference: DPM-Solver, 100 steps, eta 0, SDEdit strength in the 0.1–0.3 range. 0.1 is the
conservative setting reported in the paper; 0.3 is the default in infer_dgr.py and the value
every final inference run used. Both refine rather than redraw — above that range the diffusion
prior starts inventing structure.
Training data
Trained on distorted/undistorted DWI pairs produced by the forward physics simulator shipped in the code repository. A measured ΔB0 field is fitted per slice with a 12th-order 2-D polynomial; its high-order coefficients are then scaled by factors of 0–0.15, expanding 11 field maps measured in patients with unilateral or bilateral hip prostheses into 110 physically plausible variants. Each variant is converted to a voxel displacement map from the EPI readout geometry and applied to undistorted DWI with a conservative forward splat. (A spherical-harmonic generator ships alongside as an alternative; it was not used for these pairs.)
Source volumes: 442 retrospective MR examinations — 312 from the public fastMRI Prostate dataset and 130 in-house (Cedars-Sinai Medical Center) — of which 408 were used for training and 34 held out as the clinical test set (5 fastMRI, 29 Cedars-Sinai).
The clinical source data cannot be redistributed. The simulation half of the code repository lets the pairs be regenerated from any DWI + T2w + ΔB0 source.
Limitations
- Trained on prostate DWI at b = 50 / 1400 s/mm² on 3 T scanners. Other anatomy, other b-values, other field strengths, or markedly different EPI readout geometry are out of distribution and were not evaluated.
- Requires a co-registered T2-weighted image. Quality degrades when the T2w-to-DWI registration is poor, because the anatomical guidance is then wrong rather than merely absent.
- The stage-2 module is generative. Within the intended
strengthrange of 0.1–0.3 it refines; at higher strength it can hallucinate plausible-looking anatomy. Do not raise it without checking outputs against a reference. - Severe motion, large air pockets (post-biopsy, endorectal gas) beyond the range of the simulated B0 variants, and metal implants were not specifically modelled.
- Not a medical device. Not cleared or approved for clinical use, diagnosis, or treatment planning by any regulatory body.
Evaluation
Evaluated against FUGUE and TOPUP baselines on a held-out test set with PSNR, SSIM, NMSE and MAE
(whole-FOV and prostate-centred), plus a blinded radiologist reader study. See the paper for the
full results, and scripts/evaluation/evaluate_distortion_correction.py in the code repository
for the harness.
License
Research use only, non-commercial. See LICENSE in this repository. The accompanying code at
https://github.com/Albertlongzi/DGR is released separately under the MIT license.
Citation
@article{long2026dgr,
title = {Distortion-guided restoration: a physics-informed learning framework
to correct prostate diffusion MRI artifacts},
author = {Long, Ziyang and Binesh, Nader and Wang, Lixia and Malaji, Archana Vadiraj
and Yang, Chia-Chi and Sun, Haoran and Saouaf, Rola and Daskivich, Timothy
and Kim, Hyung and Xie, Yibin and Li, Debiao and Yang, Hsin-Jung},
journal = {Radiology Advances},
volume = {3},
number = {4},
year = {2026},
doi = {10.1093/radadv/umag031}
}
Contact
Questions about access requests or the weights: open an issue at https://github.com/Albertlongzi/DGR/issues