Chichonnade/Comparative-Analysis-3D-2.5D-MRI-Super-Resolution-EDM

MRI Super-Resolution via Elucidated Diffusion Models

Pretrained model weights for the paper: Comparative Analysis of 3D Convolutional and 2.5D Slice-Conditioned U-Net Architectures for MRI Super-Resolution via Elucidated Diffusion Models

Authors: Hendrik Chiche, Ludovic Corcos, Logan Rouge (in collaboration with GENCI)

Models

2.5D EDM Upsampler (agent_epoch_00010.pt)

• Architecture: 2D U-Net with slice conditioning, channels [64, 64, 128, 256]
• Parameters: 51.1M
• Training: 10 epochs on 59 NKI subjects (FOMO60K), AdamW lr=1e-4
• PSNR: 33.68 dB | SSIM: 0.967 (on held-out NKI test set, 2x SR)
• Inference: single-step Heun sampler (0.09s/slice on MPS)

3D EDM Upsampler (model_final.pt)

• Architecture: 3D U-Net with volumetric convolutions, channels [32, 64, 128, 256]
• Parameters: 50.7M
• Training: 10 epochs on 59 NKI subjects (FOMO60K), AdamW lr=1e-4
• PSNR: 37.77 dB | SSIM: 0.996 (on held-out NKI test set, 2x SR)
• Inference: 20-step Euler sampler with patch-based processing

Dataset

NKI cohort from FOMO60K (https://huggingface.co/datasets/FOMO-MRI/FOMO60K)

• 59 training subjects, 5 test subjects
• T1-weighted sagittal brain MRI
• 2x downsampling via block averaging (LR: 128x128, HR: 256x256)

Framework

Both models use the Elucidated Diffusion Model (EDM) framework from Karras et al. 2022, adapted from the DIAMOND game world modeling codebase.

Usage

See the repository: https://github.com/chichonnade/MRI_denoiser