| [](https://arxiv.org/abs/2410.16290) | |
| [](https://armeet.ca/nomri) | |
| # A Unified Model for Compressed Sensing MRI Across Undersampling Patterns | |
| > [**A Unified Model for Compressed Sensing MRI Across Undersampling Patterns**](https://arxiv.org/abs/2410.16290) | |
| > Armeet Singh Jatyani, Jiayun Wang, Aditi Chandrashekar, Zihui Wu, Miguel Liu-Schiaffini, Bahareh Tolooshams, Anima Anandkumar | |
| > *Paper at [CVPR 2025](https://cvpr.thecvf.com/Conferences/2025/AcceptedPapers)* | |
|  | |
| > _**(a) Unified Model:** NO works across various undersampling patterns, unlike CNNs (e.g., [E2E-VarNet](#)) that need separate models for each._ \ | |
| > _**(b) Consistent Performance:** NO consistently outperforms CNNs, especially for 2× acceleration with a single unrolled cascade._ \ | |
| > _**(c) Resolution-Agnostic:** Maintains fixed kernel size regardless of image resolution, reducing aliasing risks._ \ | |
| > _**(d) Zero-Shot Super-Resolution:** Outperforms CNNs in reconstructing high-res MRIs without retraining._ | |
|  | |
| > _**(a) Zero-Shot Extended FOV:** On 4x Gaussian undersampling, NO achieves higher PSNR and fewer artifacts than E2E-VN, despite both models being trained only on 160 x 160 FOV._ \ | |
| > _**(b) Zero-Shot Super-Resolution in Image Space:** For 2x radial with 640 x 640 input via bilinear upsampling, NO preserves quality while E2E-VN introduces artifacts._ | |
| ## Requirements | |
| We have tested training/inference on the following hardware/software versions, however there is no reason it shouldn't work on slightly older driver/cuda versions. | |
| - tested on RTX 4090 and A100 with CUDA 12.4 and NVML/Driver version 550 | |
| - Ubuntu 22.04.3 LTS & SUSE Linux Enterprise Server 15 | |
| - All python packages are in `pyproject.toml` (see Setup) | |
| ## Setup | |
| We use `uv` for environment setup. It is 10-100x faster than vanilla pip and conda. If you don't have `uv`, please install it from [here](https://docs.astral.sh/uv/getting-started/installation/) (no sudo required). If you're on a Linux environment you can install with: `curl -LsSf https://astral.sh/uv/install.sh | sh`. Of course, if you would like to use a virtual environment handled by vanilla python or conda, all package and their versions are provided in `pyproject.toml` under "dependencies." | |
| In the root directory, run | |
| ```bash | |
| uv sync | |
| ``` | |
| Then you can activate the environment with: | |
| ```bash | |
| source .venv/bin/activate | |
| ``` | |
| Note this is optional. You can run scripts with this venv without activating the environment by using `uv run python script.py` or abbreviated `uv run script.py`. | |
| `uv` will create a virtual environment for you and install all packages. | |
| Then to download the pretrained weights, run: | |
| ```bash | |
| uv run scripts/download_weights.py | |
| ``` | |
| This downloads pretrained weights into the `weights/` directory. | |
| Finally to run scripts, make them executable: | |
| ```bash | |
| chmod u+x scripts/* | |
| ``` | |
| Then you can run any script. For example: | |
| ```bash | |
| ./scripts/knee_multipatt.sh | |
| ``` | |
| By default weights & biases (WANDB) is disabled, so scripts will print results to stdout. If you want to visualize results in | |
| weights and biases, add your WANDB api key at the top of the script. We log image predictions | |
| as well as PSNR, NMSE, SSIM metrics for each epoch. | |
| ```bash | |
| export WANDB_API_KEY=*************** | |
| ``` | |
| Before you can begin training/inference, you will need to download and process the dataset. See the "Datasets" section below. | |
| ## Datasets | |
| We use the fastMRI dataset, which can be downloaded [here](https://fastmri.med.nyu.edu/). \ | |
| Dataset classes are provided in `fastmri/datasets.py`: | |
| - `SliceDatasetLMDB`: dataset in significantly faster LMDB format | |
| - `SliceDataset`: dataset class for original fastMRI dataset | |
| We convert the raw fastMRI HDF5 formatted samples into a significantly faster LMDB format. | |
| This accelerates training/validation by a significant factor. Once you have downloaded the fastMRI dataset, | |
| you will need to run `scripts/gen_lmdb_dataset.py` to convert the original fastMRI dataset into LMDB format. | |
| ```bash | |
| uv run scripts/gen_lmdb_dataset.py --body_part brain --partition val -o /path/to/lmdb/dataset | |
| ``` | |
| Do this for every dataset you need: (brain, knee) x (train, val). To choose a smaller subset for faster training/inference add `--sample_rate 0.Xx`. | |
| By default we use the LMDB format. If you want to use the original SliceDataset class, you can swap out the dataset class in `main.py`. | |
| Finally, modify your `fastmri.yaml` with the correct dataset paths | |
| ```yaml | |
| log_path: /tmp/logs | |
| checkpoint_path: /tmp/checkpoints | |
| lmdb: | |
| knee_train_path: **/**/**/knee_train_lmdb | |
| knee_val_path: **/**/**/knee_val_lmdb | |
| brain_train_path: **/**/**/brain_train_lmdb | |
| brain_val_path: **/**/**/brain_val_lmdb | |
| ``` | |
| ## Training and Validation | |
| `main.py` is used for both training and validation. We follow the original fastMRI repo | |
| and use Lightning. We provide both a simple PyTorch model `models/no_varnet.py` (if you want | |
| a thinner abstraction), as well as a Lightning wrapped `models/lightning/no_varnet_module.py` that | |
| makes distributed training across multiple GPUs easier. | |
| ## Citation | |
| If you found our work helpful or used any of our models (UDNO), please cite the following: | |
| ```bibtex | |
| @inproceedings{jatyani2025nomri, | |
| author = {Armeet Singh Jatyani* and Jiayun Wang* and Aditi Chandrashekar and Zihui Wu and Miguel Liu-Schiaffini and Bahareh Tolooshams and Anima Anandkumar}, | |
| title = {A Unified Model for Compressed Sensing MRI Across Undersampling Patterns}, | |
| booktitle = {Conference on Computer Vision and Pattern Recognition (CVPR) Proceedings}, | |
| abbr = {CVPR}, | |
| year = {2025} | |
| } | |
| ``` | |
|  | |
| https://arxiv.org/abs/2410.16290 |