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>
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
- 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
- Clone the repo:
git clone https://github.com/ziyangwang007/Weak-Mamba-UNet.git
cd Weak-Mamba-UNet
- Download Pretrained Model
Download through Google Drive for SwinUNet, and [Google Drive] for Mamba-UNet, and save in ../code/pretrained_ckpt.
- Download Dataset
Download ACDC for Weak-Supervised learning through [Google Drive], or [Baidu Netdisk] with passcode: 'rwv2', and save in ../data/ACDC folder.
- Train
cd code
- Train 2D UNet with pCE
python train_weakly_supervised_pCE_2D.py
- Train 2D SwinUNet with pCE
python train_weakly_supervised_pCE_2D_ViT.py
- Train 2D SwinUNet with MT and pCE
python train_weakly_supervised_ustm_2D_ViT.py
- Train 2D Semi-Mamba-UNet with pCE
python train_weak_mamba_unet.py
- 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