Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Efficient Vision Mamba for MRI Super-Resolution

License: MIT Medical Physics arXiv

Official PyTorch implementation for:

Efficient vision mamba for MRI super-resolution via hybrid selective scanning
Mojtaba Safari, Shansong Wang, Vanessa L. Wildman, et al.
Medical Physics, 2026;53:e70508. https://doi.org/10.1002/mp.70508

This repository contains an efficient MRI super-resolution framework based on vision Mamba blocks with hybrid selective scanning. The model is designed for high-quality reconstruction of low-resolution MRI while keeping the parameter count and computational cost low.

Framework

Efficient vision Mamba MRI super-resolution framework

Repository Layout

.
├── assets/
│   └── framework.png
├── codes/
│   ├── main.py
│   ├── build.py
│   ├── datasets/
│   └── mamba/
├── environment.yml
├── LICENSE
└── README.md

Installation

Create the Conda environment:

conda env create -f environment.yml
conda activate sr-rmamba

The code uses CUDA-enabled PyTorch and mamba-ssm. If your CUDA/PyTorch stack differs from the versions in environment.yml, install the matching PyTorch build from the official PyTorch instructions first, then install the remaining dependencies.

Data Format

The dataset loader expects a directory of HDF5 files. Each file should contain:

  • high_resolution: the target high-resolution 2D image
  • low_resolution: the input low-resolution 2D image

During loading, each image is min-max normalized to [-1, 1] and returned with a single channel dimension.

Example directory layout:

data/
├── train/
├── validation/
└── test/
    ├── brain/
    └── prostate/

Training

Run commands from the codes/ directory:

cd codes

python main.py \
  --train True \
  --train_dir ../data/train \
  --valid_dir ../data/validation \
  --out_dir ../results/train/mamba/prostate

Useful options include:

  • --epochs: number of training epochs
  • --batch_size: training batch size
  • --test_batch_size: validation or inference batch size
  • --lr: learning rate
  • --resume True: resume from content.pth in --out_dir

Inference

Place a trained checkpoint named content.pth in the directory passed to --out_dir, then run:

cd codes

python main.py \
  --train False \
  --region prostate \
  --test_dir ../data/test/prostate \
  --out_dir ../results/train/mamba/prostate \
  --out_dir_test ../results/test/mamba/prostate

Inference writes predicted images, visualizations, logs, and quantitative metrics to --out_dir_test.

Citation

If this code is useful for your research, please cite:

@article{safari2026efficient,
  title={Efficient vision mamba for MRI super-resolution via hybrid selective scanning},
  author={Safari, Mojtaba and Wang, Shansong and Wildman, Vanessa L and Hu, Mingzhe and Eidex, Zach and Chang, Chih-Wei and Middlebrooks, Erik H and Qiu, Richard LJ and Patel, Pretesh and Jani, Ashesh B and others},
  journal={Medical Physics},
  volume={53},
  number={5},
  pages={e70508},
  year={2026},
  publisher={Wiley Online Library},
  doi={10.1002/mp.70508},
  url={https://doi.org/10.1002/mp.70508}
}

License

This project is released under the MIT License. See LICENSE for details.

About

PyTorch code for Efficient Vision Mamba MRI super-resolution, Medical Physics 2026.

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages