ziyangwang007/Mamba-UNet ziyangwang007/Mamba-UNet

  • Stars
    star
    410
  • Rank 104,838 (Top 3 %)
  • Language
    Python
  • License
    Apache License 2.0
  • Created 8 months ago
  • Updated 7 days ago
Twitter logo Share LinkedIn logo Share Facebook logo Share
ziyangwang007/Mamba-UNet
github

ziyangwang007

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

Mamba-UNet Zoo

Mamba-UNet

Mamba-UNet: Unet-like Pure Visual Mamba for Medical Image Segmentation

The position of Mamba-UNet

Semi-Mamba-UNet: Pixel-Level Contrastive Cross-Supervised Visual Mamba-based UNet for Semi-Supervised Medical Image Segmentation

The position of Semi-Mamba-UNet

Weak-Mamba-UNet: Visual Mamba Makes CNN and ViT Work Better for Scribble-based Medical Image Segmentation

The introduction of Scribble Annotation

Mamba-UNet

Results

Mamba-UNet Zoo

Project

Code

  • Code for Mamba-UNet -> Link
  • Code for Semi-Mamba-UNet -> Link
  • Code for Weak-Mamba-UNet -> Link
  • Code for VMambaMorph

More Experimnents

  • Dataset with ACDC MRI Cardiac MICCAI Challenge Dataset
  • Dataset with Synapse CT Abdomen MICCAI Challenge Dataset
  • Dataset with TotalSegmentator (in Progress, released soon.)

Requirements

  • Pytorch, MONAI
  • Some basic python packages: Torchio, Numpy, Scikit-image, SimpleITK, Scipy, Medpy, nibabel, tqdm ......
cd casual-conv1d

python setup.py install

cd mamba

python setup.py install

Usage

  1. Clone the repo:
git clone https://github.com/ziyangwang007/Mamba-UNet.git 
cd Mamba-UNet
  1. Download Pretrained Model

Download through Google Drive for SwinUNet, and Google Drive for Mamba-UNet, and save in 'code/pretrained_ckpt'.

  1. Download Dataset

Download ACDC for Semi-/Supervised learning through Google Drive, and save in 'data/ACDC' folder.

Download ACDC for Weak-Supervised learning through Google Drive, and save in 'data' folder.

Download Multi-Atlas Abdomen Labeling Challenge Dataset through Official Website Link, or Google Drive, and save in 'data/Synapse' folder.

  1. Train 2D UNet
python train_fully_supervised_2D.py --root_path ../data/ACDC --exp ACDC/unet --model unet --max_iterations 10000 --batch_size 24
  1. Train SwinUNet
python train_fully_supervised_2D_ViT.py --root_path ../data/ACDC --exp ACDC/swinunet --model swinunet --max_iterations 10000 --batch_size 24
  1. Train Mamba-UNet
python train_fully_supervised_2D_VIM.py --root_path ../data/ACDC --exp ACDC/VIM --model mambaunet --max_iterations 10000 --batch_size 24
  1. Train Semi-Mamba-UNet when 5% as labeled data
python train_Semi_Mamba_UNet.py --root_path ../data/ACDC --exp ACDC/Semi_Mamba_UNet --max_iterations 30000 --labeled_num 3 --batch_size 16 --labeled_bs 8
  1. Train Semi-Mamba-UNet when 10% as labeled data
python train_Semi_Mamba_UNet.py --root_path ../data/ACDC --exp ACDC/Semi_Mamba_UNet --max_iterations 30000 --labeled_num 7 --batch_size 16 --labeled_bs 8
  1. Train UNet with Mean Teacher when 5% as labeled data
python train_mean_teacher_2D.py --root_path ../data/ACDC --model unet --exp ACDC/Mean_Teacher --max_iterations 30000 --labeled_num 3 --batch_size 16 --labeled_bs 8
  1. Train SwinUNet with Mean Teacher when 5% as labeled data
python train_mean_teacher_ViT.py --root_path ../data/ACDC --model swinunet --exp ACDC/Mean_Teacher_ViT --max_iterations 30000 --labeled_num 3 --batch_size 16 --labeled_bs 8
  1. Train UNet with Mean Teacher when 10% as labeled data
python train_mean_teacher_2D.py --root_path ../data/ACDC --model unet --exp ACDC/Mean_Teacher --max_iterations 30000 --labeled_num 7 --batch_size 16 --labeled_bs 8
  1. Train SwinUNet with Mean Teacher when 10% as labeled data
python train_mean_teacher_ViT.py --root_path ../data/ACDC --model swinunet --exp ACDC/Mean_Teacher_ViT --max_iterations 30000 --labeled_num 7 --batch_size 16 --labeled_bs 8
  1. Train UNet with Uncertainty Aware Mean Teacher when 5% as labeled data
python train_uncertainty_aware_mean_teacher_2D.py --root_path ../data/ACDC --model unet --exp ACDC/Uncertainty_Aware_Mean_Teacher --max_iterations 30000 --labeled_num 3 --batch_size 16 --labeled_bs 8
  1. Train SwinUNet with Uncertainty Aware Mean Teacher when 5% as labeled data
python train_uncertainty_aware_mean_teacher_2D_ViT.py --root_path ../data/ACDC --model swinunet --exp ACDC/Uncertainty_Aware_Mean_Teacher_ViT --max_iterations 30000 --labeled_num 3 --batch_size 16 --labeled_bs 8
  1. Train UNet with Uncertainty Aware Mean Teacher when 10% as labeled data
python train_uncertainty_aware_mean_teacher_2D.py --root_path ../data/ACDC --model unet --exp ACDC/Uncertainty_Aware_Mean_Teacher --max_iterations 30000 --labeled_num 7 --batch_size 16 --labeled_bs 8
  1. Train SwinUNet with Uncertainty Aware Mean Teacher when 10% as labeled data
python train_uncertainty_aware_mean_teacher_2D_ViT.py --root_path ../data/ACDC --model swinunet --exp ACDC/Uncertainty_Aware_Mean_Teacher_ViT --max_iterations 30000 --labeled_num 7 --batch_size 16 --labeled_bs 8
  1. Test
python test_2D_fully.py --root_path ../data/ACDC --exp ACDC/xxx --model xxx

Q&A

  1. Q: Performance: I find my results are slightly lower than your reported results.

A: Please do not worry. The performance depends on many factors, such as how the data is split, how the network is initialized, how you write the evaluation code for Dice, Accuray, Precision, Sensitivity, Specificity, and even the type of GPU used. What I want to emphasize is that you should maintain your hyper-parameter settings and test some other baseline methods(fair comparsion). If method A has a lower/higher Dice Coefficient than the reported number, it's likely that methods B and C will also have lower/higher Dice Coefficients than the numbers reported in the paper.

  1. Q: Network Block: What is the network block you used? What is the difference between Mamba-XXXNet?

A: I understand that there are so many Mamba related papers now, such as Vision Mamba, Visual Mamba, SegMamba... In this project, I integrate VMamba into U-shape network. The reference of VMamba is: Liu, Yue, et al. "Vmamba: Visual state space model." arXiv preprint arXiv:2401.10166 (2024).

  1. Q: Concurrent Work: I found similar work about the integration of Mamba into UNet.

A: I am glad to see and acknowledge that there should be similar work. Mamba is a novel architecture, and it is obviously valuable to explore integrating Mamba with segmentation, detection, registration, etc. I am pleased that we all find Mamba efficient in some cases. This GitHub repository was developed on the 6th of February 2024, and I would not be surprised if people have proposed similar work from the end of 2023 to future. Also, I have only tested a limited number of baseline methods with a single dataset. Please make sure to read other related work around Mamba/Visual Mamba with UNet/VGG/Resnet etc.

  1. Q: Other Dataset: I want to try MambaUNet with other segmentation dataset, do you have any suggestions?

A: I recommend to start with UNet, as it often proves to be the most efficient architecture. Based on my experience with various segmentation datasets, UNet can outperform alternatives like TransUNet and SwinUNet. Therefore, UNet should be your first choice. Transformer-based UNet variants, which depend on pretraining, have shown promising results, especially with larger datasets—although such extensive datasets are uncommon in medical imaging. In my view, MambaUNet not only offers greater efficiency but also more promising performance compared to Transformer-based UNet. However, it's crucial to remember that MambaUNet, like Transformer, necessitates pretraining (e.g. on ImageNet).

  1. Q: Colloboration: Could I discuss with you about other topic, like Image Registration, Human Pose Estimation, Image Fusion, and etc.

A: I would also like to do some amazing work. Connect with me via ziyang [dot] wang17 [at] gmail [dot] com.

Reference

@article{wang2024mamba,
  title={Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation},
  author={Wang, Ziyang and others},
  journal={arXiv preprint arXiv:2402.05079},
  year={2024}
}

@article{wang2024semimamba,
  title={Semi-Mamba-UNet: Pixel-Level Contrastive Cross-Supervised Visual Mamba-based UNet for Semi-Supervised Medical Image Segmentation},
  author={Wang, Ziyang and others},
  journal={arXiv preprint arXiv:2402.07245},
  year={2024}
}

@article{wang2024weakmamba,
  title={Weak-Mamba-UNet: Visual Mamba Makes CNN and ViT Work Better for Scribble-based Medical Image Segmentation},
  author={Wang, Ziyang and others},
  journal={arXiv preprint arXiv:2402.10887},
  year={2024}
}

Contact

ziyang [dot] wang17 [at] gmail [dot] com

Acknowledgement

SSL4MIS Link, Segmamba Link, SwinUNet Link, Visual Mamba Link.

More Repositories

1

CV-SSL-MIS

Exploring CNN and ViT for Semi-Supervised Medical Image Segmentation
Python
92
star
2

VMambaMorph

VMambaMorph: a Multi-Modality Deformable Image Registration Framework based on Visual State Space Model with Cross-Scan Module
Python
39
star
3

CV-WSL-MIS

Exploring CNN and ViT for Weakly-Supervised Medical Image Segmentation
Python
28
star
4

Awesome-Medical-Image-Segmentation-Dataset

A list of publicly available medical image segmentation dataset.
27
star
5

VIT4UNet

Python
17
star
6

Realtime-Openpose-on-iOS-with-Double-Robot

Swift
13
star
7

CV-WSL-Robot

Weakly Supervised Robotic Surgery Segmentation
Python
13
star
8

Weak-Mamba-UNet

Weak-Mamba-UNet: Visual Mamba Makes CNN and ViT Work Better for Scribble-based Medical Image Segmentation
Python
13
star
9

MixSegNet

MixSegNet: Mixed-Supervised Learning for Medical Image Segmentation
Python
11
star
10

Awesome-Visual-Mamba

8
star
11

TriConvUNeXt

Exploring Dilated CNN, Deformable CNN, and Depthwise CNN for medical image segmentation.
Python
6
star
12

UAMT-ViT

Python
5
star
13

TMC5160Arduino

the Arduino Libary for TMC5160 Stepper Motor Driver
C++
5
star
14

CVPixUNet

Rethinking U-Shape Segmentation Network: Towards CNN- \& ViT-based Hybrid Network with Dynamic Adaptive Pixel-Level Feature Learning for Retinal Vessel Segmentation
5
star
15

UNet-Seg

Exploration of CNN-, ViT-, Mamba-, and KAN-based UNet for Medical Image Segmentation.
1
star
16

CV-MSL-MIS

Python
1
star