[论文解读] Joint Segmentation and Uncertainty Visualization of Retinal Layers in Optical Coherence Tomography Images using Bayesian Deep Learning

Suman Sedai, Bhavna Antony+2 more|arXiv (Cornell University)|Sep 12, 2018

Retinal Imaging and Analysis参考文献 18被引用 4

一句话总结

本文提出了一种贝叶斯全卷积DenseNet（BFC-DN），用于在OCT图像中实现视网膜层分割与像素级不确定性量化。通过结合认知不确定性与随机不确定性，利用Dropout变分推断及类别加权贝叶斯损失，该方法在分割精度方面达到最先进水平，并提升了对噪声的鲁棒性，不确定性图能有效突出不可靠的分割区域。

ABSTRACT

Optical coherence tomography (OCT) is commonly used to analyze retinal layers for assessment of ocular diseases. In this paper, we propose a method for retinal layer segmentation and quantification of uncertainty based on Bayesian deep learning. Our method not only performs end-to-end segmentation of retinal layers, but also gives the pixel wise uncertainty measure of the segmentation output. The generated uncertainty map can be used to identify erroneously segmented image regions which is useful in downstream analysis. We have validated our method on a dataset of 1487 images obtained from 15 subjects (OCT volumes) and compared it against the state-of-the-art segmentation algorithms that does not take uncertainty into account. The proposed uncertainty based segmentation method results in comparable or improved performance, and most importantly is more robust against noise.

研究动机与目标

• 为解决基于深度学习的视网膜层分割中缺乏不确定性量化的问题，该问题常将模型视为黑箱且输出未经验证。
• 提升分割模型在图像噪声或分布偏移情况下的鲁棒性，特别是在测试数据与训练数据不一致时。
• 开发一种方法，提供像素级不确定性图，以指导下游临床决策或主动学习，识别低置信度区域。
• 在不牺牲分割精度的前提下，将不确定性估计集成到全卷积架构中。
• 验证不确定性图是否与分割误差相关，并可用于标记需要人工审查或重新评估的区域。

提出的方法

• 将全卷积DenseNet（FC-DN）与密集块及跳跃连接相结合，用于端到端视网膜层分割。
• 通过Dropout变分推断引入贝叶斯推断，利用多次前向传播估计推理过程中的认知不确定性。
• 通过预测输出方差与分割logits并行，建模随机不确定性，实现对固有图像噪声的不确定性估计。
• 采用类别加权贝叶斯损失函数，结合预测不确定性与类别不平衡问题进行训练。
• 在推理过程中使用温度缩放技术（T），以校准不确定性估计，其中T > 1可实现更优的不确定性校准。
• 通过在启用Dropout的情况下进行多次前向传播并计算方差，生成不确定性图，其中高方差表示低置信度。

实验结果

研究问题

• RQ1 贝叶斯深度学习是否能通过提供不确定性估计，提升OCT中视网膜层分割的可靠性？
• RQ2 同时包含认知不确定性与随机不确定性的方法是否能提升模型对噪声和分布偏移的鲁棒性？
• RQ3 不确定性图是否能有效突出分割不准确的区域，特别是在噪声或分布外图像中？
• RQ4 与非贝叶斯最先进的分割模型相比，该方法在Dice分数和边界准确性方面表现如何？
• RQ5 不确定性量化是否能增强视网膜成像中的下游临床决策或主动学习流程？

主要发现

• 所提出的BFC-DN方法在GCL+IPL层上实现了0.97的Dice系数，在OPL和IS层上实现了0.91的Dice系数，优于非贝叶斯的FC-DN基线模型。
• 与Lang等人[12]提出的最先进方法相比，边界误差平均降低了30%，IS-OS边界误差最低达到2.1 μm。
• 在噪声测试条件下，BFC-DN的分割性能显著优于FC-DN，表现出对随机块噪声更强的鲁棒性。
• 模型生成的不确定性图与分割误差具有强烈的空域相关性，颜色越暖（数值越高）表示不确定性越高，准确度越低。
• BFC-DN的平均推理时间为每幅图像2.5秒（在Tesla K40上），由于需要T=10次前向传播进行不确定性估计，因此比FC-DN（300 ms）更慢。
• 该方法表明，不确定性量化可指导临床决策，并通过识别需要专家验证的区域，支持主动学习。