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

Ziyang Wang,Chao Ma

2024-02-17

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 UNet for detailed local feature extraction, a Swin Transformer-based SwinUNet 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 a publicly available MRI cardiac segmentation dataset with processed scribble annotations, where it surpasses the performance of a similar WSL framework utilizing only UNet or SwinUNet. This highlights its potential in scenarios with sparse or imprecise annotations. The source code is made publicly accessible.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address the issue of the need for a large amount of precisely annotated data in medical image segmentation, especially when dealing with scribble annotations. Specifically, the paper proposes the Weak-Mamba-UNet framework, an innovative weakly supervised learning (WSL) method that combines the capabilities of Convolutional Neural Networks (CNN), Vision Transformers (ViT), and the Vision Mamba architecture. This framework leverages three different architectures—CNN-based UNet, Swin Transformer-based SwinUNet, and Mamba-based Mamba-UNet—and promotes iterative learning and optimization among the networks through a collaborative and cross-supervision mechanism.

This method improves segmentation performance under limited supervision signals by converting sparse annotation information into dense signal labels through the generation of pseudo-labels. Experimental results show that on publicly available cardiac MRI segmentation datasets, Weak-Mamba-UNet outperforms similar WSL frameworks that use only UNet or SwinUNet, especially in scenarios with scarce or imprecise annotations. This highlights the potential of this method in resource-constrained settings.