Skip to content

(ARXIV24) This is the official code repository for "VM-UNet: Vision Mamba UNet for Medical Image Segmentation".

License

Notifications You must be signed in to change notification settings

JCruan519/VM-UNet

Repository files navigation

VM-UNet

This is the official code repository for "VM-UNet: Vision Mamba UNet for Medical Image Segmentation". {Arxiv Paper}

Abstract

In the realm of medical image segmentation, both CNN-based and Transformer-based models have been extensively explored. However, CNNs exhibit limitations in long-range modeling capabilities, whereas Transformers are hampered by their quadratic computational complexity. Recently, State Space Models (SSMs), exemplified by Mamba, have emerged as a promising approach. They not only excel in modeling long-range interactions but also maintain a linear computational complexity. In this paper, leveraging state space models, we propose a U-shape architecture model for medical image segmentation, named Vision Mamba UNet (VM-UNet). Specifically, the Visual State Space (VSS) block is introduced as the foundation block to capture extensive contextual information, and an asymmetrical encoder-decoder structure is constructed. We conduct comprehensive experiments on the ISIC17, ISIC18, and Synapse datasets, and the results indicate that VM-UNet performs competitively in medical image segmentation tasks. To our best knowledge, this is the first medical image segmentation model constructed based on the pure SSM-based model. We aim to establish a baseline and provide valuable insights for the future development of more efficient and effective SSM-based segmentation systems.

0. Main Environments

conda create -n vmunet python=3.8
conda activate vmunet
pip install torch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu117
pip install packaging
pip install timm==0.4.12
pip install pytest chardet yacs termcolor
pip install submitit tensorboardX
pip install triton==2.0.0
pip install causal_conv1d==1.0.0  # causal_conv1d-1.0.0+cu118torch1.13cxx11abiFALSE-cp38-cp38-linux_x86_64.whl
pip install mamba_ssm==1.0.1  # mmamba_ssm-1.0.1+cu118torch1.13cxx11abiFALSE-cp38-cp38-linux_x86_64.whl
pip install scikit-learn matplotlib thop h5py SimpleITK scikit-image medpy yacs

The .whl files of causal_conv1d and mamba_ssm could be found here. {Baidu}

1. Prepare the dataset

ISIC datasets

  • The ISIC17 and ISIC18 datasets, divided into a 7:3 ratio, can be found here {Baidu or GoogleDrive}.

  • After downloading the datasets, you are supposed to put them into './data/isic17/' and './data/isic18/', and the file format reference is as follows. (take the ISIC17 dataset as an example.)

  • './data/isic17/'

    • train
      • images
        • .png
      • masks
        • .png
    • val
      • images
        • .png
      • masks
        • .png

Synapse datasets

  • For the Synapse dataset, you could follow Swin-UNet to download the dataset, or you could download them from {Baidu}.

  • After downloading the datasets, you are supposed to put them into './data/Synapse/', and the file format reference is as follows.

  • './data/Synapse/'

    • lists
      • list_Synapse
        • all.lst
        • test_vol.txt
        • train.txt
    • test_vol_h5
      • casexxxx.npy.h5
    • train_npz
      • casexxxx_slicexxx.npz

2. Prepare the pre_trained weights

  • The weights of the pre-trained VMamba could be downloaded here or Baidu. After that, the pre-trained weights should be stored in './pretrained_weights/'.

3. Train the VM-UNet

cd VM-UNet
python train.py  # Train and test VM-UNet on the ISIC17 or ISIC18 dataset.
python train_synapse.py  # Train and test VM-UNet on the Synapse dataset.

4. Obtain the outputs

  • After trianing, you could obtain the results in './results/'

5. Acknowledgments

  • We thank the authors of VMamba and Swin-UNet for their open-source codes.

About

(ARXIV24) This is the official code repository for "VM-UNet: Vision Mamba UNet for Medical Image Segmentation".

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages