[논문 리뷰] Brain Tumor Segmentation using 3D-CNNs with Uncertainty Estimation

Laura Mora Ballestar, Verónica Vilaplana|arXiv (Cornell University)|2020. 09. 24.

Brain Tumor Detection and Classification참고 문헌 18인용 수 13

한 줄 요약

이 논문은 뇌 종양 MRI를 위한 3D V-Net 기반의 의미 분할 모델을 제안하며, 테스트 시점 드롭아웃(지식 기반)과 데이터 증강(확률 기반)을 통한 병렬 볼륨 수준의 불확실성 추정을 통합한다. 이는 종양 분할의 정확성과 신뢰도를 향상시키며, 특히 저신뢰 영역을 식별하는 데 기여한다. 후처리를 거친 BraTS 2020 검증 세트에서 DICE 스코어는 WT 0.8463, TC 0.7526, ET 0.6179를 기록하였고, 낮은 불확실성 필터링 비율(FTP: WT 기준 0.0118)을 유지하였다.

ABSTRACT

Automation of brain tumors in 3D magnetic resonance images (MRIs) is key to assess the diagnostic and treatment of the disease. In recent years, convolutional neural networks (CNNs) have shown improved results in the task. However, high memory consumption is still a problem in 3D-CNNs. Moreover, most methods do not include uncertainty information, which is specially critical in medical diagnosis. This work proposes a 3D encoder-decoder architecture, based on V-Net \cite{vnet} which is trained with patching techniques to reduce memory consumption and decrease the effect of unbalanced data. We also introduce voxel-wise uncertainty, both epistemic and aleatoric using test-time dropout and data-augmentation respectively. Uncertainty maps can provide extra information to expert neurologists, useful for detecting when the model is not confident on the provided segmentation.

연구 동기 및 목표

• 뇌 종양 분할을 위한 3D-CNN에서 높은 메모리 소비와 불확실성 정량화의 부재 문제를 해결한다.
• 패치 기반 기법과 데이터 증강을 활용하여, 다중 모odal MRI에서의 불균형 데이터에 대한 분할 정확도를 향상시킨다.
• 의사들이 신뢰할 수 없는 예측을 식별하는 데 도움이 되는 볼륨 수준의 불확실성 지ap (지식 기반 및 확률 기반)을 제공한다.
• 거짓 양성 및 종양 경계의 저신뢰 영역을 탐지할 수 있도록 하여 임상적 신뢰도를 향상시킨다.
• 후처리 및 불확실성 인식 추론을 통해 BraTS 2020 데이터셋에서 모델 성능을 최적화한다.

제안 방법

• 메모리 소비를 줄이기 위해 특징 맵을 감소시키고 배치 정규화 대신 인스턴스 정규화를 사용하여 V-Net 아키텍처를 변형한다.
• 큰 3D MRI 볼륨을 처리하고 클래스 불균형 문제를 완화하기 위해 패치 기반 학습 기법을 적용한다.
• 테스트 시점 드롭아웃(TTD)을 통해 지식 기반 불확실성 추정을 수행하며, 추론 중 다중 순차 전방 전파의 분산을 계산한다.
• 데이터 증강(예: 회전, 반전 등)을 활용하여 확률 기반 불확실성 추정을 수행하고, 증강된 예측의 엔트로피를 계산한다.
• BraTS에서 요구하는 바와 같이, 0(가장 확신됨)에서 100(가장 불확실함)으로 정규화된 세 가지 불확실성 지도(전체 종양(WT), 종양 핵심부(TC), 강화 종양(ET))를 생성한다.
• 예측을 정밀화하고 특히 강화 종양(ET) 영역에서의 거짓 양성 수를 줄이기 위해 후처리를 적용한다.

실험 결과

연구 질문

• RQ1 불확실성 추정 기능을 갖춘 3D-CNN가 다중 모달 MRI에서 뇌 종양 분할의 신뢰도를 향상시킬 수 있는가?
• RQ2 테스트 시점 드롭아웃은 종양 분할 예측의 지식 기반 불확실성 추정에 어떻게 기여하는가?
• RQ3 데이터 증강은 저신뢰 영역에서의 확률 기반 불확실성 추정에 어느 정도 기여하는가?
• RQ4 불확실성 지도는 임상 의사결정 지원에서 거짓 양성 및 모호한 종양 경계를 효과적으로 강조할 수 있는가?
• RQ5 불확실성 지표(FTP 및 FTN 비율)는 BraTS 2020 검증 세트에서의 분할 성능과 어떻게 상관관계가 있는가?

주요 결과

• 후처리를 거친 V-Net 모델은 BraTS 2020 검증 세트에서 전체 종양(WT)의 DICE 스코어가 0.8463, 종양 핵심부(TC)는 0.7526, 강화 종양(ET)은 0.6179를 기록하였다.
• 검증 세트에서 ET의 하우스도르프 거리가 크게 증가하여 학습 시 31.18 mm에서 검증 시 49.74 mm로 증가하였으며, 이는 ET 분할에서 높은 거짓 양성 비율을 시사한다.
• 검증 세트에서 FTP(0.0118)와 FTN(0.0027) 비율이 낮게 나타나, 모델이 참 양성 및 참 음성 예측에 대해 여전히 확신을 가짐을 나타낸다.
• 불확실성 지도는 특히 종양 경계와 거짓 양성 영역에서 저신뢰 영역을 명확히 강조하여 임상 검토를 지원한다.
• 다중 불확실성 인식 전방 전파의 예측 평균화는 DICE 스코어를 약간 감소시켰지만, 불확실성 인식 기반의 정확도 향상을 증가시켰다.
• 모델은 특히 ET 분할에서 어려움을 겪었으며, 단일 거짓 양성 볼륨이 DICE 스코어를 0으로 만들고 하우스도르프 거리를 약 373 mm로 크게 상승시켜 평균 값에 심각한 영향을 주었다.