[논문 리뷰] D-TrAttUnet: Dual-Decoder Transformer-Based Attention Unet Architecture for Binary and Multi-classes Covid-19 Infection Segmentation

Fares Bougourzi, Cosimo Distante외 2명|arXiv (Cornell University)|2023. 03. 27.

COVID-19 diagnosis using AI인용 수 5

한 줄 요약

이 논문은 CT 영상에서 이중 및 다중 클래스 코로나19 감염 부위를 세분화하기 위한 이중 디코더를 갖춘 Transformer-CNN 하이브리드 U-Net 아키텍처인 D-TrAttUnet을 제안한다. 복합적인 CNN-Transformer 인코더와 두 개의 주의 메커니즘을 갖춘 디코더를 결합하여 기존의 최상위 성능(SOTA)을 초월하는 정확도와 강건성을 확보했으며, 제한된 데이터에서의 성능이 향상되었다.

ABSTRACT

In the last three years, the world has been facing a global crisis caused by Covid-19 pandemic. Medical imaging has been playing a crucial role in the fighting against this disease and saving the human lives. Indeed, CT-scans has proved their efficiency in diagnosing, detecting, and following-up the Covid-19 infection. In this paper, we propose a new Transformer-CNN based approach for Covid-19 infection segmentation from the CT slices. The proposed D-TrAttUnet architecture has an Encoder-Decoder structure, where compound Transformer-CNN encoder and Dual-Decoders are proposed. The Transformer-CNN encoder is built using Transformer layers, UpResBlocks, ResBlocks and max-pooling layers. The Dual-Decoder consists of two identical CNN decoders with attention gates. The two decoders are used to segment the infection and the lung regions simultaneously and the losses of the two tasks are joined. The proposed D-TrAttUnet architecture is evaluated for both Binary and Multi-classes Covid-19 infection segmentation. The experimental results prove the efficiency of the proposed approach to deal with the complexity of Covid-19 segmentation task from limited data. Furthermore, D-TrAttUnet architecture outperforms three baseline CNN segmentation architectures (Unet, AttUnet and Unet++) and three state-of-the-art architectures (AnamNet, SCOATNet and CopleNet), in both Binary and Mutli-classes segmentation tasks.

연구 동기 및 목표

• 코로나19 CT 영상 세분화에서 부족한 레이블이 부여된 데이터 문제를 해결한다.
• 흐린 경계와 낮은 대비를 보이는 이질적인 감염 패atters(GGO 및 복합성)를 세분화하는 데 어려움을 극복한다.
• CNN의 국소적 특징과 Transformer의 장거리 의존성 특징을 활용하여 세분화 정확도를 향상시킨다.
• 폐 부위와 감염 유형을 동시에 세분화하여 주의 메커니즘을 안내하고 잠재적 오진을 줄인다.
• 저자료 환경에서도 이중 및 다중 클래스 세분화 작업에 효과적인 강건한 아키텍처를 개발한다.

제안 방법

• 인코더는 이중 경로 아키텍처를 사용한다: 한 경로는 2D 자기주의를 통해 트랜스포머 레이어에서 영상 패치를 처리하고, 다른 경로는 잔차 블록을 갖춘 컨볼루션 레이어를 적용한다.
• 트랜스포머 경로에서 유도된 네 가지 다중 스케일 특징이 업샘플링 블록(UpResBlocks)을 통해 CNN 경로와 융합되어 계층적 표현을 풍부하게 한다.
• 이중 디코더는 주의 메커니즘을 갖춘 동일한 U-Net 스타일의 두 개의 디코더로 구성되며, 각각 폐 부위와 감염 마스크 예측을 담당한다.
• 두 디코더는 둘 다 폐와 감염 세분화를 동시에 최적화하는 병합된 손실 함수를 통해 공동으로 훈련된다.
• 디코더 경로에서는 선형 업샘플링과 컨볼루션 레이어를 사용하여 공간적 세부 정보를 유지하고 세분화 경계를 정밀하게 조정한다.
• 이 아키텍처는 이중 및 다중 클래스 작업 모두에 대해 세 개의 공개 코로나19 CT 세분화 데이터셋에서 엔드 투 엔드로 훈련된다.

실험 결과

연구 질문

• RQ1 하이브리드 CNN-Transformer 인코더는 코로나19 감염 세분화에 있어 국소적, 전반적, 장거리 맥락적 특징을 효과적으로 캡처할 수 있는가?
• RQ2 폐 부위와 감염 유형의 동시 세분화가 모델의 강건성과 정확도를 향상시키는가?
• RQ3 복잡한 감염 패턴을 다룰 때 이중 디코더 주의 메커니즘은 단일 디코더 아키텍처보다 우수한가?
• RQ4 제안된 아키텍처는 다양한 데이터 분포에 일반화되어 제한된 레이블 데이터에서 최상위 성능을 달성할 수 있는가?
• RQ5 각 구성 요소(예: 주의 게이트, 이중 경로 인코더, 병합된 손실)가 전체 성능에 기여하는 정도는 어떠한가?

주요 결과

• D-TrAttUnet은 이중 및 다중 클래스 세분화 작업 모두에서 세 가지 베이스라인 CNN 아키텍처(UNet, AttUNet, UNet++)와 세 가지 최상위 성능 모델(AnamNet, SCOATNet, CopleNet)을 모두 능가했다.
• 특히 흐린 경계를 가진 소형 및 주변 감염 부위를 탐지하는 데 뛰어난 세분화 정확도를 확보했다.
• 시각적 비교 결과, 혼합된 GGO 및 복합성 패턴을 포함한 복잡한 케이스에서 D-TrAttUnet의 예측 마스크가 진짜 마스크와 매우 유사하게 나타났다.
• 제거 실험을 통해 이중 디코더와 이중 경로 인코더의 조합이 개별 구성 요소보다 성능 향상에 크게 기여하는 것으로 확인되었다.
• 폐와 감염 세분화의 공동 훈련은 다양한 환자 케이스에서 더 일관되고 안정적인 예측을 이끌어냈다.
• 세 가지 다른 공개 데이터셋에서 강력한 일반화 능력을 보이며, 데이터 변동성에 대비한 강건성을 입증했다.