[論文レビュー] 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.

研究の動機と目的

• カーブ状セグメンテーションにおける短いブランチや結合部など、エラーが起こりやすい構造レベルのセグメントを強調することで、効率的な校正を促進する。
• 離散モース理論を用いてセグメンテーション尤度を構成要素となる構造に分解し、トポロジーを捉える。
• 構造変動をサンプリングして Probabilistic DMT による構造内不確実性をモデル化する。
• 構造間不確実性を、構造間の結合グラフ推論によって捉える。
• 2Dおよび3Dデータセット全体で、構造ごとの不確実性マップの改善と堅牢な性能を示す。

提案手法

• セグメンテーション尤度マップを離散モース理論を用いてモース骨格に分解し、安定多様体、サドル、極大値といった構造の集合を得る。
• 尤度マップを摂動し、誘導付き摂動と探索を実行して各構造につき複数の骨格をサンプリングし、構造内不確実性をモデル化する Probabilistic DMT (Prob. DMT) を導入する。
• 構造をノードとし、重なりを表すエッジを持つグラフを構築する。グラフニューラルネットワークを用いて、構造が真である確率と不確実性を共同で予測し、構造間の文脈を組み込む。
• 実データの重なりと共に、減衰型損失を用いた回帰目的で構造ごとの不確実性を監督し、明示的な不確実性ラベルなしで学習する。
• 構造別の不確実性ヒートマップと重ね合わせた改善セグメンテーションマップを生成する後処理を行い、構造中心の校正を可能にする。

実験結果

リサーチクエスチョン

• RQ1 トポロジー認識フレームワークは、画素単位ではなく個々の曲線状構造（ブランチ、接続部）のレベルで不確実性を識別・定量化できるか。
• RQ2 Probabilistic DMTは各構造内のばらつきをモデル化することで、構造内不確実性を効果的に捉えられるか。
• RQ3 すべての構造に対する結合グラフ推論は、構造間不確実性を正確にモデルし、構造レベルの不確実性マップを改善できるか。
• RQ4 提案手法は、人間のアノテータを高不確実性構造へ誘導することで、校正効率を実質的に向上させられるか。
• RQ5 このアプローチは、2Dおよび3Dデータセットとさまざまなバックボーンセグメンテーションモデルに対して頑健か。

主な発見

• 本手法は、2D/3Dデータセット全体で、不確実性の較正（ECE）とトポロジー感度指標において、ピクセル単位および構造単位のベースラインを上回る。
• Probabilistic DMT は決定論的 DMT に比べて構造内不確実性推定を大幅に改善し、ECEが低いことから示される。
• 構造間不確実性を、構造を独立に推論するよりも、構造間の結合グラフ推論を用いた GNN ベースの推論がよりよく捉える。
• この手法は、構造ごとの不確実性マップをより忠実に生成し、再較正時にセグメンテーションのバックボーン（UNet、DeepVesselNet、CS2-Net など）を改善する。
• 校正シミュレーションは、構造ごとの不確実性が前処理方法よりも不確実な構造へユーザーを効果的に導くことを示唆している。