Home

Awesome

<div align="center"> <h1> <b>Weak-Mamba-UNet</b>: <br /> Visual Mamba Makes CNN and ViT Work Better for Scribble-based Medical Image Segmentation </h1>

arXiv

</div>

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

Contents

Graphical Abstract

The introduction of Scribble Annotation

<img src="img/wslintro.png" width="50%" height="auto">

The proposed Framework

<img src="img/wslframework.png" width="50%" height="auto">

Results

<img src="img/results.png" width="50%" height="auto">

Requirements

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
  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 Weak-Supervised learning through [Google Drive], or [Baidu Netdisk] with passcode: 'rwv2', and save in ../data/ACDC folder.

  1. Train
cd code
  1. Train 2D UNet with pCE
python train_weakly_supervised_pCE_2D.py 
  1. Train 2D SwinUNet with pCE
python train_weakly_supervised_pCE_2D_ViT.py 
  1. Train 2D SwinUNet with MT and pCE
python train_weakly_supervised_ustm_2D_ViT.py 
  1. Train 2D Semi-Mamba-UNet with pCE
python train_weak_mamba_unet.py 
  1. 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