Recurrent Variational Network: A Deep Learning Inverse Problem Solver applied to the task of Accelerated MRI Reconstruction (CVPR 2022)

This folder contains the training code specific for reproduction of our experiments as presented in our paper Recurrent Variational Network: A Deep Learning Inverse Problem Solver applied to the task of Accelerated MRI Reconstruction (pre-print version) accepted in CVPR 2022.

Datasets

• For the proposed model, the comparison, and ablation studies we used the Calgary-Campinas public brain multi-coil MRI dataset which was released as part of an accelerated MRI reconstruction challenge. The dataset is consisted of 67 3D raw k-space volumes. After cropping the 100 outer slices, these amount to 10,452 slices of fully sampled k-spaces which we randomly split into training (47 volumes), validation (10 volumes) and test (10 volumes) sets (see lists/). Sub-sampling was performed by applying the Poisson disk distribution sub-sampling masks provided by the challange.

• For additional experiments we used the AXT1 brain FastMRI dataset. The dataset was consisted of 3D raw k-space volumes:

  • Training Set: 248 volumes (3844 slices)

  • Validation Set: 92 volumes (1428 slices) split in half to create new Validation and Test sets.

  For this dataset random Cartesian sub-sampling was performed.

Training

After downloading the data to <data_root>` the standard training command ``direct train can be used for training. Configurations can be found in the project folder.

To train our proposed model on the Calgary Campinas Dataset:

direct train <output_folder> \
            --training-root <training_data_directory> \
            --validation-root <validation_data_directory>  \
            --cfg projects/cvpr2022_recurrentvarnet/calgary_campinas/configs/base_recurrentvarnet.yaml \
            --num-gpus <number_of_gpus> \
            --num-workers <number_of_workers> \

To train a model used for the comparison or ablation studies in the paper (Section 4) a command such as the one below is used:

direct train <output_folder> \
            --training-root <training_data_directory> \
            --validation-root <validation_data_directory>  \
            --cfg projects/cvpr2022_recurrentvarnet/calgary_campinas/configs/<ablation_or_comparisons>/base_<model_name>.yaml \
            --num-gpus <number_of_gpus> \
            --num-workers <number_of_workers> \

To train a model used for the additional experiments on the FastMRI AXT1 brain Dataset as in the paper (Appendix B) a command such as the one below is used:

direct train <output_folder> \
            --training-root <training_data_directory> \
            --validation-root <validation_data_directory>  \
            --cfg projects/cvpr2022_recurrentvarnet/fastmri/AXT1_brain/configs/base_<model_name>.yaml \
            --num-gpus <number_of_gpus> \
            --num-workers <number_of_workers> \

For further information about training see Training.

During training, training loss, validation metrics and validation image predictions are logged. Additionally, Tensorboard allows for visualization of the above.

Inference

Validation

To perform inference on the Calgary Campinas validation set run:

python3 predict_val.py <output_directory> \
            --checkpoint <path_or_url_to_checkpoint> \
            --cfg projects/cvpr2022_recurrentvarnet/calgary_campinas/configs/<...>/base_<model_name>.yaml \
            --data-root <data_directory_containing_val_samples_as_in `.../calgary_campinas/lists/val/val.lst`>
            --filenames-filter projects/cvpr2022_recurrentvarnet/calgary_campinas/lists/val/val.lst
            --num-gpus <number_of_gpus> \
            --num-workers <number_of_workers> \
            --validation-index <validation_set_index> \
            [--other-flags]

To perform inference on the fastMRI AXT1 validation set run:

python3 predict_val.py <output_directory> \
            --checkpoint <path_or_url_to_checkpoint> \
            --cfg projects/cvpr2022_recurrentvarnet/fastmri/AXT1_brain/configs/base_<model_name>.yaml \
            --data-root <data_directory_containing_val_samples_as_in `.../fastmri/AXT1_brain/lists/val.lst`>
            --filenames-filter projects/cvpr2022_recurrentvarnet/fastmri/AXT1_brain/lists/val.lst
            --num-gpus <number_of_gpus> \
            --num-workers <number_of_workers> \
            --validation-index <validation_set_index> \
            [--other-flags]

Test

To perform inference on the Calgary Campinas test set run:

direct predict <output_directory>
            --checkpoint <path_or_url_to_checkpoint> \
            --cfg projects/cvpr2022_recurrentvarnet/calgary_campinas/configs_inference/<...>/base_<model_name>.yaml \
            --data-root <data_directory_containing_test_samples_as_in_`.../calgary_campinas/lists/test/test.lst`>
            --filenames-filter projects/cvpr2022_recurrentvarnet/calgary_campinas/lists/test/test.lst
            --num-gpus <number_of_gpus> \
            --num-workers <number_of_workers> \
            [--other-flags]

To perform inference on the fastMRI AXT1 test set run:

direct predict <output_directory>
            --checkpoint <path_or_url_to_checkpoint> \
            --cfg projects/cvpr2022_recurrentvarnet/fastmri/AXT1_brain/configs_inference/<R>x/base_<model_name>.yaml \
            --data-root <data_directory_containing_test_samples_as_in_`.../AXT1_brain/lists/test.lst`>
            --filenames-filter projects/cvpr2022_recurrentvarnet/fastmri/AXT1_brain/lists/test.lst
            --num-gpus <number_of_gpus> \
            --num-workers <number_of_workers> \
            [--other-flags]