[논문 리뷰] BEFUnet: A Hybrid CNN-Transformer Architecture for Precise Medical Image Segmentation

Omid Nejati Manzari, Javad Mirzapour Kaleybar외 2명|arXiv (Cornell University)|2024. 02. 13.

Brain Tumor Detection and Classification인용 수 9

한 줄 요약

BEFUnet은 Local Cross-Attention Fusion 및 Double-Level Fusion 모듈을 통해 엣지 중심 CNN 인코딩과 바디 중심 Swin Transformer 인코딩을 융합하여 다중 데이터셋에서 의료 영상 분할 성능을 향상시킨다.

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.

연구 동기 및 목표

• CNN의 한계와 의료 영상 분할에서의 순수 Transformer 및 특히 엣지와 글로벌 맥락의 문제점을 동기화하고 해결한다.
• 엣지와 바디 특징을 동시에 포착하기 위한 이중 브랜치 인코더를 제안한다.
• Cross-modal 특징을 효율적으로 융합하기 위한 Local Cross-Attention Fusion (LCAF)을 도입한다.
• 다단계에서 거칠고 섬세한 특징을 결합하기 위한 Double-Level Fusion (DLF) 모듈을 제시한다.
• 다양한 의료 영상 데이터셋에서 우수한 분할 성능을 시연한다.

제안 방법

• BEFUnet을 듀얼-브랜치 인코더로 제안: Pixel Different Convolution (PDC) 블록을 사용하는 CNN 기반 엣지 인코더와 Swin Transformer 바디 인코더를 포함한다.
• 엣지와 바디 특징을 지역 교차-attention으로 융합하는 Local Cross-Attention Fusion (LCAF)을 도입하여 전역 교차-attention 대비 계산을 줄인다.
• 가장 작은 단계와 가장 큰 단계의 특징을 교차-attention 및 클래스 토큰 기반 상호 작용으로 융합하는 Double-Level Fusion (DLF) 모듈을 구현한다.
• 엣지 및 바디에 대해 두 가지 손실 구성요소로 학습한다: 엣지 감독 손실과 바디 감독 손실(Binary Cross-Entropy + Dice Loss).
• 최종 손실 L = L_body + gamma * L_edge로 두 독립 브랜치를 함께 최적화한다.

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

실험 결과

연구 질문

• RQ1 엣지와 바디 의미를 분리적으로 모델링하는 듀얼-브랜치 인코더가 단일 브랜치 CNN 또는 Transformer 모델보다 분할 정확도를 향상시킬 수 있는가?
• RQ2 Local Cross-Attention Fusion이 정확도를 유지하고 계산량을 줄이면서 cross-modal 특징을 효율적으로 융합하는가?
• RQ3 Double-Level Fusion 모듈이 다중 스케일 특징을 효과적으로 통합하여 분할 성능을 높일 수 있는가?
• RQ4 BEFUnet의 다양한 의료 영상 데이터셋에서 최신 방법과 비교할 때 성능 향상은 어느 정도인가?

주요 결과

• BEFUnet은 Synapse에서 DSC 80.47 및 HD 16.26으로 우수한 분할 성능을 달성하며 다수의 SOTA 방법을 능가한다.
• SegPC 데이터셋에서 BEFUnet은 95.1% 정확도와 0.871 F1-점수를 달성하며 다수의 베이스라인을 능가한다.
• ISIC 2017 피부 병변 분할에서 BEFUnet은 최상의 Dice 0.868, SE 0.853, SP 0.985, ACC 0.946, IoU 0.768를 달성한다.
• PDC 블록을 갖춘 엣지 인코더로 인해 경계 예측이 향상되어 엣지 분할 능력이 강하게 나타난다.
• 테이블 4에 언급된 제거(추출) 결과에서 엣지 및 바디 브랜치와 융합 모듈의 공헌도가 전체 성능에 기여한다는 점을 시사한다.

Figure 2 : A block diagram of the LCAF.