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

Abstract

Medical image segmentation is increasingly reliant on deep learning techniques, yet the promising performance often come with high annotation costs. This paper introduces Weak-Mamba-UNet, an innovative weakly-supervised learning (WSL) framework that leverages the capabilities of Convolutional Neural Network (CNN), Vision Transformer (ViT), and the cutting-edge Visual Mamba (VMamba) architecture for medical image segmentation, especially when dealing with scribble-based annotations. The proposed WSL strategy incorporates three distinct architecture but same symmetrical encoder-decoder networks: a CNN-based U-Net for detailed local feature extraction, a Swin Transformer-based Swin-UNet for comprehensive global context understanding, and a VMamba-based Mamba-UNet for efficient long-range dependency modeling. The key concept of this framework is a collaborative and cross-supervisory mechanism that employs pseudo labels to facilitate iterative learning and refinement across the networks. The effectiveness of Weak-Mamba-UNet is validated on two publicly available datasets with processed scribble annotations, where it surpasses the performance of a similar WSL framework utilizing only U-Net or Swin-UNet, as well as other baseline methods. This paper highlights the potential of Mamba for medical image segmentation in scenarios with sparse or imprecise annotations. The source code, dataset, and all baseline methods are made publicly accessible https://github.com/ziyangwang007/Mamba-UNet.