[论文解读] Globally Optimal Segmentation of Mutually Interacting Surfaces using Deep Learning

Hui Xie, Zhe Pan+6 more|arXiv (Cornell University)|Jul 2, 2020

Anomaly Detection Techniques and Applications参考文献 24被引用 4

一句话总结

该论文提出了一种深度学习框架，通过使用原始-对偶内点法（IPM）将基于图的最优表面分割与端到端可微优化相结合。通过参数化表面代价函数并利用U-Net主干网络进行学习，该方法实现了在OCT和IVUS图像中多个相互作用表面的全局最优分割，在显式建模表面相互作用的同时实现了高效的反向传播通过IPM层，取得了最先进精度。

ABSTRACT

Segmentation of multiple surfaces in medical images is a challenging problem, further complicated by the frequent presence of weak boundary and mutual influence between adjacent objects. The traditional graph-based optimal surface segmentation method has proven its effectiveness with its ability of capturing various surface priors in a uniform graph model. However, its efficacy heavily relies on handcrafted features that are used to define the surface cost for the "goodness" of a surface. Recently, deep learning (DL) is emerging as powerful tools for medical image segmentation thanks to its superior feature learning capability. Unfortunately, due to the scarcity of training data in medical imaging, it is nontrivial for DL networks to implicitly learn the global structure of the target surfaces, including surface interactions. In this work, we propose to parameterize the surface cost functions in the graph model and leverage DL to learn those parameters. The multiple optimal surfaces are then simultaneously detected by minimizing the total surface cost while explicitly enforcing the mutual surface interaction constraints. The optimization problem is solved by the primal-dual Internal Point Method, which can be implemented by a layer of neural networks, enabling efficient end-to-end training of the whole network. Experiments on Spectral Domain Optical Coherence Tomography (SD-OCT) retinal layer segmentation and Intravascular Ultrasound (IVUS) vessel wall segmentation demonstrated very promising results. All source code is public to facilitate further research at this direction.

研究动机与目标

• 解决在弱边界和相互依赖关系下，医学图像中多个相互作用表面分割精度降低的挑战。
• 克服传统基于图的方法在表面代价建模中依赖手工设计特征的局限性。
• 将深度学习的特征表示能力与凸优化的全局最优性相结合，用于多表面分割。
• 通过神经网络层实现IPM，使约束优化过程可微，从而实现端到端训练。
• 在真实世界医学影像任务（如SD-OCT和IVUS数据中的视网膜层和血管壁分割）中展示该方法的有效性。

提出的方法

• 使用深度神经网络（U-Net）对基于图的分割模型中的表面代价函数进行参数化，以学习图像相关的代价参数。
• 将多表面分割建模为一个约束凸优化问题，以强制实施表面相互作用约束并保证全局最优性。
• 将原始-对偶内点法（IPM）实现为可微层，以求解优化问题，从而实现通过优化过程的反向传播。
• 利用KKT条件和矩阵微积分，推导出IPM层相对于网络权重、输入特征和约束的闭式梯度。
• 使用PyTorch实现IPM层，支持GPU加速的批量处理，并对奇异矩阵采用伪逆处理。
• 将IPM层集成到深度学习流水线中，实现U-Net与优化组件的联合训练，支持完整的反向传播。

实验结果

研究问题

• RQ1 深度学习能否与凸优化有效结合，实现在医学图像中多个相互作用表面的全局最优分割？
• RQ2 如何使用深度神经网络对基于图的分割模型中的表面代价函数进行参数化并实现端到端学习？
• RQ3 在弱边界存在的情况下，可微优化层（如IPM）能在多大程度上提升多表面分割的鲁棒性和准确性？
• RQ4 所提出的方法能否在不同医学成像模态（如SD-OCT和IVUS）之间实现泛化？
• RQ5 通过约束显式建模表面相互作用，相比独立表面检测，能多大程度提升分割性能？

主要发现

• 所提方法在SD-OCT图像中的视网膜层分割任务上达到了最先进性能，显著优于基线U-Net和FCN模型。
• 该方法在IVUS图像中血管壁分割任务中表现出鲁棒性，即使在图像对比度差和边界弱的区域也表现良好。
• IPM层的集成实现了端到端训练，收敛稳定，每次前向传播最多仅需7次IPM迭代。
• 定量评估显示，与非优化和非交互感知基线相比，Dice分数更高，边界误差更小。
• 消融研究证实，建模表面相互作用显著提升了分割精度，尤其在模糊或低对比度区域。
• 所有代码和数据已在[26]处开源，促进了可复现性，并推动了可微优化在医学图像分割领域的进一步研究。