[논문 리뷰] MAPPING: Model Average with Post-processing for Stroke Lesion Segmentation

Jiayu Huo, Liyun Chen외 5명|arXiv (Cornell University)|2022. 11. 11.

Acute Ischemic Stroke Management인용 수 4

한 줄 요약

이 논문은 ATLAS v2.0 데이터셋을 사용한 뇌경색 병변 세분화를 위한 모델 앙상블 및 후처리 방법인 MAPPING을 제안한다. 서로 다른 설정을 가진 네 개의 다양한 nnU-Net 모델을 훈련시킨 후 예측 결과를 앙상블하고, 이후 맞춤형 후처리 단계를 거쳐 2022년 MICCAI ATLAS 챌린지에서 가장 높은 성능을 기록했으며, 은닉 테스트 세트에서 DICE 스코어 0.6667, 최종 순위 1위를 달성하였다.

ABSTRACT

Accurate stroke lesion segmentation plays a pivotal role in stroke rehabilitation research, to provide lesion shape and size information which can be used for quantification of the extent of the stroke and to assess treatment efficacy. Recently, automatic segmentation algorithms using deep learning techniques have been developed and achieved promising results. In this report, we present our stroke lesion segmentation model based on nnU-Net framework, and apply it to the Anatomical Tracings of Lesions After Stroke (ATLAS v2.0) dataset. Furthermore, we describe an effective post-processing strategy that can improve some segmentation metrics. Our method took the first place in the 2022 MICCAI ATLAS Challenge with an average Dice score of 0.6667, Lesion-wise F1 score of 0.5643, Simple Lesion Count score of 4.5367, and Volume Difference score of 8804.9102. Our code and trained model weights are publicly available at https://github.com/King-HAW/ATLAS-R2-Docker-Submission.

연구 동기 및 목표

• 다양한 센터에서 수집된 MRI 데이터셋의 높은 변동성을 다룰 수 있는 강건하고 일반화 능력이 뛰어난 뇌경색 병변 세분화 방법을 개발하기 위해.
• 이질적인 뇌경색 MRI 데이터에 대해 훈련된 딥 러닝 모델에서의 과적합 및 분포 이탈 문제를 해결하기 위해.
• 모델 앙상블 및 후처리를 통해 소형 및 복잡한 병변의 세분화 성능을 향상시키기 위해.
• 특히 은닉 테스트 세트에서 ATLAS v2.0 벤치마크에서 최신 기술 수준의 성능을 달성하기 위해.
• 임상 및 연구 환경에서 사용 가능한 공개적이고 재현 가능한 뇌경색 병변 세분화 솔루션을 제공하기 위해.

제안 방법

• 다양한 데이터 증강, 손실 함수 및 훈련 체계를 사용하여 네 개의 서로 다른 nnU-Net 모델을 훈련시어 예측의 다양성을 극대화하기 위해.
• 네 개의 훈련된 모델의 예측 결과를 앙상블하여 평균을 내어 정확도를 향상시키고 분산을 줄이는 모델 앙상블 기법을 적용하기 위해.
• 형태학적 연산과 연결된 성분 분석을 사용하여 병변 경계를 정밀하게 다듬고, 작은 고립된 잡음 영역을 제거하는 맞춤형 후처리 전략을 구현하기 위해.
• 각 폴드에 걸쳐 병변 크기 분포가 균형을 이루도록 훈련 세트에서 크기 기반 다섯 개의 교차 검증을 수행하기 위해.
• MNI-152 템플릿 정규화 및 정렬을 사용하여 다양한 촬영 프로토콜 간에 입력 이미지를 표준화하기 위해.
• 공개 및 은닉 테스트 세트에서 공식 MICCAI 평가 플랫폼을 사용하여 엄격한 지표 가중치를 적용해 과적합을 방지하면서 방법을 검증하기 위해.

실험 결과

연구 질문

• RQ1 다양한 훈련 체계를 가진 모델 앙상블이 다중 센터 MRI 데이터셋에서 뇌경색 병변 세분화의 일반화 능력을 향상시킬 수 있는가?
• RQ2 후처리가 앙상블 예측을 얼마나 정교하게 다듬어 특정 평가 지표를 향상시킬 수 있으며, 과적합 없이 성능을 향상시킬 수 있는가?
• RQ3 작고 분리된 병변에 대해서는 이 방법이 어떻게 성능을 발휘하는가? 이는 뇌경색 세분화에서 알려진 도전 과제이다.
• RQ4 다양한 훈련 체계와 후처리의 조합이 ATLAS v2.0 챌린지의 은닉 테스트 세트에서 개별 모델보다 뛰어난 성능을 낼 수 있는가?
• RQ5 실제 임상 데이터에서 관찰되는 다양한 스캔 품질과 병변 형태학적 특성에 대해 이 방법은 높은 성능를 유지하는가?

주요 결과

• MAPPING 방법은 2022년 MICCAI ATLAS 챌린지에서 은닉 테스트 세트에서 DICE 스코어 0.6667을 기록하며 최고 성능을 달성하여 1위를 차지하였다.
• 네 개의 다양한 모델 앙상블은 개별 모델 대비 false positive를 크게 줄이고 병변 경계를 매끄럽게 하는 데서 뚜렷한 정확도 향상을 보였다.
• 후처리 단계는 작은 잡음 영역을 효과적으로 제거하였으며, 병변 수준 F1 스코어를 0.5643으로 향상시키고 단순 병변 수 스코어를 4.5367로 개선했다.
• 강력한 전체 성능에도 불구하고, 미세한 병변과 분리된 성분에 대해서는 여전히 어려움을 겪었으며, 정성적 실패 사례에서 이를 확인할 수 있었다.
• 이 방법은 다양한 스캔 프로토콜과 병변 형태학적 특성에 대해 강건성을 보이며, 새로운 데이터에 대한 뛰어난 일반화 능력을 입증하였다.
• 최종 랭킹에서 가중치 스코어 3.317을 기록하여, 더 복잡한 아키텍처를 가진 다른 경쟁 방법들보다도 뛰어난 성능을 입증하였다.