[Paper Review] BEFUnet: A Hybrid CNN-Transformer Architecture for Precise Medical Image Segmentation

Omid Nejati Manzari, Javad Mirzapour Kaleybar+2 more|arXiv (Cornell University)|Feb 13, 2024

Brain Tumor Detection and Classification9 citations

TL;DR

BEFUnet fuses edge-focused CNN encoding with body-focused Swin Transformer encoding via Local Cross-Attention Fusion and a Double-Level Fusion module to improve medical image segmentation performance across multiple datasets.

ABSTRACT

The accurate segmentation of medical images is critical for various healthcare applications. Convolutional neural networks (CNNs), especially Fully Convolutional Networks (FCNs) like U-Net, have shown remarkable success in medical image segmentation tasks. However, they have limitations in capturing global context and long-range relations, especially for objects with significant variations in shape, scale, and texture. While transformers have achieved state-of-the-art results in natural language processing and image recognition, they face challenges in medical image segmentation due to image locality and translational invariance issues. To address these challenges, this paper proposes an innovative U-shaped network called BEFUnet, which enhances the fusion of body and edge information for precise medical image segmentation. The BEFUnet comprises three main modules, including a novel Local Cross-Attention Feature (LCAF) fusion module, a novel Double-Level Fusion (DLF) module, and dual-branch encoder. The dual-branch encoder consists of an edge encoder and a body encoder. The edge encoder employs PDC blocks for effective edge information extraction, while the body encoder uses the Swin Transformer to capture semantic information with global attention. The LCAF module efficiently fuses edge and body features by selectively performing local cross-attention on features that are spatially close between the two modalities. This local approach significantly reduces computational complexity compared to global cross-attention while ensuring accurate feature matching. BEFUnet demonstrates superior performance over existing methods across various evaluation metrics on medical image segmentation datasets.

Motivation & Objective

• Motivate and address limitations of CNNs and pure Transformers in medical image segmentation, particularly edges and global context."
• Propose a dual-branch encoder to simultaneously capture edge and body features.
• Introduce Local Cross-Attention Fusion (LCAF) to efficiently fuse cross-modal features.
• Develop a Double-Level Fusion (DLF) module to merge coarse and fine-grained features across scales.
• Demonstrate superior segmentation performance on multiple medical image datasets.

Proposed method

• Propose BEFUnet with a dual-branch encoder: a CNN-based edge encoder using Pixel Different Convolution (PDC) blocks and a Swin Transformer body encoder.
• Introduce Local Cross-Attention Fusion (LCAF) to fuse edge and body features using local cross-attention, reducing computation versus global cross-attention.
• Implement a Double-Level Fusion (DLF) module that fuses smallest and largest level features via cross-attention and class-token based interaction.
• Train with two loss components for edge and body: edge supervision loss and body supervision loss (Binary Cross-Entropy + Dice Loss).
• Use a combined final loss L = L_body + gamma * L_edge to optimize both branches.

Figure 1 : The structure of our proposed BEFUnet. BEFUnet consists of a dual-encoder that simultaneously extracts edge and body features. Moreover, it incorporates an efficient fusion module called LCAF, which facilitates the merging of edge and body features. Additionally, there is a DLF module int

Experimental results

Research questions

• RQ1 Can a dual-branch encoder that separately models edges and body semantics improve segmentation accuracy over single-branch CNN or Transformer models?
• RQ2 Does Local Cross-Attention Fusion efficiently fuse cross-modal features while preserving accuracy and reducing computation?
• RQ3 Can the Double-Level Fusion module effectively integrate multi-scale features to boost segmentation performance?
• RQ4 What are the performance gains of BEFUnet across diverse medical imaging datasets compared to state-of-the-art methods?

Key findings

• BEFUnet achieves superior segmentation performance on Synapse with DSC 80.47 and HD 16.26, outperforming several SOTA methods.
• On SegPC dataset, BEFUnet achieves 95.1% accuracy and 0.871 F1-score, outperforming multiple baselines.
• On ISIC 2017 skin lesion segmentation, BEFUnet attains the best Dice 0.868, SE 0.853, SP 0.985, ACC 0.946, IoU 0.768.
• BEFUnet demonstrates strong edge segmentation capabilities due to the edge encoder with PDC blocks, improving boundary predictions.
• Ablation results (referenced in Table 4) show the contribution of the edge and body branches and fusion modules to overall performance.

Figure 2 : A block diagram of the LCAF.