H-vmunet: High-order Vision Mamba UNet for Medical Image Segmentation
January 23, 2025 ยท View on GitHub
H-vmunet: High-order Vision Mamba UNet for Medical Image Segmentation
Renkai Wu, Yinghao Liu, Pengchen Liang*, and Qing Chang*
News๐
(2025.01.12) The paper has been accepted by Neurocomputing๐ฅ๐ฅ
(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
- 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
- images
- In the 'Prepare_your_dataset.py' file, change the number of training sets, validation sets and test sets you want.
- 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{wu2025h,
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={Neurocomputing},
pages={129447},
year={2025},
publisher={Elsevier}
}
Acknowledgement
Thanks to Vim, HorNet and VM-UNet for their outstanding work.