{"id":27086975,"url":"https://github.com/icon-lab/mambaroll","last_synced_at":"2025-04-06T05:49:05.151Z","repository":{"id":270746229,"uuid":"904591151","full_name":"icon-lab/MambaRoll","owner":"icon-lab","description":"Official implementation of MambaRoll: A Physics-Driven Autoregressive State Space Model for Medical Image Reconstruction (https://arxiv.org/abs/2412.09331)","archived":false,"fork":false,"pushed_at":"2025-02-14T21:39:17.000Z","size":321,"stargazers_count":14,"open_issues_count":2,"forks_count":1,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-02-14T22:28:55.730Z","etag":null,"topics":["autoregressive","deep-learning","deep-neural-networks","image-reconstruction","mamba","medical-imaging","next-scale-prediction","python","pytorch","sequence-models","state-space-models","unrolled-reconstruction-algorithm"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/icon-lab.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-12-17T07:28:32.000Z","updated_at":"2025-02-14T21:39:20.000Z","dependencies_parsed_at":"2025-01-02T20:29:07.641Z","dependency_job_id":"1af6eeb0-7897-4876-8006-f7c354bc7af4","html_url":"https://github.com/icon-lab/MambaRoll","commit_stats":null,"previous_names":["icon-lab/mambaroll"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/icon-lab%2FMambaRoll","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/icon-lab%2FMambaRoll/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/icon-lab%2FMambaRoll/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/icon-lab%2FMambaRoll/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/icon-lab","download_url":"https://codeload.github.com/icon-lab/MambaRoll/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247441007,"owners_count":20939236,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["autoregressive","deep-learning","deep-neural-networks","image-reconstruction","mamba","medical-imaging","next-scale-prediction","python","pytorch","sequence-models","state-space-models","unrolled-reconstruction-algorithm"],"created_at":"2025-04-06T05:49:04.700Z","updated_at":"2025-04-06T05:49:05.134Z","avatar_url":"https://github.com/icon-lab.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003chr\u003e\n\u003ch1 align=\"center\"\u003e\n  MambaRoll \u003cbr\u003e\n  \u003csub\u003ePhysics-Driven Autoregressive State Space Models for Medical Image Reconstruction\u003c/sub\u003e\n\u003c/h1\u003e\n\n\u003cdiv align=\"center\"\u003e\n  \u003ca href=\"https://bilalkabas.github.io/\" target=\"_blank\"\u003eBilal\u0026nbsp;Kabas\u003c/a\u003e\u003csup\u003e1,2\u003c/sup\u003e \u0026ensp; \u003cb\u003e\u0026middot;\u003c/b\u003e \u0026ensp;\n  \u003ca href=\"https://github.com/fuat-arslan\" target=\"_blank\"\u003eFuat\u0026nbsp;Arslan\u003c/a\u003e\u003csup\u003e1,2\u003c/sup\u003e \u0026ensp; \u003cb\u003e\u0026middot;\u003c/b\u003e \u0026ensp;\n  \u003ca href=\"https://github.com/Valiyeh\" target=\"_blank\"\u003eValiyeh\u0026nbsp;A. Nezhad\u003c/a\u003e\u003csup\u003e1,2\u003c/sup\u003e \u0026ensp; \u003cb\u003e\u0026middot;\u003c/b\u003e \u0026ensp;\n  \u003ca href=\"https://scholar.google.com/citations?hl=en\u0026user=_SujLxcAAAAJ\" target=\"_blank\"\u003eSaban\u0026nbsp;Ozturk\u003c/a\u003e\u003csup\u003e1,2\u003c/sup\u003e \u0026ensp; \u003cb\u003e\u0026middot;\u003c/b\u003e \u0026ensp;\n  \u003ca href=\"https://kilyos.ee.bilkent.edu.tr/~saritas/\" target=\"_blank\"\u003eEmine U.\u0026nbsp;Saritas\u003c/a\u003e\u003csup\u003e1,2\u003c/sup\u003e \u0026ensp; \u003cb\u003e\u0026middot;\u003c/b\u003e \u0026ensp;\n  \u003ca href=\"https://kilyos.ee.bilkent.edu.tr/~cukur/\" target=\"_blank\"\u003eTolga\u0026nbsp;Çukur\u003c/a\u003e\u003csup\u003e1,2\u003c/sup\u003e \u0026ensp;\n  \n  \u003cspan\u003e\u003c/span\u003e\n  \n  \u003csup\u003e1\u003c/sup\u003eBilkent University \u0026emsp; \u003csup\u003e2\u003c/sup\u003eUMRAM \u003cbr\u003e\n\u003c/div\u003e\n\u003chr\u003e\n\n\u003ch3 align=\"center\"\u003e[\u003ca href=\"https://arxiv.org/abs/2412.09331\"\u003earXiv\u003c/a\u003e]\u003c/h3\u003e\n\nOfficial PyTorch implementation of **MambaRoll**, a novel physics-driven autoregressive state space model for enhanced fidelity in medical image reconstruction. In each cascade of an unrolled architecture, MambaRoll employs an autoregressive framework based on physics-driven state space modules (PSSM), where PSSMs efficiently aggregate contextual features at a given spatial scale while maintaining fidelity to acquired data, and autoregressive prediction of next-scale feature maps from earlier spatial scales enhance capture of multi-scale contextual features\n\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"figures/architecture.png\" alt=\"architecture\"\u003e\n\u003c/p\u003e\n\n## ⚙️ Installation\n\nThis repository has been developed and tested with `CUDA 12.2` and `Python 3.12`. Below commands create a conda environment with required packages. Make sure conda is installed.\n\n```\nconda env create --file requirements.yaml\nconda activate mambaroll\n```\n\n\u003cdetails\u003e\n\u003csummary\u003e[Optional] Setting Up Faster and Memory-efficient Radon Transform\u003c/summary\u003e\u003cbr\u003e\n\nWe use a faster (over 100x) and memory-efficient (~4.5x) implementation of Radon transform ([torch-radon](https://github.com/matteo-ronchetti/torch-radon)). To install, run commands below within `mambaroll` conda environment.\n\n```\ngit clone https://github.com/matteo-ronchetti/torch-radon.git\ncd torch-radon\npython setup.py install\n```\n\n\u003c/details\u003e\n\n## 🗂️ Prepare dataset\n\nMambaRoll supports reconstructions for MRI and CT modalities. Therefore, we have two dataset classes: (1) `MRIDataset` and (2) `CTDataset` in `datasets.py`.\n\n### 1. MRI dataset folder structure\n\nMRI dataset has subfolders for each undersampling rate, e.g. us4x, us8x, etc. There is a separate `.npz` file for each contrast.\n\n\u003cdetails\u003e\n\u003csummary\u003eDetails for npz files\u003c/summary\u003e\u003cbr\u003e\n  \nA `\u003ccontrast\u003e.npz` file has the following keys:\n\n| Variable key    | Description                               | Shape                                     | \n|-----------------|-------------------------------------------|-------------------------------------------|\n| `image_fs`      | Coil-combined fully-sampled MR image.     | n_slices x 1 x height x width             |\n| `image_us`      | Multi-coil undersampled MR image.         | n_slices x n_coils x height x width       |\n| `us_masks`      | K-space undersampling masks.              | n_slices x 1 x height x width             |\n| `coilmaps`      | Coil sensitivity maps.                    | n_slices x n_coils x height x width       |\n| `subject_ids`   | Corresponding subject ID for each slice.  | n_slices                                  |\n| `us_factor`     | Undersampling factor.                     | (Single integer value)                    |\n\n\u003c/details\u003e\n\n```\nfastMRI/\n├── us4x/\n│   ├── train/\n│   │   ├── T1.npz\n│   │   ├── T2.npz\n│   │   └── FLAIR.npz\n│   ├── test/\n│   │   ├── T1.npz\n│   │   ├── T2.npz\n│   │   └── FLAIR.npz\n│   └── val/\n│       ├── T1.npz\n│       ├── T2.npz\n│       └── FLAIR.npz\n├── us8x/\n│   ├── train/...\n│   ├── test/...\n│   └── val/...\n├── ...\n```\n\n\n\n### 2. CT dataset folder structure\n\nEach split in CT dataset contains images with different undersampling rates.\n\n\u003cdetails\u003e\n\u003csummary\u003eDetails for npz files\u003c/summary\u003e\u003cbr\u003e\n\n`image_fs.npz` files have the fully-sampled data with the following key:\n\n| Variable key    | Description                               | Shape                                     |\n|-----------------|-------------------------------------------|-------------------------------------------|\n| `image_fs`      | Fully-sampled CT image.                   | n_slices x 1 x height x width             |\n\nA `us\u003cus_factor\u003ex.npz` file has the following keys:\n\n| Variable key          | Description                                                                                                                | Shape                                             |\n|-----------------------|----------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------|\n| `image_us`            | Undersampled CT image.                                                                                                     | n_slices x 1 x height x width                     |\n| `sinogram_us`         | Corresponding sinograms for undersampled CTs.                                                                              | n_slices x 1 x detector_positions x n_projections |\n| `projection_angles`   | Projection angles in degrees at which the Radon transform performed on fully-sampled images to obtain undersampled ones.   | n_slices x n_projections                          |\n| `subject_ids`         | Corresponding subject ID for each slice.                                                                                   | n_slices                                          |\n| `us_factor`           | Undersampling factor.                                                                                                      | (Single integer value)                            |\n\n\u003c/details\u003e\n\n```\nlodopab-ct/\n├── train/\n│   ├── image_fs.npz\n│   ├── us4x.npz\n│   └── us6x.npz\n├── test/\n│   ├── image_fs.npz\n│   ├── us4x.npz\n│   └── us6x.npz\n└── val/\n    ├── image_fs.npz\n    ├── us4x.npz\n    └── us6x.npz\n```\n\n\n\n## 🏃 Training\n\nRun the following command to start/resume training. Model checkpoints are saved under `logs/$EXP_NAME/MambaRoll/checkpoints` directory, and sample validation images are saved under `logs/$EXP_NAME/MambaRoll/val_samples`. The script supports both single and multi-GPU training. By default, it runs on a single GPU. To enable multi-GPU training, set `--trainer.devices` argument to the list of devices, e.g. `0,1,2,3`. Be aware that multi-GPU training may lead to convergence issues. Therefore, it is only recommended during inference/testing.\n\n```\npython main.py fit \\\n    --config $CONFIG_PATH \\\n    --trainer.logger.name $EXP_NAME \\\n    --model.mode $MODE \\\n    --data.dataset_dir $DATA_DIR \\\n    --data.contrast $CONTRAST \\\n    --data.us_factor $US_FACTOR \\\n    --data.train_batch_size $BS_TRAIN \\\n    --data.val_batch_size $BS_VAL \\\n    [--trainer.max_epoch $N_EPOCHS] \\\n    [--ckpt_path $CKPT_PATH] \\\n    [--trainer.devices $DEVICES]\n\n```\n\n\u003cdetails\u003e\n\u003csummary\u003eExample Commands\u003c/summary\u003e\n\nMRI reconstruction using fastMRI dataset:\n\n```\npython main.py fit \\\n  --config configs/config_fastmri.yaml \\\n  --trainer.logger.name fastmri_t1_us8x \\\n  --data.dataset_dir ../datasets/fastMRI \\\n  --data.contrast T1 \\\n  --data.us_factor 8 \\\n  --data.train_batch_size 1 \\\n  --data.val_batch_size 16 \\\n  --trainer.devices [0]\n```\n\nCT reconstruction using [LoDoPaB-CT](https://zenodo.org/records/3384092) dataset:\n\n```\npython main.py fit \\\n  --config configs/config_ct.yaml \\\n  --trainer.logger.name ct_us4x \\\n  --data.dataset_dir ../datasets/lodopab-ct/ \\\n  --data.us_factor 4 \\\n  --data.train_batch_size 1 \\\n  --data.val_batch_size 16 \\\n  --trainer.devices [0]\n```\n\u003c/details\u003e\n\n### Argument descriptions\n\n| Argument                    | Description                                                                                                                    |\n|-----------------------------|--------------------------------------------------------------------------------------------------------------------------------|\n| `--config`                  | Config file path. Available config files: 'configs/config_mri.yaml' and 'configs/config_ct.yaml'                           |\n| `--trainer.logger.name`     | Experiment name.                                                                                                               |\n| `--model.mode`              | Mode depending on data modality. Options: 'mri', 'ct'.                                                                      |\n| `--data.dataset_dir`        | Data set directory.                                                                                                            |\n| `--data.contrast`           | Source contrast, e.g. 'T1', 'T2', ... for MRI. Should match the folder name for that contrast.                                 |\n| `--data.us_factor`          | Undersampling factor, e.g 4, 8.                                                                                                |\n| `--data.train_batch_size`   | Train set batch size.                                                                                                          |\n| `--data.val_batch_size`     | Validation set batch size.                                                                                                     |\n| `--trainer.max_epoch`       | [Optional] Number of training epochs (default: 50).                                                                            |\n| `--ckpt_path`               | [Optional] Model checkpoint path to resume training.                                                                           |\n| `--trainer.devices`         | [Optional] Device or list of devices. For multi-GPU set to the list of device ids, e.g `0,1,2,3` (default: `[0]`).             |\n\n\n## 🧪 Testing\n\nRun the following command to start testing. The predicted images are saved under `logs/$EXP_NAME/MambaRoll/test_samples` directory. By default, the script runs on a single GPU. To enable multi-GPU testing, set `--trainer.devices` argument to the list of devices, e.g. `0,1,2,3`.\n\n```\npython main.py test \\\n    --config $CONFIG_PATH \\\n    --model.mode $MODE \\\n    --data.dataset_dir $DATA_DIR \\\n    --data.contrast $CONTRAST \\\n    --data.us_factor $US_FACTOR \\\n    --data.test_batch_size $BS_TEST \\\n    --ckpt_path $CKPT_PATH\n```\n\n### Argument descriptions\n\nSome arguments are common to both training and testing and are not listed here. For details on those arguments, please refer to the training section.\n\n| Argument                    | Description                                |\n|-----------------------------|--------------------------------------------|\n| `--data.test_batch_size`    | Test set batch size.                       |\n| `--ckpt_path`               | Model checkpoint path.                     |\n\n\n## ✒️ Citation\nYou are encouraged to modify/distribute this code. However, please acknowledge this code and cite the paper appropriately.\n```\n@article{kabas2024mambaroll,\n  title={Physics-Driven Autoregressive State Space Models for Medical Image Reconstruction}, \n  author={Bilal Kabas and Fuat Arslan and Valiyeh A. Nezhad and Saban Ozturk and Emine U. Saritas and Tolga Çukur},\n  year={2024},\n  journal={arXiv:2412.09331}\n}\n```\n\n\n### 💡 Acknowledgments\n\nThis repository uses code from the following projects:\n\n- [mamba](https://github.com/state-spaces/mamba)\n- [deepinv](https://github.com/deepinv/deepinv)\n\n\u003chr\u003e\nCopyright © 2024, ICON Lab.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ficon-lab%2Fmambaroll","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ficon-lab%2Fmambaroll","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ficon-lab%2Fmambaroll/lists"}