H-vmunet: High-order Vision Mamba UNet for Medical Image Segmentation

Renkai Wu, Yinghao Liu, Pengchen Liang*, and Qing Chang*

News🚀

(2024.03.21) Model weights have been uploaded for download🔥🔥

(2024.03.21) The project code has been uploaded🔥

(2024.03.20) The first edition of our paper has been uploaded to arXiv 📃

0. Main Environments.
The environment installation procedure can be followed by VM-UNet, or by following the steps below:

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

1. Datasets.

A.ISIC2017
1- Download the ISIC 2017 train dataset from this link and extract both training dataset and ground truth folders inside the /data/dataset_isic17/.
2- Run Prepare_ISIC2017.py for data preparation and dividing data to train,validation and test sets.

B.Spleen
1- Download the Spleen dataset from this link.

C.CVC-ClinicDB
1- Download the CVC-ClinicDB dataset from this link.

D. Prepare your own dataset

1. The file format reference is as follows. (The image is a 24-bit png image. The mask is an 8-bit png image. (0 pixel dots for background, 255 pixel dots for target))

• './your_dataset/'
  • images
    • 0000.png
    • 0001.png
  • masks
    • 0000.png
    • 0001.png
  • Prepare_your_dataset.py

2. In the 'Prepare_your_dataset.py' file, change the number of training sets, validation sets and test sets you want.
   
3. Run 'Prepare_your_dataset.py'.
   

2. Train the H_vmunet.

python train.py

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

3. Test the H_vmunet.
First, in the test.py file, you should change the address of the checkpoint in 'resume_model'.

python test.py

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

4. Get model weights

A.ISIC2017
Google Drive

B.Spleen
Google Drive

C.CVC-ClinicDB
Google Drive

Citation

If you find this repository helpful, please consider citing:

@article{wu2024h,
  title={H-vmunet: High-order Vision Mamba UNet for Medical Image Segmentation},
  author={Wu, Renkai and Liu, Yinghao and Liang, Pengchen and Chang, Qing},
  journal={arXiv preprint arXiv:2403.13642},
  year={2024}
}

Acknowledgement

Thanks to Vim, HorNet and VM-UNet for their outstanding work.