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}
}