Awesome
DiffuseRecon
This codebase is modified based on Improved DDPM
Installation
Clone this repository and navigate to it in your terminal. Then run:
pip install -e .
This should install the improved_diffusion
python package that the scripts depend on.
Data Preparation and Pre-Trained Checkpoints
A pre-trained checkpoint can be downloaded via this link or link.
For FastMRI, the simplified h5 data can be downloaded by following the instructions in ReconFormer, i.e. through Link. DiffuseRecon converts it to a normalized format in scripts/data_process.py
python scripts/data_process.py
Sampling
python scripts/image_sample_complex_duo.py --model_path img_space_dual/ema_0.9999_150000.pt --data_path EVAL_PATH \
--image_size 320 --num_channels 128 --num_res_blocks 3 --learn_sigma False --dropout 0.3 --diffusion_steps 4000 \
--noise_schedule cosine --timestep_respacing 100 --save_path test/ --num_samples 1 --batch_size 5
Note that timestep_respacing indicates the initial coarse sampling steps.
Training
mpiexec -n GPU_NUMS python scripts/image_train.py --data_dir TRAIN_PATH --image_size 320 --num_channels 128\
--num_res_blocks 3 --learn_sigma False --dropout 0.3 --diffusion_steps 4000 --noise_schedule cosine --lr 1e-4 --batch_size 1\
--save_dir img_space_dual
TODO
- Upload PSNR evaluation
- Currently, the refinement step is fixed at 20 (line 592, gaussian_diffusion_duo.py); make this an adjustable input.
- Graphics.