CSWin-UNet: Transformer UNet with Cross-Shaped Windows for Medical Image Segmentation

TL;DR

CSWin-UNet is a novel Transformer-based U-shaped segmentation model that improves efficiency and accuracy in medical image segmentation by integrating cross-shaped window self-attention and a content-aware reassembly decoder.

Contribution

It introduces a new self-attention mechanism with horizontal and vertical stripes and a content-aware decoder, enhancing segmentation performance with low model complexity.

Findings

Achieves high segmentation accuracy across multiple datasets.

Maintains low computational complexity compared to other Transformer models.

Demonstrates superior performance in diverse medical imaging scenarios.

Abstract

Deep learning, especially convolutional neural networks (CNNs) and Transformer architectures, have become the focus of extensive research in medical image segmentation, achieving impressive results. However, CNNs come with inductive biases that limit their effectiveness in more complex, varied segmentation scenarios. Conversely, while Transformer-based methods excel at capturing global and long-range semantic details, they suffer from high computational demands. In this study, we propose CSWin-UNet, a novel U-shaped segmentation method that incorporates the CSWin self-attention mechanism into the UNet to facilitate horizontal and vertical stripes self-attention. This method significantly enhances both computational efficiency and receptive field interactions. Additionally, our innovative decoder utilizes a content-aware reassembly operator that strategically reassembles features, guided by predicted kernels, for precise image resolution restoration. Our extensive empirical evaluations on diverse datasets, including synapse multi-organ CT, cardiac MRI, and skin lesions, demonstrate that CSWin-UNet maintains low model complexity while delivering high segmentation accuracy. Codes are available at https://github.com/eatbeanss/CSWin-UNet.