[논문 리뷰] Topology-Aware Uncertainty for Image Segmentation

Saumya Gupta, Yikai Zhang외 3명|PubMed|2023. 06. 09.

Topological and Geometric Data Analysis인용 수 9

한 줄 요약

이 논문은 Probabilistic Discrete Morse Theory(Prob. DMT)를 사용한 곡선형 구조 분할의 구조적 불확실성 추정과 intra- 및 inter-구조적 불확실성을 포착하기 위한 공동 그래프 기반 추론을 제시하여 기존 방법들보다 구조 수준의 불확실성 맵을 향상시킨다.

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 데이터셋 전반에서 향상된 구조별 불확실성 맵과 견고한 성능을 시연한다.

제안 방법

• 분할 가능도 맵을 Discrete Morse Theory를 사용한 모스 골격으로 분해하여 구조들의 모음을 얻는다(안정 매니폴드, 안장점, 극대점).
• 가능도 맵을 섭동하고 가이드된 섭동-도보를 수행하여 구조당 여러 골격을 샘플링하는 방식으로 Probabilistic DMT(Prob. DMT)를 도입하여 구조 내 불확실성을 모델링한다.
• 구조를 노드로, 중첩(overlaps)을 표현하는 엣지를 가지는 그래프를 구축하고, Graph Neural Network를 사용해 구조가 참일 확률과 그 불확실성을 공동으로 예측하며 구조 간 맥락을 반영한다.
• 실제 겹침(ground-truth overlap)과 함께 구조별 불확실성을 감독하기 위해 감쇠 기반 손실을 포함한 회귀 목표를 사용하되 명시적 불확실성 라벨은 사용하지 않는다.
• 구조별 불확실성 히트맵과 겹친 개선된 분할 맵을 생성하기 위한 후처리를 수행하여 구조 중심의 교정이 가능하도록 한다.

실험 결과

연구 질문

• RQ1 토폴로지 인식 프레임워크가 픽셀 단위가 아니라 개별 곡선형 구조(가지, 연결) 수준에서 불확실성을 식별하고 정량화할 수 있는가?
• RQ2 Probabilistic DMT가 각 구조 내의 변이를 모델링함으로써 구조 내 불확실성을 효과적으로 포착하는가?
• RQ3 모든 구조에 대한 공동 그래프 기반 추론이 구조 간 불확실성을 정확하게 모델링하고 구조 수준의 불확실성 맵을 향상시키는가?
• RQ4 제안된 방법이 높은 불확실성을 보이는 구조로 인간 주석자의 교정 작업을 지시함으로써 교정 효율을 의미 있게 향상시킬 수 있는가?
• RQ5 이 방법이 2D 및 3D 데이터셋과 다양한 백본 분할 모델에서 견고한가?

주요 결과

• 본 방법은 2D/3D 데이터셋에서 불확실성 보정(ECE) 및 위상 민감 지표에서 픽셀-보다 및 구조-기반 베이스라인보다 우수하다.
• Probabilistic DMT는 결정론적 DMT에 비해 구조 내 불확실성 추정에 현저한 향상을 보였으며, 더 낮은 ECE로 나타난다.
• 구조들에 대한 GNN 기반의 공동 추론은 독립적 구조 추론보다 구조 간 불확실성을 더 잘 포착한다.
• 이 방식은 구조별 불확실성 맵의 충실도를 높이고 재 보정 시 분할 백본을 개선하며, 다수의 백본(UNet, DeepVesselNet, CS2-Net)에서 향상된다.
• 교정 시뮬레이션은 구조별 불확실성이 사용자를 과거 방법보다 불확실한 구조로 더 효과적으로 안내함을 시사한다.