Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Ziyang Wang, Chao Ma

TL;DR

This work tackles the high annotation burden in medical image segmentation by proposing Weak-Mamba-UNet, a scribble-based weakly-supervised framework that fuses CNN (UNet), ViT (SwinUNet), and Visual Mamba (MambaUNet) under a multi-view cross-supervision scheme. It computes a dense pseudo label Y_{pseudo} = \alpha f_{cnn}(X;\theta) + \beta f_{vit}(X;\theta) + \gamma f_{mamba}(X;\theta) with random per-iteration weights while optimizing a combined loss of partial cross-entropy and dense Dice losses. Key contributions include the integration of a Mamba-based network into WSL for medical segmentation, a novel tri-network cross-supervision framework, and demonstrated gains on MRI cardiac segmentation with scribble annotations. The results indicate that Visual Mamba enhances long-range dependency modeling and, when fused with CNN and ViT, yields state-of-the-art performance under limited supervision, with code publicly available for reproducibility.

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.

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

TL;DR

This work tackles the high annotation burden in medical image segmentation by proposing Weak-Mamba-UNet, a scribble-based weakly-supervised framework that fuses CNN (UNet), ViT (SwinUNet), and Visual Mamba (MambaUNet) under a multi-view cross-supervision scheme. It computes a dense pseudo label Y_{pseudo} = \alpha f_{cnn}(X;\theta) + \beta f_{vit}(X;\theta) + \gamma f_{mamba}(X;\theta) with random per-iteration weights while optimizing a combined loss of partial cross-entropy and dense Dice losses. Key contributions include the integration of a Mamba-based network into WSL for medical segmentation, a novel tri-network cross-supervision framework, and demonstrated gains on MRI cardiac segmentation with scribble annotations. The results indicate that Visual Mamba enhances long-range dependency modeling and, when fused with CNN and ViT, yields state-of-the-art performance under limited supervision, with code publicly available for reproducibility.

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.
Paper Structure (6 sections, 4 equations, 3 figures, 2 tables)

This paper contains 6 sections, 4 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Scribble-Supervised Learning
  4. Multi-View Cross-Supervised Learning
  5. Experiments
  6. Conclusion

Figures (3)

  • Figure 1: The Example Images of MRI Cardiac Scans, with the Corresponding Ground Truth, and Scribble-based Annotations.
  • Figure 2: Semi-Mamba-UNet: The Framework of Contrastive Cross-Supervised Visual Mamba-based UNet for Semi-Supervised Medical Image Segmentation.
  • Figure 3: The Example Segmentation Results when 5% of Data are Assumed as Labeled Data.