[论文解读] Uncertainty-based graph convolutional networks for organ segmentation refinement

Roger D. Soberanis-Mukul, Nassir Navab+1 more|arXiv (Cornell University)|Jun 5, 2019

Advanced X-ray and CT Imaging被引用 14

一句话总结

该论文提出了一种新颖的不确定性感知图卷积网络（GCN）框架，用于在CT影像中对CNN输出的器官分割结果进行优化。通过利用蒙特卡洛 dropout 识别高置信度和低置信度体素，构建了一个半监督图结构，其中高置信度预测作为标记节点；随后，GCN对不确定区域进行优化，使胰腺和脾脏的Dice分数相比基线U-Net预测分别提升了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中应用蒙特卡洛 dropout（MCDO），以估计模型不确定性并生成预测置信度图。
• 基于不确定性阈值，将体素划分为高置信度前景/背景区域以及低置信度区域（潜在假阳性/假阴性区域）。
• 利用不确定性感知的空间约束，在器官周围定义一个3D形状自适应的感兴趣区域（ROI）。
• 构建一个稀疏的、半标记的图结构：高置信度体素作为标记节点，低置信度体素作为未标记节点。
• 使用标记节点以半监督方式训练图卷积网络（GCN），以预测未标记节点的标签。
• 通过在训练好的GCN中对完整图进行推理，获得优化后的分割结果，节点表示融合了空间、强度和期望特征。

实验结果

研究问题

• RQ1 CNN的不确定性估计是否能有效提升对器官分割中假阳性和假阴性区域的识别能力？
• RQ2 半监督图学习是否能有效利用高置信度预测作为监督信号，对低置信度分割区域进行优化？
• RQ3 与基于CRF的基线优化方法相比，所提出的GCN优化方法在分割精度上表现如何？
• RQ4 在数据量较少的场景下，GCN优化策略是否优于CNN自身的期望预测（平均预测）？
• RQ5 该框架能否作为即插即用的后处理模块，在不微调基础CNN的前提下应用？

主要发现

• 与原始U-Net预测相比，所提出的GCN优化方法使胰腺分割的Dice分数提升了1%，脾脏分割提升了2%。
• 在NIH胰腺数据集上，该方法在Dice分数上优于最先进的基于CRF的优化方法。
• 在低数据场景（如仅10个训练体素）下，GCN优化方法的Dice分数优于CNN自身的期望预测，表明其对不可靠不确定性估计具有鲁棒性。
• 相对改进分析显示，在20个体素的胰腺中，有13个体素的Dice分数在GCN优化后优于或等于期望预测，尤其在模型训练不佳时表现更优。
• GCN能够有效减少CNN期望预测中由伪影引起的假阳性区域，证明其具备纠正模型错误的能力。
• 该方法兼容2D和3D CNN，并作为推理阶段的后处理模块运行，无需微调基础网络。