[Paper Review] ParaTransCNN: Parallelized TransCNN Encoder for Medical Image Segmentation

Hongkun Sun, Jing Xu+1 more|arXiv (Cornell University)|Jan 27, 2024

Brain Tumor Detection and Classification11 citations

TL;DR

This paper introduces ParaTransCNN, a parallelized encoder combining a CNN and a Transformer with a pyramid structure and channel attention to fuse local and global features for medical image segmentation, achieving superior results especially on small organs across multiple datasets.

ABSTRACT

The convolutional neural network-based methods have become more and more popular for medical image segmentation due to their outstanding performance. However, they struggle with capturing long-range dependencies, which are essential for accurately modeling global contextual correlations. Thanks to the ability to model long-range dependencies by expanding the receptive field, the transformer-based methods have gained prominence. Inspired by this, we propose an advanced 2D feature extraction method by combining the convolutional neural network and Transformer architectures. More specifically, we introduce a parallelized encoder structure, where one branch uses ResNet to extract local information from images, while the other branch uses Transformer to extract global information. Furthermore, we integrate pyramid structures into the Transformer to extract global information at varying resolutions, especially in intensive prediction tasks. To efficiently utilize the different information in the parallelized encoder at the decoder stage, we use a channel attention module to merge the features of the encoder and propagate them through skip connections and bottlenecks. Intensive numerical experiments are performed on both aortic vessel tree, cardiac, and multi-organ datasets. By comparing with state-of-the-art medical image segmentation methods, our method is shown with better segmentation accuracy, especially on small organs. The code is publicly available on https://github.com/HongkunSun/ParaTransCNN.

Motivation & Objective

• Motivate improved medical image segmentation by capturing both local and global context.
• Develop a parallelized encoder that fuses CNN and Transformer representations at multiple scales.
• Incorporate a pyramid-structured Transformer and a channel attention module for effective feature fusion.
• Evaluate the approach on aortic vessel tree, cardiac, and multi-organ segmentation tasks against state-of-the-art methods.

Proposed method

• Two-branch encoder: a CNN (ResNet) for local features and a Transformer (ViT) for global features.
• Pyramid Transformer to learn global information at multiple resolutions (stages with downsampling factors 4, 8, 16).
• Patch embedding with stage-wise patch sizes to align CNN and Transformer features.
• Channel attention module to fuse CNN and Transformer features at each stage before passing to the decoder.
• U-Net–like encoder–decoder architecture with skip connections and a convolutional decoder.
• Combined loss of Dice and Cross-Entropy (equal weights 0.5).

Figure 1: Conceptual comparison of the three most popular models used for medical image segmentation, where (a) classical U-Net [ 39 ] ; (b) Swin U-Net [ 8 ] ; (c) TransUNet [ 10 ] ; (d) Our parallelized TransCNN encoder.

Experimental results

Research questions

• RQ1 Does a parallel CNN+Transformer encoder with pyramid global features improve segmentation accuracy over single-branch models?
• RQ2 Can channel attention effectively fuse multi-scale local and global features for accurate segmentation, including small organs?
• RQ3 What is the impact of patch embedding strategies and Transformer depth on segmentation performance?
• RQ4 How does ParaTransCNN perform on multi-center medical image datasets (AVT, ACDC, Synapse) compared to state-of-the-art methods?

Key findings

• ParaTransCNN achieves state-of-the-art or competitive results across AVT, ACDC, and Synapse datasets.
• On AVT, ParaTransCNN attains DSC = 87.82% and 95% HD = 4.70, outperforming the second-best method in DSC and HD.
• On ACDC, ParaTransCNN achieves DSC = 91.31% and HD = 1.16, surpassing several baselines in both metrics.
• On Synapse, ParaTransCNN reaches DSC = 83.86% and HD = 15.86, with notable improvements for pancreas and stomach segments.
• Ablation studies show the pyramid Transformer and channel attention are crucial for best performance, while patch overlap and deeper fourth-stage downsampling are not beneficial.
• Qualitative results indicate more continuous vessel structures and better small-organ segmentation (pancreas and spleen) compared to baselines.

Figure 2: The overall structure of the proposed parallelized TransCNN encoder network. The parallelized encoder consists of three stages, each of which extracts local and global information at different resolutions. The Channel Attention is applied to provide rich and useful semantic information for