This repository offers the Adaptive Sampling of k-space in Magnetic Resonance for Rapid Pathology Prediction (ASMR), introduced in [ICML24 Paper]. [Citation]

1. Clone the repository

2. Install the dependencies

conda create -n "asmr" python=3.10 pip install -r requirements.txt pip install -e . 

The dataset can be downloaded from the fastMRI website. Users need to provide a csv file that includes the pathology information. Once the data is ready, set the datadir and split_csv_file according to the data path. Tools for processing the data are in data_modules.

We provide sample running scripts in the scripts/ directory.

To enable the use of ASMR's environment, we first need to train a reward model. The configuration and hyperparamters for training a reward model is located in the configs/ directory. Set the datadir and split_csv_file in the YAML file to your data path.

Trained the reward model:

bash train_reward_model.sh 

Once the reward model is ready, we can use the ASMR environment for reinforcement learning. The configuration for the ASMR environment is located in the rl/cfgs/env/ directory. Set the reward_model_ckpt in the YAML file to the reward model checkpoint from the previous step.

1. Generate the weighted sampler for training,

bash generate_sampler.sh 

and set the train_sampler_filename in the rl/cfgs/env/*.yaml accordingly.

2. Train the ASMR policy

bash train_asmr.sh 

To evaluate the trained policy

cd rl python eval_asmr.py load_from_snapshot_base_dir=<policy-checkpoint-directory> eval_range=[<start>,<end>] 

If you use this repo in you research, please cite the paper as follows

@misc{yen2024adaptive, title={Adaptive Sampling of k-Space in Magnetic Resonance for Rapid Pathology Prediction}, author={Chen-Yu Yen and Raghav Singhal and Umang Sharma and Rajesh Ranganath and Sumit Chopra and Lerrel Pinto}, year={2024}, eprint={2406.04318}, archivePrefix={arXiv}, primaryClass={cs.LG} }