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

This repo provides an implementation of the training and inference pipeline for Weak-Mamba-UNet.

Contents

• Graphical Abstract
• Results
• Requirements
• Usage
• Reference
• Contact

Graphical Abstract

The introduction of Scribble Annotation

The proposed Framework

Results

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/Weak-Mamba-UNet.git
cd Weak-Mamba-UNet

2. Download Pretrained Model

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

3. Download Dataset

Download ACDC for Weak-Supervised learning through [Google Drive], or [Baidu Netdisk] with passcode: 'rwv2', and save in ../data/ACDC folder.

4. Train

cd code

5. Train 2D UNet with pCE

python train_weakly_supervised_pCE_2D.py 

6. Train 2D SwinUNet with pCE

python train_weakly_supervised_pCE_2D_ViT.py 

7. Train 2D SwinUNet with MT and pCE

python train_weakly_supervised_ustm_2D_ViT.py 

8. Train 2D Semi-Mamba-UNet with pCE

python train_weak_mamba_unet.py 

9. Test

Test CNN-based model

python test_2D.py -root_path ../data/XXX --exp ACDC/XXX

Test ViT/Mamba-based model

python test_2D_fully.py -root_path ../data/XXX --exp ACDC/XXX

Reference

Wang, Ziyang, et al. "Mamba-unet: Unet-like pure visual mamba for medical image segmentation." arXiv preprint arXiv:2402.05079 (2024).

Wang, Ziyang, and Chao Ma. "Weak-Mamba-UNet: Visual Mamba Makes CNN and ViT Work Better for Scribble-based Medical Image Segmentation." arXiv preprint arXiv:2402.10887 (2024).

@article{wang2024mamba,
  title={Mamba-unet: Unet-like pure visual mamba for medical image segmentation},
  author={Wang, Ziyang and Zheng, Jian-Qing and Zhang, Yichi and Cui, Ge and Li, Lei},
  journal={arXiv preprint arXiv:2402.05079},
  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 Ma, Chao},
  journal={arXiv preprint arXiv:2402.10887},
  year={2024}
}

Contact

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