Official repository for "LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation".

• 🔥News: 2024/3/29: Received feedback on issue#6 and updated the erroneous code in the original code repository.

• News: 2024/3/17: LightM-UNet released.

• A regrettable notification: 2024/3/12 Thank you for your attention! The author has been quite busy lately (；′⌒`), so only the main code has been uploaded to this repository for now. Detailed code explanations, data, and configuration details will be completed by March 17, 2024. Thanks again for your interest! o(￣▽￣)ブ

LightM-UNet is a lightweight fusion of UNet and Mamba, boasting a mere parameter count of 1M. Through validation on both 2D and 3D real-world datasets, LightM-UNet surpasses existing state-of-the-art models. In comparison to the renowned nnU-Net and contemporaneous U-Mamba, LightM-UNet reduces the parameter count by 116X and 224X, respectively.

Requirements: CUDA ≥ 11.6

1. Create a virtual environment: conda create -n lightmunet python=3.10 -y and conda activate lightmunet
2. Install Pytorch 2.0.1: pip install torch==2.0.1 torchvision==0.15.2
3. Install Mamba: pip install causal-conv1d==1.1.1 and pip install mamba-ssm
4. Download code: git clone https://github.com/MrBlankness/LightM-UNet
5. cd LightM-UNet/lightm-unet and run pip install -e .

sanity test: Enter python command-line interface and run

import torch
import mamba_ssm

Download LiTs dataset here and Montgomery&Shenzhen dataset here, then put them into the LightM-Unet/data/nnUNet_raw folder. LightM-UNet is built on the popular nnU-Net framework. If you want to train LightM-UNet on your own dataset, please follow this guideline to prepare the dataset.

Please organize the dataset as follows:

data/
├── nnUNet_raw/
│   ├── Dataset801_LiverCT/
│   │   ├── imagesTr
│   │   │   ├── Liver_0001_0000.nii.gz
│   │   │   ├── Liver_0002_0000.nii.gz
│   │   │   ├── ...
│   │   ├── labelsTr
│   │   │   ├── Liver_0001.nii.gz
│   │   │   ├── Liver_0002.nii.gz
│   │   │   ├── ...
│   │   ├── dataset.json
│   ├── Dataset802_LungXray/
│   │   ├── imagesTr
│   │   │   ├── Lung_0001_0000.png
│   │   │   ├── Lung_0002_0000.png
│   │   │   ├── ...
│   │   ├── labelsTr
│   │   │   ├── Lung_0001.png
│   │   │   ├── Lung_0001.png
│   │   │   ├── ...
│   │   ├── dataset.json
│   ├── ...

Based on nnUNet, preprocess the data and generate the corresponding configuration files (the generated results can be found in the LightM-Unet/data/nnUNet_preprocessed folder).

nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity

• Train 2D LightM-Unet model

nnUNetv2_train DATASET_ID 2d all -tr nnUNetTrainerLightMUNet

• Train 3D LightM-Unet model

nnUNetv2_train DATASET_ID 3d_fullres all -tr nnUNetTrainerLightMUNet

• Inference 2D LightM-Unet model

nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d DATASET_ID -c 2d -tr nnUNetTrainerLightMUNet --disable_tta

• Inference 3D LightM-Unet model

nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d DATASET_ID -c 3d_fullres -tr nnUNetTrainerLightMUNet --disable_tta

If you find our work helpful, please consider citing the following papers

@misc{liao2024lightmunet,
      title={LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation}, 
      author={Weibin Liao and Yinghao Zhu and Xinyuan Wang and Chengwei Pan and Yasha Wang and Liantao Ma},
      year={2024},
      eprint={2403.05246},
      archivePrefix={arXiv},
      primaryClass={eess.IV}
}

We acknowledge all the authors of the employed public datasets, allowing the community to use these valuable resources for research purposes. We also thank the authors of nnU-Net, Mamba and U-Mamba for making their valuable code publicly available.