[논문 리뷰] Neuro-TransUNet: Segmentation of stroke lesion in MRI using transformers

Muhammad Nouman, Mohamed A. Mabrok외 1명|arXiv (Cornell University)|2024. 06. 10.

Brain Tumor Detection and ClassificationNeuroscience인용 수 3

한 줄 요약

Neuro-TransUNet은 U-Net의 공간적 특징 추출 능력과 SwinUNETR의 전역적 맥락 모델링 능력을 융합하는 하이브리드 딥러닝 프레임워크를 제안한다. 이는 3D MRI 뇌경색 병변 분할을 향상시키기 위해 고도화된 특징 융합과 종합적인 사전처리 파이프라인을 통합함으로써 ATLAS v2.0 벤치마크에서 테스트 Dice 스코어 0.730을 달성하였으며, 기존 방법들을 능가하고 뇌경색 병변 분할 분야에서 새로운 최고 성능 기록을 수립하였다.

ABSTRACT

Accurate segmentation of the stroke lesions using magnetic resonance imaging (MRI) is associated with difficulties due to the complicated anatomy of the brain and the different properties of the lesions. This study introduces the Neuro-TransUNet framework, which synergizes the U-Net's spatial feature extraction with SwinUNETR's global contextual processing ability, further enhanced by advanced feature fusion and segmentation synthesis techniques. The comprehensive data pre-processing pipeline improves the framework's efficiency, which involves resampling, bias correction, and data standardization, enhancing data quality and consistency. Ablation studies confirm the significant impact of the advanced integration of U-Net with SwinUNETR and data pre-processing pipelines on performance and demonstrate the model's effectiveness. The proposed Neuro-TransUNet model, trained with the ATLAS v2.0 \emph{training} dataset, outperforms existing deep learning algorithms and establishes a new benchmark in stroke lesion segmentation.

연구 동기 및 목표

• 해부학적 복잡성으로 인해 작은 크기이거나 형태가 다양하여 감지가 어려운 3D MRI에서의 뇌경색 병변 정확한 분할에 도전한다.
• 작은 병변이나 다중 병변에 특히 취약한 기존 모델의 병변 크기 변동성 및 데이터 불균형 문제를 해결한다.
• 두개골 제거 없이도 데이터 품질과 일관성을 향상시키는 종합적인 사전처리 파이프라인을 개발하여 모델의 강건성과 일반화 능력을 향상시킨다.
• U-Net(국소적 공간적 특징)과 SwinUNETR(전역적 맥락 이해)의 강점을 고도화된 특징 융합 및 분할 합성 기법을 통해 통합한다.
• 어려운 ATLAS v2.0 데이터셋을 기반으로 모델을 검증하여 뇌경색 병변 분할 분야의 새로운 기준을 설정하며, 이전 딥러닝 방법들보다 뛰어난 성능을 입증한다.

제안 방법

• U-Net의 인코더-디코더 아키텍처와 SwinUNETR의 트랜스포머 기반 디코더를 융합하여 국소적 공간적 특징 학습과 장거리 상관관계 모델링의 시너지를 유도한다.
• 다중 해상도 수준에서 U-Net 및 SwinUNETR 브랜치의 특징을 다중 단계 특징 융합 메커니즘을 통해 집계함으로써 표현 학습 능력을 향상시킨다.
• 해상도 조정, 왜곡 필드 보정, 데이터 표준화를 포함한 종합적인 사전처리 파이프라인을 적용하여 환자 간 입력 데이터 품질과 일관성을 향상시킨다.
• 기존 2D/3D 후처리 방식을 회피하는 볼륨 기반 추론 방식을 사용하여 공간 일관성이 향상된 엔드 투 엔드 3D 분할을 가능하게 한다.
• 다양한 병변 크기와 해부학적 위치를 가진 뇌경색 사례를 포함하는 ATLAS v2.0 데이터셋을 학습 및 평가에 활용한다.
• 클래스 불균형을 고려한 손실 함수를 최적화하여 작은 병변 및 낮은 빈도의 병변에 대한 민감도를 향상시킨다.

실험 결과

연구 질문

• RQ1 U-Net과 SwinUNETR를 융합한 하이브리드 아키텍처가 3D MRI에서 모든 크기의 뇌경색 병변 분할 정확도를 향상시킬 수 있는가?
• RQ2 두개골 제거 없이도 종합적인 사전처리 파이프라인이 실제 MRI 데이터에서 모델 성능과 강건성을 얼마나 향상시키는가?
• RQ3 고도화된 특징 융합 및 분할 합성 기법은 국소적 및 전역적 특징을 통합하여 병변 경계를 개선하는 데 얼마나 효과적인가?
• RQ4 Neuro-TransUNet 프레임워크는 ATLAS v2.0 벤치마크에서 기존 최고 성능 모델들을 능가하는가?
• RQ5 모델는 복잡하거나 다양하거나 희귀한 병변 패턴을 다루는 데에서 어떤 한계를 지니며, 향후 연구는 일반화 능력을 어떻게 향상시킬 수 있는가?

주요 결과

• Neuro-TransUNet는 볼륨 기반 추론 방식을 사용하여 ATLAS v2.0 벤치마크 데이터셋에서 테스트 Dice 스코어 0.730을 달성하였으며, 기존 딥러닝 모델들을 능가하였다.
• 제거 실험을 통해 U-Net과 SwinUNETR의 고도화된 특징 융합 통합이 단독 아키텍처보다 분할 성능을 크게 향상시킴을 확인하였다.
• 해상도 조정, 왜곡 보정, 표준화를 포함한 종합적인 사전처리 파이프라인은 모델 성능에 결정적인 영향을 미쳤으며, 이를 제거할 경우 Dice 스코어가 0.65로 감소하였다.
• 모델는 병변 크기에 관계없이 작은 병변, 중간 크기 병변, 큰 병변, 다중 병변을 한 명의 환자에서 모두 성공적으로 감지하는 데 강력한 일반화 능력을 보였다.
• 높은 성능에도 불구하고, 매우 복잡하거나 희귀한 병변 패턴을 다루는 데서 여전히 어려움을 겪고 있어, 더 크고 다양한 훈련 데이터가 필요하다는 점을 시사한다.
• 모델의 높은 파rameter 수(1억 개 이상)는 세밀한 특징 학습을 가능하게 하지만, 임상 환경에서의 계산 가능성에 대한 우려를 제기한다.