[论文解读] Topology-Aware Uncertainty for Image Segmentation

Saumya Gupta, Yikai Zhang+3 more|PubMed|Jun 9, 2023

Topological and Geometric Data Analysis被引用 9

一句话总结

本文提出使用 Probabilistic Discrete Morse Theory (Prob. DMT) 的结构级不确定性估计来对曲线状分割进行不确定性评估，并通过基于图的联合推断捕获结构内和结构间的不确定性，从而在现有方法之上提升结构级不确定性图。

ABSTRACT

Segmentation of curvilinear structures such as vasculature and road networks is challenging due to relatively weak signals and complex geometry/topology. To facilitate and accelerate large scale annotation, one has to adopt semi-automatic approaches such as proofreading by experts. In this work, we focus on uncertainty estimation for such tasks, so that highly uncertain, and thus error-prone structures can be identified for human annotators to verify. Unlike most existing works, which provide pixel-wise uncertainty maps, we stipulate it is crucial to estimate uncertainty in the units of topological structures, e.g., small pieces of connections and branches. To achieve this, we leverage tools from topological data analysis, specifically discrete Morse theory (DMT), to first capture the structures, and then reason about their uncertainties. To model the uncertainty, we (1) propose a joint prediction model that estimates the uncertainty of a structure while taking the neighboring structures into consideration (inter-structural uncertainty); (2) propose a novel Probabilistic DMT to model the inherent uncertainty within each structure (intra-structural uncertainty) by sampling its representations via a perturb- and-walk scheme. On various 2D and 3D datasets, our method produces better structure-wise uncertainty maps compared to existing works. Code available at https://github.com/Saumya-Gupta-26/struct-uncertainty.

研究动机与目标

• 通过突出易出错且处于结构级的分段（例如短分支、连接）来提高校勘效率。
• 通过离散 Morse 理论将分割似然分解为组成结构以捕捉拓扑。
• 通过对结构变体进行采样，用 Probabilistic DMT 建模结构内不确定性。
• 通过对结构进行联合图推断来捕捉结构间不确定性。
• 展示在 2D/3D 数据集上改进的结构级不确定性图和鲁棒性能。

提出的方法

• 利用 Discrete Morse Theory 将分割似然映射分解为 Morse 骨架，以获得一组结构（稳定流形、鞍点、极大值）。
• 通过扰动似然映射并执行引导性扰动-走动来为每个结构采样多个骨架，建模结构内不确定性，从而引入 Probabilistic DMT (Prob. DMT)。
• 构建一个以结构为节点、以重叠关系为边的图；使用图神经网络对结构为真概率及其不确定性进行联合预测，融入结构间上下文。
• 使用带衰减的回归目标函数来监督结构级不确定性，同时结合真实重叠信息进行训练，无需显式的不确定性标签。
• 后处理以生成结构级不确定性热图以及叠加的改进分割图，支持面向结构的校勘。

实验结果

研究问题

• RQ1 一个具拓扑感知的框架是否能够在单个曲线结构（分支、连接）级别而非像素级别识别并量化不确定性？
• RQ2 Probabilistic DMT 是否通过对每个结构内的变异进行建模来有效捕捉结构内不确定性？
• RQ3 对所有结构进行联合图推理是否能准确建模结构间不确定性并提升结构级不确定性图？
• RQ4 所提出的方法是否能够通过将人工标注者引导到高不确定性结构来实质性提升校勘效率？
• RQ5 该方法在 2D/3D 数据集以及不同主干分割模型下是否具有鲁棒性？

主要发现

• 我们的方法在不确定性校准（ECE）和拓扑敏感指标上超过像素级和结构级基线，适用于 2D/3D 数据集。
• Probabilistic DMT 在结构内不确定性估计上显著优于确定性 DMT。
• 基于 GNN 的结构联合推断比独立结构推理更好地捕获结构间不确定性。
• 该方法在再标定后，在多个骨干网络（UNet、DeepVesselNet、CS2-Net）上，产生更真实的结构级不确定性图并改善分割骨干。
• 校勘仿真表明结构级不确定性比先前方法更有效地引导用户关注不确定结构。