Skip to content

Latest commit

 

History

History
123 lines (108 loc) · 5.32 KB

README.md

File metadata and controls

123 lines (108 loc) · 5.32 KB
Raw

i-RIM for fastMRI

Official implementation of the i-RIM applied to the fastMRI dataset as described in Invert to Learn to Invert and i-RIM applied to the fastMRI challenge. Pre-trained models can be found under Releases.

See some example reconstructions here:

And some numbers:

~------ 4x ------~ ~------ 8x ------~
i-RIM single-coil NMSE PSNR SSIM NMSE PSNR SSIM
Validation 0.0342 32.43 0.751 0.0446 30.92 0.692
Test 0.0272 33.65 0.781 0.0421 30.56 0.687
Challenge n/a n/a n/a 0.031 33 0.754
i-RIM multi-coil NMSE PSNR SSIM NMSE PSNR SSIM
Validation 0.0062 38.84 0.916 0.0103 36.19 0.886
Test 0.0052 39.52 0.928 0.0093 36.53 0.887
Challenge 0.006 39 0.925 0.010 37 0.899

Installation

To use this code, please run the following commands (preferably in a virtualenv):

git clone --recurse-submodules https://github.com/pputzky/irim_fastMRI.git
cd irim_fastMRI
pip install -r requirements.txt
./install.sh

The above commands will clone this repository with all submodules. Running ./install.sh will install irim as a package in your current Python environment.

Usage

This repository includes two scripts that allow training (scripts.train_model) and running (scripts.run_model) of an i-RIM. Both scripts are derived from the train and run scripts in the fastMRi code base.

To train models as used in our fastMRI challenge submission (see i-RIM applied to the fastMRI challenge), run the following commands (make sure to set $DATA_PATH and $CHECKPOINT_PATH before) :

Single-coil

Train the model
python -m scripts.train_model \
--challenge singlecoil --batch_size 8 --n_steps 8 \
--n_hidden 64 64 64 64 64 64 64 64 64 64 64 64 \
--n_network_hidden 64 64 128 128 256 1024 1024 256 128 128 64 64 \
--dilations 1 1 2 2 4 8 8 4 2 2 1 1 \
--multiplicity 4 --parametric_output \
--loss ssim --resolution 320 --train_resolution 368 368 --lr_gamma 0.1 \
--lr 0.0001 --lr_step_size 30 --num_epochs 50 --optimizer Adam \
--num_workers 8 --report_interval 100 --data_parallel --resume \
--data-path $DATA_PATH --exp_dir $CHECKPOINT_DIR
Run the model
python -m scripts.run_model --challenge singlecoil --batch-size 8 \
--data-path $DATA_PATH --checkpoint $CHECKPOINT_DIR/best_model.pt \
--out-dir $OUTPUT_DIR --data-split val --mask-kspace
Evaluate reconstructions
python -m external.fastMRI.common.evaluate --challenge singlecoil \
--target-path $DATA_PATH/singlecoil_val/ --predictions-path $OUTPUT_DIR

Multi-coil

python -m scripts.train_model \
--challenge multicoil --batch_size 32 --n_steps 8 \
--n_hidden 96 96 96 96 96 96 96 96 96 96 96 96 \
--n_network_hidden 64 64 128 128 256 1024 1024 256 128 128 64 64 \
--dilations 1 1 2 2 4 8 8 4 2 2 1 1 \
--multiplicity 1 --parametric_output \
--loss ssim --resolution 320 --train_resolution 368 368 --lr_gamma 0.1 \
--lr 0.0001 --lr_step_size 30 --num_epochs 50 --optimizer Adam \
--num_workers 8 --report_interval 100 --data_parallel --resume \
--data-path $DATA_PATH --exp_dir $CHECKPOINT_DIR

Running and evaluating as above.

RIM baseline on single-coil

python -m scripts.train_model \
--use_rim --challenge singlecoil --batch_size 8 --n_steps 8 \
--loss ssim --resolution 320 --train_resolution 368 368 --lr_gamma 0.1 \
--lr 0.0001 --lr_step_size 30 --num_epochs 50 --optimizer Adam \
--num_workers 8 --report_interval 100 --data_parallel --resume \
--data-path $DATA_PATH --exp_dir $CHECKPOINT_DIR

Running and evaluating as above.

References

If you use this code or derivatives thereof, please cite the following works

@incollection{pputzky2019,
title = {Invert to Learn to Invert},
author = {Putzky, Patrick and Welling, Max},
booktitle = {Advances in Neural Information Processing Systems 32},
editor = {H. Wallach and H. Larochelle and A. Beygelzimer and F. d\textquotesingle Alch\'{e}-Buc and E. Fox and R. Garnett},
pages = {444--454},
year = {2019},
publisher = {Curran Associates, Inc.},
url = {http://papers.nips.cc/paper/8336-invert-to-learn-to-invert.pdf}
}
@misc{pputzky2019fastMRI,
    title={i-RIM applied to the fastMRI challenge},
    author={Patrick Putzky and Dimitrios Karkalousos and Jonas Teuwen and Nikita Miriakov and Bart Bakker and Matthan Caan and Max Welling},
    year={2019},
    eprint={1910.08952},
    archivePrefix={arXiv},
    primaryClass={eess.IV}
}