Home

Awesome

<div align="center"> <h1> UNet Zoo for Medical Image Segmentation </h1> <h2> CNN vs ViT vs Mamba vs KAN </h2> </div>

Introduction

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

Over 10 Segmentation Networks, 7 public benchmark datasets, 6 evaluation metrics are public available!

Contents

Networks

ArchitectureNameReferenceAvailability
CNNU-netRonneberger et al. (2015)
CNNAttention U-netOktay et al. (2018)
CNNResUnetDiakogiannis et al. (2020)
CNNDenseUnetLi et al. (2020)
CNNConvUNeXTHan et al. (2022)
CNNCMUNeXtTang et al. (2023)
CNNAMSUnetYin et al. (2023)
CNNMBSNetJin et al. (2023)
CNNTriConvUNeXTMa et al. (2024)
ViTTransUNETChen et al. (2021)
ViTSwin-UNETHu et al. (2021)
ViTH2FormerHe et al. (2023)
ViTUDTransNetWang et al. (2024)
ViTCSWin-UNetLiu et al. (2025)
MambaMamba-UNetWang et al. (2024)
MambaVM-UNetRuan et al. (2024)
KANU-KANLi et al. (2024)
KANKAN-UNetWang et al. (2024)

Datasets

DatasetReferenceGoogle DriveBaidu NetdiskStatus
GLASOfficialLinkLink (Passcode: fp42)
BUSIOfficialLinkLink (Passcode: 5m5m)
2018DSBOfficialLinkLink (Passcode: yavx)
CVC-ClinicDBOfficialLinkLink (Passcode: 3tpy)
Kvasir-SEGOfficialLinkLink (Passcode: 6fgs)
ISIC2016OfficialLinkLink (Passcode: m2fw)
PH2OfficialLinkLink (Passcode: aiax)
ACDCOfficial
SynapseOfficial

Metrics

MetricDescriptionFormulaStatus
DiceMeasures overlap between predicted and ground truth masks$\frac{2TP}{2TP + FP + FN}$
IoUIntersection over Union, also known as Jaccard Index$\frac{TP}{TP + FP + FN}$
AccuracyProportion of correctly classified elements$\frac{TP + TN}{TP + FP + FN + TN}$
PrecisionProportion of true positive predictions$\frac{TP}{TP + FP}$
SensitivityTrue positive rate (Recall)$\frac{TP}{TP + FN}$
SpecificityTrue negative rate$\frac{TN}{TN + FP}$

Environment

Usage

  1. Download the Code.
git clone https://github.com/ziyangwang007/UNet-Seg.git 
cd UNet-Seg
  1. Download the Dataset via Google Drive or Baidu Netdisk to UNet-Seg/data folder.

  2. Download Pretrained Model

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

  1. Train the model.
CUDA_VISIBLE_DEVICES=0 python -u train.py --network UNet --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network UNet --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network UNet --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network UNet --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network UNet --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network UNet --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network UNet --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u train.py --network DenseUnet --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network DenseUnet --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network DenseUnet --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network DenseUnet --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network DenseUnet --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network DenseUnet --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network DenseUnet --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u train.py --network AttU_Net --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network AttU_Net --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network AttU_Net --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network AttU_Net --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network AttU_Net --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network AttU_Net --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network AttU_Net --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u train.py --network ConvUNeXt --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network ConvUNeXt --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network ConvUNeXt --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network ConvUNeXt --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network ConvUNeXt --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network ConvUNeXt --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network ConvUNeXt --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u train.py --network SwinUnet --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network SwinUnet --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network SwinUnet --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network SwinUnet --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network SwinUnet --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network SwinUnet --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network SwinUnet --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u train.py --network TransUNet --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network TransUNet --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network TransUNet --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network TransUNet --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network TransUNet --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network TransUNet --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network TransUNet --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u train.py --network KANUSeg --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network KANUSeg --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network KANUSeg --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network KANUSeg --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network KANUSeg --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network KANUSeg --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u train.py --network KANUSeg --datasets 2018DSB
  1. Test the model.
CUDA_VISIBLE_DEVICES=0 python -u test.py --network UNet --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network UNet --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network UNet --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network UNet --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network UNet --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network UNet --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network UNet --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u test.py --network DenseUnet --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network DenseUnet --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network DenseUnet --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network DenseUnet --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network DenseUnet --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network DenseUnet --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network DenseUnet --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u test.py --network AttU_Net --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network AttU_Net --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network AttU_Net --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network AttU_Net --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network AttU_Net --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network AttU_Net --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network AttU_Net --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u test.py --network ConvUNeXt --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network ConvUNeXt --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network ConvUNeXt --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network ConvUNeXt --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network ConvUNeXt --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network ConvUNeXt --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network ConvUNeXt --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u test.py --network SwinUnet --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network SwinUnet --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network SwinUnet --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network SwinUnet --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network SwinUnet --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network SwinUnet --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network SwinUnet --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u test.py --network TransUNet --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network TransUNet --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network TransUNet --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network TransUNet --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network TransUNet --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network TransUNet --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network TransUNet --datasets 2018DSB && \

CUDA_VISIBLE_DEVICES=0 python -u test.py --network KANUSeg --datasets PH2  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network KANUSeg --datasets isic16  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network KANUSeg --datasets BUSI  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network KANUSeg --datasets GLAS  && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network KANUSeg --datasets CVC-ClinicDB && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network KANUSeg --datasets Kvasir-SEG && \
CUDA_VISIBLE_DEVICES=0 python -u test.py --network KANUSeg --datasets 2018DSB

  1. You can check the evaluation metrics directly after testing, and find prediction images in UNet-Seg/output.

Reference

If you find this repository useful for your research, please consider citing the following papers.

@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{zhang2024survey,
  title={A survey on visual mamba},
  author={Zhang, Hanwei and Zhu, Ying and Wang, Dan and Zhang, Lijun and Chen, Tianxiang and Wang, Ziyang and Ye, Zi},
  journal={Applied Sciences},
  volume={14},
  number={13},
  pages={5683},
  year={2024},
  publisher={MDPI}
}