Source-linked AI summary
CSWin-UNet: Transformer UNet with Cross-Shaped Windows for Medical Image Segmentation
Xiao Liu, Peng Gao, Tao Yu, Fei Wang, Ru-Yue Yuan
TL;DR
Medical image segmentation requires both detailed boundary information and efficient long-range modeling, but CNNs and Transformers each have important limitations. CSWin-UNet integrates CSWin self-attention into a U-shaped encoder-decoder and uses CARAFE for content-aware upsampling. Across CT, MRI, and skin-lesion datasets, it is reported to maintain low complexity while achieving high segmentation accuracy, though challenging low-contrast and variable cases remain.
Problem
CNNs have limited global and long-range semantic modeling, whereas Transformer-based methods can require high computational resources for medical image segmentation.
Method
CSWin-UNet is a U-shaped Transformer architecture using horizontal and vertical stripe self-attention, with CARAFE in the decoder for content-aware feature upsampling.
Results
CSWin-UNet reportedly surpasses other state-of-the-art methods across CT, MRI, and skin-lesion segmentation datasets while maintaining low model complexity.
Takeaways & Limitations
The method combines efficient receptive-field interactions with retained fine details and improved organ-edge segmentation precision.
Takeaways & Limitations
Performance remains deficient for some gallbladder and kidney samples and for low-contrast skin-lesion images.
Abstract
from arXiv · showhide
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.
1 Introduction
Medical image segmentation benefits from deep learning, but CNNs have limited global and long-range modeling while Transformers impose high computational demands. CSWin-UNet addresses this trade-off with a Transformer U-shaped architecture designed for efficient contextual modeling and detailed boundary recovery.
- Medical image segmentation extracts diseased-organ and tissue structure to support more accurate diagnostic and treatment recommendations.
- CNN-based methods classify pixels effectively but are limited by inductive biases and weak capture of global, long-range semantic information.
- Transformers use self-attention to model long-range dependencies and contextual relationships across image patches.
- CSWin-UNet is a pure Transformer U-shaped encoder-decoder that uses CSWin blocks in both encoder and decoder with scale-specific block counts.
- CSWin self-attention learns horizontal and vertical stripes, broadening each token’s focus area for more comprehensive contextual integration.
- CARAFE replaces conventional upsampling strategies in the decoder to produce more accurate pixel-level segmentation masks.
2 Related works
Related work contrasts CNNs’ detailed local modeling with Transformers’ broader contextual capacity and reviews increasingly structured attention mechanisms. CSWin-UNet builds on this progression by applying parallel cross-shaped stripe attention in a medical-image segmentation architecture.
- Axial and criss-cross attention compute attention along horizontal and vertical stripes, but face sequential, window-size, or overlapping-window limitations.
- CSWin Transformer computes horizontal and vertical stripe self-attention in parallel, motivating its use for efficient medical image segmentation.
- CNN-based segmentation architectures such as FCN and UNet are widely used, with UNet’s symmetric encoder-decoder supporting accurate medical image segmentation.
- CNNs capture detailed local features but often struggle to access global and long-range semantic context in high-resolution medical images.
- CSWin-UNet applies CSWin self-attention in a medical image segmentation model to conserve computational resources while elevating segmentation accuracy.
3 Methodology
CSWin-UNet is a four-stage U-shaped encoder-decoder using CSWin Transformer blocks, stripe-based self-attention, skip connections, and CARAFE upsampling. Its design combines horizontal and vertical contextual modeling with content-aware feature reassembly for resolution restoration.
- Overall architecture: CSWin-UNet uses an encoder, decoder, and skip connections, with CSWin Transformer blocks as its basic unit.Convolutional token embedding produces H/4×W/4 patch tokens with C channels, and both encoder and decoder contain four stages.
- CSWin Transformer block: A CSWin Transformer block combines CSWin self-attention with LayerNorm, an MLP, and skip connections.The block is defined around the CSWin-attention output and the preceding block or convolutional-layer output.
- CSWin Transformer block: The stripe width sw balances computational complexity and contextual coverage, with larger widths enabling longer-range pixel correlations within each stripe.The encoder changes sw across stages as feature resolution decreases.
- CSWin Transformer block: CSWin self-attention splits the heads into two groups that learn horizontal and vertical stripe attention in parallel before concatenating their outputs.The concatenated representation incorporates contextual information along both spatial directions.
- Decoder: The decoder uses CARAFE in its last three stages to increase resolution and channels through content-aware feature reassembly.CARAFE predicts kernels from encoded features and applies weighted reassembly over local neighborhoods, while skip connections restore spatial detail.
- Decoder: CARAFE predicts a distinct reassembly kernel for each output pixel from a local feature neighborhood, then applies that kernel to reassemble nearby features.The content encoder’s receptive field is denoted by kencoder, while kup specifies the reassembly-kernel size.
4 Experiments
Experiments evaluate CSWin-UNet across Synapse CT, ACDC cardiac MRI, and skin-lesion datasets using segmentation metrics and comparative tables or figures.
- Comparative evaluation: The Synapse experiments compare recent methods using DSC and HD, with first- and second-best values highlighted.
- Comparative evaluation: Synapse uncertainty is visualized with 95% confidence intervals for mean DSC, mean HD, and per-organ DSC.
- Evaluation metrics: Skin-lesion segmentation is evaluated with mean DSC, sensitivity, specificity, and accuracy.
4.3 Results on Synapse dataset
On Synapse, CSWin-UNet improves mean DSC and HD over TransUNet and Swin-UNet, while visualizations show stronger handling of difficult organ boundaries and backgrounds.
- Quantitative results: Mean DSC increases by 3.64% over TransUNet and 1.99% over Swin-UNet on Synapse.
- Quantitative results: Mean HD improves by 12.83% over TransUNet and 2.69% over Swin-UNet on Synapse.
- Quantitative results: CSWin-UNet achieves notably higher pancreas DSC than other segmentation methods, despite the pancreas having blurry and variable boundaries.
- Visual analysis: Competing methods show errors on small organs, incomplete stomach segmentation, or incorrect pancreas labeling in complex cases.
- Visual analysis: Visual comparisons cover eight abdominal organs, including the aorta, gallbladder, kidneys, liver, pancreas, spleen, and stomach.
4.4 Results on ACDC dataset
On ACDC cardiac MRI, CSWin-UNet is compared with advanced methods for three cardiac structures and reports 91.40% accuracy with claimed generalization and robustness.
- Evaluation: ACDC evaluation compares CSWin-UNet with other advanced methods on right ventricle, myocardium, and left ventricle segmentation.
- Results: Figure 7 represents 95% confidence intervals for mean DSC and per-structure DSC on ACDC.
- Results: 91.40% accuracy is reported for CSWin-UNet on the ACDC dataset.
4.5 Results on skin lesion segmentation datasets
Across ISIC2017, ISIC2018, and PH2, CSWin-UNet outperforms other methods on most metrics and preserves edges and detailed features, but errors remain under low contrast or occlusion.
- Quantitative results: CSWin-UNet outperforms other methods on most evaluation metrics across the three skin-lesion datasets.
- Quantitative results: CSWin-UNet performs better than Swin-UNet on most metrics, indicating satisfactory generalization capability.
- Visual analysis: Visual results show advantages in preserving segmented-object edges and detailed features compared with Swin-UNet.
- Limitations: Low contrast or occlusion can produce significant segmentation errors.
4.6 Comparison of computational efficiency
CSWin-UNet is evaluated for segmentation accuracy alongside parameter count and computational complexity, using FLOPs and parameters as measures. On Synapse, it achieves excellent segmentation performance under the lowest complexity.
- Table 5 presents an ablation study of upsampling strategies, with first- and second-best values highlighted.
- FLOPs and parameters measure computational complexity in the Synapse comparison.Parameters are reported in millions (M).
- CSWin-UNet achieves excellent segmentation performance under the lowest complexity on the Synapse dataset.
4.7 Ablation studies
The ablation studies examine decoder upsampling, skip connections, network architecture, and combined-loss hyperparameters. Results support CARAFE upsampling, three skip connections, and combined loss for segmentation performance.
- Ablation studies: The study varies upsampling strategies, skip connections, network architecture, and combined-loss hyperparameters on Synapse.
- Skip connection: Table 6 compares different numbers of skip connections on the Synapse dataset.
- Network architecture: Table 7 compares different numbers of CSWin Transformer blocks at each stage on the Synapse dataset.
- Upsampling strategy: CARAFE upsampling achieves the highest segmentation accuracy with little computational overhead compared with transposed convolution.CARAFE uses input content to guide feature reassembly for more accurate and efficient upsampling.
- Skip connection: Segmentation accuracy generally improves as the number of skip connections increases, with three selected for optimal performance.Improvements are more significant for smaller organs than for larger organs.
- Network architecture: Balancing network depth and model performance avoids insufficient feature representation and excessive computational burden.
- Combined loss: Combined loss yields higher segmentation accuracy than Dice or cross-entropy loss alone, especially compared with Dice-only loss.
4.8 Discussions
CSWin-UNet performs across CT, MRI, and skin-lesion datasets, but its results vary across challenging samples and low-contrast images, while pre-training strongly affects performance.
- CSWin-UNet shows different segmentation accuracy across gallbladder and kidney samples in Synapse.
- Low-contrast skin-lesion images remain a challenging setting with room for improved segmentation performance.
- Pre-training dramatically impacts CSWin-UNet performance.The study initializes the encoder and decoder with CSWin Transformer weights trained on ImageNet.
- The experiments cover CT, MRI, and skin-lesion medical image segmentation datasets.
5 Conclusion
CSWin-UNet combines CSWin self-attention with a U-shaped encoder-decoder and CARAFE upsampling to target receptive-field limitations, computational cost, and boundary detail. Across three large-scale datasets, it surpasses state-of-the-art methods in accuracy while remaining lightweight.
- CSWin-UNet integrates CSWin self-attention into a U-shaped encoder-decoder to address receptive-field interaction limitations.
- CARAFE upsampling helps retain intricate details and improve organ-edge segmentation precision.
- CSWin-UNet surpasses other state-of-the-art methods in segmentation accuracy across three large-scale medical image segmentation datasets.
- CSWin-UNet is lightweight in model parameters and computational load.