Skip to content

[arXiv] The official code for "H-vmunet: High-order Vision Mamba UNet for Medical Image Segmentation".

Notifications You must be signed in to change notification settings

wurenkai/H-vmunet

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

16 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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

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

arXiv

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
  1. In the 'Prepare_your_dataset.py' file, change the number of training sets, validation sets and test sets you want.
  2. 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.

About

[arXiv] The official code for "H-vmunet: High-order Vision Mamba UNet for Medical Image Segmentation".

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages