[논문 리뷰] Uncertainty-based graph convolutional networks for organ segmentation refinement

Roger D. Soberanis-Mukul, Nassir Navab외 1명|arXiv (Cornell University)|2019. 06. 05.

Advanced X-ray and CT Imaging인용 수 14

한 줄 요약

이 논문은 CT 영상에서 CNN 출력값을 보정하기 위한 새로운 불확실성 인식 그래프 컬러리션 네트워크(GCN) 프레임워크를 제안한다. 몬테 카를로 드롭아웃을 활용해 높고 낮은 신뢰도의 볼렉스를 식별하고, 높은 신뢰도 예측값을 레이블이 부여된 노드로 삼아 반감성 그래프를 구성한다. 이후 GCN는 불확실한 영역을 보정하여 기준 U-Net 예측 대비 측정 지표(Dice score)를 췌장에서 1%, 비장에서 2% 향상시킨다.

ABSTRACT

Organ segmentation in CT volumes is an important pre-processing step in many computer assisted intervention and diagnosis methods. In recent years, convolutional neural networks have dominated the state of the art in this task. However, since this problem presents a challenging environment due to high variability in the organ's shape and similarity between tissues, the generation of false negative and false positive regions in the output segmentation is a common issue. Recent works have shown that the uncertainty analysis of the model can provide us with useful information about potential errors in the segmentation. In this context, we proposed a segmentation refinement method based on uncertainty analysis and graph convolutional networks. We employ the uncertainty levels of the convolutional network in a particular input volume to formulate a semi-supervised graph learning problem that is solved by training a graph convolutional network. To test our method we refine the initial output of a 2D U-Net. We validate our framework with the NIH pancreas dataset and the spleen dataset of the medical segmentation decathlon. We show that our method outperforms the state-of-the art CRF refinement method by improving the dice score by 1% for the pancreas and 2% for spleen, with respect to the original U-Net's prediction. Finally, we discuss the results and current limitations of the model for future work in this research direction. For reproducibility purposes, we make our code publicly available.

연구 동기 및 목표

• 의료 CT 영상에서 형태의 변동성과 조직 유사성으로 인해 발생하는 잘못된 양성 및 잘못된 음성 결과를 해결하기 위해.
• 모델의 불확실성을 활용해 CNN 기반 분할의 신뢰성을 향상시키기 위해.
• 부피 분할의 후처리 전략으로 반감성 그래프 학습을 탐색하기 위해.
• 기본 CNN을 재학습하지 않고도 추론 시점에 작동하는 프레임워크를 개발하기 위해.
• 낮은 데이터 환경에서 불확실성이 중요한 상황에서의 성능 평가를 위해.

제안 방법

• 2D U-Net에 몬테 카를로 드롭아웃(MCDO)을 적용하여 모델 불확실성을 추정하고 예측 신뢰도 맵을 생성한다.
• 불확실성 임계값을 기반으로 볼렉스를 높은 신뢰도의 전경/배경 및 낮은 신뢰도(잠재적 FP/FN) 영역으로 분류한다.
• 불확실성 인식 공간 제약 조건을 활용해 기관 주변에 3D 형태 적응형 관심 영역(ROI)을 정의한다.
• 희박하고 반감성 레이블이 부여된 그래프를 구성한다: 높은 신뢰도 볼렉스는 레이블이 부여된 노드로, 낮은 신뢰도 볼렉스는 레이블이 없는 노드로 간주한다.
• 레이블이 부여된 노드를 활용해 반감성 방식으로 그래프 컬러리션 네트워크(GCN)를 훈련하고, 레이블이 없는 노드의 레이블을 예측한다.
• 학습된 GCN을 통해 전체 그래프를 평가하여 보정된 분할 결과를 도출한다. 이 과정에서 노드 표현에 공간, 강도 및 기대치 특징을 통합한다.

실험 결과

연구 질문

• RQ1 CNN에서 유도한 불확실성 추정이 장기 분할에서 잘못된 양성 및 잘못된 음성 영역을 식별하는 데 효과적인가?
• RQ2 높은 신뢰도 예측값을 감독으로 삼아 반감성 그래프 학습이 낮은 신뢰도 분할 영역을 효과적으로 보정할 수 있는가?
• RQ3 제안된 GCN 보정 방법은 CRF와 같은 기준 보정 기법 대비 분할 정확도에서 어떻게 비교되는가?
• RQ4 낮은 데이터 환경에서 GCN 보정 전략이 CNN 자체의 기대값(평균 예측)을 초월하는가?
• RQ5 이 프레임워크는 기본 CNN을 재학습하지 않고도 플러그인 형태의 후처리 모듈로 적용 가능한가?

주요 결과

• 제안된 GCN 보정 방법은 원본 U-Net 예측 대비 췌장 분할에서 DICE 스코어를 1% 향상시키고 비장 분할에서는 2% 향상시켰다.
• NIH 췌장 데이터셋에서, 이 방법은 최신 기술인 CRF 기반 보정 방법보다 DICE 스코어에서 뛰어난 성능을 보였다.
• 낮은 데이터 환경(예: 10개의 학습 영상)에서 GCN 보정은 CNN의 기대값보다 더 높은 DICE 스코어를 달성하여 신뢰도가 떨어지는 불확실성 추정에 대해 강건함을 입증했다.
• 상대적 향상 분석 결과, 20개의 췌장 영상 중 13개에서 GCN 보정이 기대값과 동일하거나 더 높은 DICE 스코어를 기록했으며, 특히 모델 학습이 부족한 경우에 두드러진 성능 향상을 보였다.
• GCN는 CNN 기대값에서 기인한 아티팩트로 인한 잘못된 양성 영역을 줄이는 데 성공하여 모델 오류를 수정할 수 있는 능력을 입증했다.
• 이 방법은 2D 및 3D CNN과 호환되며, 기본 네트워크를 재학습하지 않고도 추론 시점 후처리 모듈로 작동한다.