Skip to content

README.md

Latest commit

 

History

History
154 lines (111 loc) · 5.02 KB

README.md

File metadata and controls

154 lines (111 loc) · 5.02 KB
Raw

EDA — MRI Bias Field Correction (BFC)

Official implementation of bias field correction for the paper:

Elucidating the Design Space of Arbitrary-Noise-Based Diffusion Models (EDA)
Xingyu Qiu, Mengying Yang, Xinghua Ma, Dong Liang, Fanding Li, Gongning Luo, Wei Wang, Kuanquan Wang, Shuo Li
CVPR 2026

BFC is Task 1 in EDA: it removes global smooth bias fields from MRI by diffusing structured low-frequency basis noise (EDA Case 1, η = 0) instead of pixel-wise Gaussian noise, and restores the image in as few as 5 deterministic sampling steps.

Acknowledgement

This module is built on top of the Denoising Diffusion Implicit Models (DDIM) PyTorch codebase. We thank Jiaming Song, Chenlin Meng, and Stefano Ermon for releasing their implementation, which provides the UNet backbone, diffusion scheduling utilities, and the overall training/sampling structure used here.

Original DDIM repository: https://github.com/ermongroup/ddim

Method Summary (Paper Settings)

Item Setting
Noise pattern Unified basis set H_{3,5} (Legendre + trigonometric functions)
EDA case Case 1 — fixed basis set, sample-independent
η 0 (maximal stochasticity in basis-weighted noise)
Forward process s(t) = 1, σ(t) = √(1 − ᾱ_t)
Training target Predict diffused noise N (DDPM/DDIM objective)
Loss L_BFC = 𝔼 ‖N_θ(x_t, t) − N‖²
Diffusion steps T 100 (linear β: 0.0001 → 0.02)
Input size 256 × 256, grayscale
Optimizer Adam, lr = 2×10⁻⁵, β = (0.9, 0.999), EMA = 0.9999
Training 500 epochs, batch size 1
Sampling 5-step deterministic Euler (Algorithm 2), η = 0
Preprocessing Log transform (multiplicative bias → additive in log domain)

Repository Layout

BC/
├── main.py                 # Entry point (train / sample)
├── configs/mydataset.yml   # HCP BFC hyperparameters
├── runners/diffusion.py    # EDA BFC training & sampling
├── functions/
│   ├── losses.py           # Noise-prediction loss
│   └── denoising.py        # Deterministic sampling steps
├── datasets/
│   └── data_loading.py     # BFCDataset (log-domain preprocessing)
└── models/diffusion.py     # UNet denoiser (from DDIM)

Requirements

This module is built on DDIM. Install Python, PyTorch, and other dependencies according to the baseline repository:

Baseline GitHub: https://github.com/ermongroup/ddim

Data Preparation

We follow the Human Connectome Project (HCP) protocol described in the paper (2,206 slices for training/validation, 1,000 for testing after preprocessing).

  1. Place preprocessed clean MRI slices for training under:

    ./data/HCP/train/

  2. Place LQ (bias-corrupted) test images under:

    ./data/HCP/test/LQ/

    Optionally provide paired GT/ and LQ/ under ./data/HCP/test/ for evaluation.

  3. Update paths in configs/mydataset.yml if needed:

    data:
    train_dir: "./data/HCP/train"
    test_dir: "./data/HCP/test"
    image_size: 256

Supported formats: .png, .jpg, .jpeg, .npy.

Training

python main.py --config mydataset.yml --doc bfc --ni

Resume from checkpoint:

python main.py --config mydataset.yml --doc bfc --ni --resume_training

Checkpoints are saved under exp/logs/bfc/ (ckpt.pth, ckptbest.pth).

Sampling (Inference)

Run 5-step deterministic restoration from LQ images (paper default):

python main.py \
  --config mydataset.yml \
  --exp exp \
  --doc bfc \
  --sample \
  --sequence \
  --timesteps 5 \
  --eta 0 \
  --ni \
  --input_dir ./data/HCP/test/LQ

Outputs are written to exp/image_samples/bfc/{num_steps}/.

Log-domain pipeline: LQ images are log-preprocessed before diffusion (same as training). Restored results are mapped back with exp(·) and normalized for saving. See runners/diffusion.py → sample_sequence() for details.

Citation

If you use this code, please cite EDA and the original DDIM work:

@inproceedings{qiu2026eda,
  title={Elucidating the Design Space of Arbitrary-Noise-Based Diffusion Models},
  author={Qiu, Xingyu and Yang, Mengying and Ma, Xinghua and Liang, Dong and Li, Fanding and Luo, Gongning and Wang, Wei and Wang, Kuanquan and Li, Shuo},
  booktitle={CVPR},
  year={2026}
}

@article{song2020denoising,
  title={Denoising Diffusion Implicit Models},
  author={Song, Jiaming and Meng, Chenlin and Ermon, Stefano},
  journal={arXiv:2010.02502},
  year={2020}
}

License

This directory retains the MIT License from the original DDIM repository. Please also respect the licenses and terms of the HCP dataset when using the data.