Title: Topology-Aware Uncertainty for Image Segmentation
Venue: NeurIPS 2023
Accepted: yes
Keywords: Topological Representation

Discrete Morse Theory

Structural Uncertainty

Image Segmentation
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

Topology-Aware Uncertainty for Image Segmentation

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 perturband-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

1 Introduction

Curvilinear segmentation is an essential initial step in various medical and non-medical applications, involving the precise extraction of fine-scale structures, such as blood vessels, nerves, and other elongated objects [21, 33]. For example, extraction of retinal vasculature is an essential precursor to understanding disease progression and assessing therapeutic effects [15]. In civil engineering, road network and railway track segmentation can support urban planning and transportation system optimization [45]. Despite the success of deep learning [5, 6, 25, 37], automatic segmentation of thin structures remains challenging due to their relatively low visibility and complex topology. Existing segmentation methods often make topological errors such as broken connections or missing branches.

As a cost-efficient alternative, in many applications, one employs semi-automatic techniques, e.g., iterative proofreading by human annotators [23]. On the other hand, proofreading of complex finescale structures can be extremely time-consuming [53]. This necessitates a better strategy to direct the annotators' attention towards locations that are more error-prone. Following the classic active learning principle [18, 55, 61], we may estimate the uncertainty [19] and concentrate on the locations where a neural network is the least certain.

Despite many existing studies on segmentation uncertainty [12, 51, 60], most existing uncertainty estimation methods do not apply to curvilinear structure segmentation. Existing methods typically generate pixel-wise uncertainty maps. Such maps highlight pixels along the boundary of all structures,

a) Input image

c) UNet

d) PHiSeg

f) Ours

b) GT

10

1.0

- 0.8

- 0.8

- 0.8

- 0.6

- 0.6

- 0.6

0.4

- 0.4

- 0.4

- 0.2

0.2

- 0.2

0.0

0.0

0.0

regardless of their width or thickness (see Fig. 1(d)). This offers limited information for human annotators; a desirable uncertainty map should instead highlight the error-prone 'structures', e.g., small vessels/branches or short stretches of roads that tend to be disconnected or missed.

In this paper, we propose a new topology-aware uncertainty estimation method that highlights errorprone structures as a whole (such as in Fig. 1(f)). By highlighting structures with high uncertainty, our method empowers annotators to accept or reject/correct structural proposals efficiently, thus streamlining the proofreading process. To capture the uncertainty of a given segmentation network's prediction at a structural level, we require the realization of two key components: a) decompose the prediction into a set of constituent structures, including false positives and false negatives, and b) estimate uncertainties of all the structures. Furthermore, we need to consider two types of structural uncertainty, intra-structural and inter-structural. The intra-structural uncertainty of a structure is due to its intrinsic composition, e.g., geometry, intensity, and the segmentation network's confidence. The inter-structural uncertainty is more contextual; it is due to interactions between neighboring structures. Our method explicitly models the two types of uncertainty.

Given a segmentation model, our method uses its likelihood map (Fig. 2(i)) plus the input image to estimate structural uncertainties. First, we obtain a structural decomposition of the prediction, i.e., a collection of one-pixel-wide pieces/structures (see Fig. 2(ii)). Each structure represents a potential branch/connection according to the segmentation model. We employ the principles of the classic discrete Morse theory (DMT) [14, 43]. Intuitively, DMT treats the likelihood function as a terrain function and extracts landscape features, e.g., mountain ridges or valleys, as structures. Note we are capturing all possible structures visible in the likelihood map, including the ones that do not appear in the segmentation due to low probability.

Next, we estimate uncertainties for all these structures. Existing uncertainty estimation approaches often sample multiple hypotheses and calculate the variance across them [17, 34, 69]. However, this principle is not feasible in our problem; with N structures, the space of all their combinations is of exponential size ( 2 N ). Sampling from such a space is very challenging. An alternative is to make independent uncertainty estimation on each structure. However, this is also suboptimal as it ignores inter-structural uncertainty. Fig. 2(vi) shows three false positive structures. Treating all structures independently will incorrectly assign the horizontal structure with low uncertainty (see Fig. 2(vii)).

We propose a joint inference model (i.e., a graph neural network [59]) to jointly predict uncertainties on all the structures. This joint inference framework avoids explicit enumeration/sampling over the exponential size space of hypotheses. It also takes into account the inter-structural uncertainty. In Fig. 2(viii), our method correctly assigns high uncertainties to all three false positive structures. To supervise the training process, we use the attenuation loss proposed in [30] to learn the uncertainty. As there is no 'uncertainty label', it is implicitly learned from the loss function.

An important contribution in this paper is a novel Probabilistic DMT modeling the intra-structural uncertainty, due to the geometry, topology, as well as the segmentation model's confidence. The classic DMT cannot capture the inherent uncertainty of a structure; it produces a single representation for the structure, i.e., a one-pixel-wide skeleton. However, this skeleton is computed in a deterministic manner. Such deterministic computation is rigid and fails to capture possible variations of a structure.

Instead, our Probabilistic DMT (Prob. DMT) represents each structure as sample skeletons drawn from an underlying generative process guided by the likelihood map (the skeleton resulting from the

i) Likelihood map

iii) DMT vs. GT

ii) Morse skeleton

vi) Likelihood map

viii) w/ inter

vii) w/o inter

overconfident FP

low confidence FP

v) Input image

0.6

5.0

Figure 2: (a) Intra-structural uncertainty: In the likelihood map (i), we highlight a false negative (FN) structure missed by the segmentation network. In (ii), we show the skeletons representing structures from classic DMT. In (iii), we highlight the GT (green) and the incorrect DMT skeleton (red) for the FN structure. In (iv), we show skeleton samples by Prob. DMT (blue); (b) Inter-structural uncertainty ( inter ): a retinal image with very weak signal (v). The likelihood map (vi) shows three potential false positive ( FP ) structures, two vertical and one horizontal. Without inter-structural uncertainty (vii), the horizontal structure has high confidence. With inter-structural uncertainty (viii), the horizontal structure gets higher uncertainty influenced by the two vertical structures.

original DMT being one of the samples). As illustrated in Fig. 2(iii), the original DMT generates a skeleton that significantly deviates from the true structure due to the uncertainty inherent in the likelihood. In contrast, Prob. DMT effectively captures potential variations (Fig. 2(iv)), offering valuable insights into the impact of uncertainty on the structural composition. The greater the variation, the greater the intra-structural uncertainty. During training, we repeatedly sample skeletons for each structure, and feed the samples to the joint inference model for uncertainty prediction. Indeed, as shown in Fig. 2(iv), sampling multiple skeletons also leads to a better chance of uncovering the true structure. This observation inspires us to use the uncertainty estimation method to re-calibrate the original segmentation model and achieve even better segmentation performance.

We note that the method in [28] (which we refer to as Hu et al.) also used DMT to decompose structures and estimate their uncertainty. However, this method used the classic DMT to deterministically generate skeletons, and thus failed to model intra-structural uncertainty. Furthermore, instead of joint inference, the method sampled from all configurations; and to reduce the computational burden, it pruned the exponential-sized configuration space using a saliency measure called persistence [9, 64, 71]. The pruning was very coarse, thus resulting in suboptimal uncertainty estimation. As illustrated in Fig. 1(e), Hu et al. produces overconfident maps; most structures, including many false negatives and false positives, are assigned zero uncertainty. In contrast, our method produces much better uncertainty estimates (Fig. 1(f)), owing to the proper modeling of both intra-structural and inter-structural uncertainties.

We summarize the main contributions of this paper as follows:

We propose a novel method to estimate the uncertainty of a given segmentation network at a structural level .

We propose Probabilistic DMT, a probabilistic method to generate structural variations and to capture intra-structural uncertainty.

We propose a joint prediction model on all the structures in order to capture inter-structural uncertainty.

Empirical evaluation shows that our method achieves much better uncertainty estimates on both 2D and 3D datasets, outperforming existing methods.

2 Related work

Topology-guided image segmentation. Several works focus on maintaining the correct connectivity or topology of thin structures. Topology-aware loss functions [46, 62, 7, 27, 72, 22, 26] impose per-pixel constraints to improve topological integrity. Discrete Morse theory has also been used to improve the topological awareness of segmentation networks [29, 9, 11, 56, 71, 3]. These approaches use topological tools to improve segmentation at a pixel level, which is a weaker constraint compared to the structural level. In contrast, our method performs joint reasoning directly over the structures.

Uncertainty quantification. In recent years, there has been significant work on uncertainty quantification (UQ) of deep neural networks [1, 20, 36]. Here we review UQ techniques tailored for semantic segmentation. Pixel-wise uncertainty: Semantic segmentation is a per-pixel classification

Segmentation

network

Input image

Likelihood

Structure-wise

uncertainty

Probabilistic

DMT

...

...

...

Structural variations

Graph construction

Joint

reasoning

Likelihood

Intra-structural uncertainty

Inter-structural uncertainty

GT

Seg. map

External uncertainty quantification framework

Fe

0.6

4

- 0.2

task and naturally most UQ methods produce per-pixel uncertainty estimates. In [30], the authors propose a Bayesian framework using MC dropout [16] and a learned loss attenuation to respectively capture model and data uncertainty. Recent methods have turned to generative models to generate multiple hypotheses, and the per-pixel variance across the hypotheses is treated as uncertainty. Some works in this direction are an ensemble of M networks [34], a single network with M heads [58], Prob.-UNet [32], and PHiSeg [4]. Prob.-UNet integrates a conditional variational autoencoder [63] with UNet [57], generating multiple hypotheses via latent variable sampling. PHiSeg extends this by introducing latent variables at every UNet level, thereby producing more diverse samples. Structurewise uncertainty: Methods such as [41, 60] compute structure (volume) uncertainty by averaging over the pixel-wise uncertainty estimates. The method closest to ours is Hu et al. [28]. It is a generative model derived from Prob.-UNet where the latent variable has meaning in topology (specifically, a global persistence threshold). This threshold severely limits the structure space, overlooking several false positive/negative structures. Thus they tend to produce overconfident uncertainty estimates.

3 Method

Given a trained segmentation network, our goal is to capture the uncertainty of its prediction at a structural level. Note that we do not modify the network in any way; instead, we propose an external module that reasons the uncertainty of each structure in the segmentation. Fig. 3 provides an overview of our method. Let F θ denote the trained segmentation network, and M ϕ denote our proposed external uncertainty quantification framework. M ϕ takes as input the likelihood map of F θ and the input image. It generates a set of structures, and estimates an uncertainty value for each of them. During training, M ϕ is trained by comparing with the ground truth (GT) annotation.

M ϕ consists of two primary modules to capture intra-structural and inter-structural uncertainty. The first module, Probabilistic DMT (Prob. DMT), generates structures based on the likelihood map. For each structure, it samples a set of skeletons representing different variations. Details are provided in Sec. 3.1. The second module jointly predicts the uncertainties of all the structures. At each training iteration, it takes one sample skeleton for each structure, plus the likelihood map and input image, as input. Details are described in Sec. 3.2. Throughout the sections, we consider one data sample ( x, y ) where x is an input image and y is the segmentation GT. The likelihood map is f = F θ ( x ) .

3.1 Modelling the structural space

In this section, we first describe how DMT obtains the constituent structures of a likelihood map. Then we propose our Prob. DMT formulation to capture intra-structural uncertainty.

Discrete Morse theory. Consider the likelihood map f generated from the segmentation network F θ . We wish to decompose f into a set of structures, capturing not only the salient structures but also the faint ones. In the segmentation map, salient and faint structures broadly correspond to true positive and false negative structures. In Fig. 4(b), we highlight the false negative (FN) structures. These structures are missed in the segmentation, but will be captured by DMT (Fig. 4(d)).

DMT treats the likelihood map f as a terrain function, decomposing it into a Morse complex consisting of critical points, paths connecting them, patches in between paths, and volumes enclosed by patches (for 3D images). Critical points are locations w with zero gradients ( ∇ f ( w ) = 0 ), i.e., minima, maxima, or saddle points. Paths, called V-paths , are routes connecting critical points via the non-critical ones. A V-path connecting a saddle point to a maxima is called a stable manifold. These stable manifolds are the underlying terrain's mountain ridges, and delineate structures of interest. In Fig. 4(c), we show the locations of saddles and maximas in the Morse complex, and in Fig. 4(d), we show the union of all the stable manifolds connecting them. In this paper, we only focus on the zero- and onedimensional Morse structures, i.e., the union of all stable manifolds and their associated saddle and maxima. We call the collection of such structures the Morse skeleton.

By default, DMT generates stable manifolds in a completely deterministic manner, failing to take into account the intra-structural uncertainty in the likelihood f . Therefore, these stable manifolds may fail to correctly delineate the true structure, as shown in Fig. 5.

Probabilistic DMT. To account for the inherent uncertainty, we explicitly model the structure as a collection of samples from an underlying generative process. The skeleton from the original DMT is just one possibility out of many. The method is achieved via a perturb-and-walk algorithm, in which we iteratively perturb the likelihood map, and regenerate the skeleton.

The rationale is that the likelihood map is a weighted aggregation of all possible skeleton representations. To inverse the aggregation and recover these skeletons is challenging. Instead, we follow the classic perturb-and-map principle , which was used to efficiently sample from a complex discrete graphical model distribution [52, 24, 35]. We randomly perturb the likelihood function. For each perturbed likelihood, we compute a skeleton as a sample. See Fig. 5 for an illustration. The sampled skeletons will reflect the uncertainty properly. For a structure that is less salient in the likelihood map, the sample skeletons will have large variations, generating a large uncertainty. For a salient structure in

GT

DMT:

Prob. DMT:

region of

structural

uncertainty

Likelihood

Prob. DMT:

c m

c s

intermediate steps

(#1)

(#2)

the likelihood map, the sample skeletons will be less variant, resulting in a low uncertainty.

Assume a given likelihood function f and one of its structures, represented by a V-path e connecting a saddle-maximum pair ( c s , c m ) . We generate a sample skeleton of the structure by first perturbing the likelihood with random noise. Next, we generate a path connecting c s and c m . Recall in the original DMT, the skeleton is generated by following the mountain ridge. In other words, we start from the saddle point, and 'walk' towards the maximum. At every step, we always walk to the neighboring pixel with the highest likelihood value. In Prob. DMT, we follow the same principle on the perturbed likelihood. However, the noisy perturbation of likelihood can cause the path to grow astray. Therefore, we additionally apply a distance-based regularizer to guide the walk towards the target c m . We describe the process in detail below.

Let e denote the structure obtained by following the V-path between ( c s , c m ) in the original DMT. In order to generate its sample skeleton ˆ e , we first draw a likelihood f n from a distribution centered on f

a) Likelihood map

b) Seg. map

d) Morse skeleton

c) Critical points

as f n ∼ f + r . This process is independent of the perturbation model r used, and we use a Gaussian model in this work. As the variance of the Gaussian model is unknown, we use Bayesian probability theory to sample the variance from the Inverse Gamma distribution (its conjugate prior [49]).

Once we obtain f n , we regenerate the path between ( c s , c m ) . We take inspiration from random walk [38] as well as probability regularized walk [47] to generate the variant structure ˆ e from f n . Our walk algorithm continuously grows ˆ e starting from c s and ending at c m , one pixel at a time. The algorithm considers both the terrain f n and the distance to the destination c m to ensure path completeness. During the walk, given the current pixel location c , the next location c ′′ is chosen as c ′′ = argmax ( Q ( c ′ )) , where, c ′ ∈ neighborhood ( c ) 2 and Q ( c ′ ) = γQ d ( c ′ ) + (1 -γ ) f n ( c ′ ) , and, Q d ( c ′ ) = 1 ∥ c m -c ′ ∥ 2 . We begin with c := c s and continue in this manner c := c ′′ till we reach c m 3 . In Fig. 5, we show a deterministic structure obtained from DMT along with sample variations produced by our method. We demonstrate the intermediate steps in the algorithm: the red arrow denotes the next step without considering the distance regularizer Q d , while the green arrow denotes the next step using our formulation Q . Notice how only considering f n without Q d can prevent the path from reaching c m . We thus require Q d to guide the path to completeness.

The structure ˆ e is a different realization of e , making each run of the Prob. DMT a stochastic one. We are thus able to explicitly model the structures as samples from a probability distribution. We also note that DMT is a special case of Prob. DMT when r = 0 and γ = 0 . In practice, with some probability, we consider the original structure e from DMT over generating its variant ˆ e . Specifically, following a Bernoulli distribution, with a small probability u we retain e , while with probability 1 -u we sample its variant ˆ e using the perturb-and-walk algorithm outlined above. This process is done separately and in parallel for every structure. The structures taken together form a Morse skeleton. The output of Prob. DMT is effectively one sample skeleton from the space of Morse skeletons. We provide further information regarding Prob .DMT in the Appendix: A outlines the pseudocode, B discusses the hyperparameters, and C reports its computational complexity.

3.2 Joint estimation of structural uncertainty

The Prob. DMT module gives us a set E of structures. Our final step is to jointly reason about the uncertainty of all of them. To achieve this, we use a regression network that takes as input each structure e ∈ E , and outputs whether it is a true positive and the uncertainty of F θ in predicting it.

Details of the network. Structures interact with each other in the image space and are not independent. During uncertainty estimation, it is therefore crucial to consider their spatial context, i.e., interstructural uncertainty. Hence, we use Graph Neural Networks (GNN) [59], specifically Graph Convolution Networks (GCN) [31], to jointly reason about the structures and capture the high-order spatial interactions. In the graph, each node represents a structure, and edges between nodes exist when corresponding structures have non-zero overlaps (typically at endpoints). The input feature vector for each node is constructed as shown in Fig. 6. For every structure, we first concatenate [ x c , f c , m ] , where x c comes from the original input x ; f c from the likelihood map f (not f n ); and m is a binary map indicating the presence of the structure. These x c , f c , m are smaller crops/bounding boxes centered on the structure. After passing them through convolution blocks, we apply channelwise pooling to obtain a fixed-length feature vector for training. We further concatenate the persistence value of the saddle point associated with the original DMT structure (aka stable manifold). Persistence value (from persistent homology [13]) is defined as the difference of function (likelihood) values of 2 critical cells (saddle-maxima pair). It captures the importance of a structure, thus making it a valuable feature in our framework. Note that we do not use the perturbed f n from the Prob. DMT method when constructing the feature vector.

Training the network. We train the regression network using the attenuation loss proposed in [30]. As there are no labels to learn uncertainty, it is implicitly learned during regression optimization. We fix a Gaussian likelihood, and so variance ˆ δ 2 is used as a measure of uncertainty. The network's head is split into two - to predict ˆ p ( e ) of being a true positive structure and its associated uncertainty ˆ δ 2 e . For numerical stability, we actually predict the log variance s e = log ˆ δ 2 e . The training loss is given in Eq. 1. The structures that we obtain from Prob. DMT may not always fully overlap with the true GT structures, that is, some structures may only have partial overlap. We thus do not impose any hard

Convolution

blocks

Persistence value

input feature vector

Channel-wise

max pooling

...

...

constraints in Eq. 1, instead, z e is a soft label, and is given by: z e = ( ∑ y ⊙ m ) / ( ∑ m ) , where y is the GT and ⊙ is the Hadamard product. This value simply represents the proportion of the structure that overlaps with the GT, i.e., the fraction of the structure that is a true positive.

In [30], ˆ δ 2 denoted the pixel-wise uncertainty of the framework's input. In our setting, the input to our framework is f = F θ ( x ) , and so ˆ δ 2 is modeled to capture the structure-wise uncertainty inherent in data x and model F θ . Training ˆ δ 2 in this manner ensures that the network does not trivially predict high or low uncertainty, rather, predicts an uncertainty estimate that is dependent on the input.

3.3 Proposed module M ϕ

For Eq. 1 to hold, we require M ϕ to be a probabilistic network. We already show in Sec. 3.1 our formulation for Prob. DMT. Additionally, the regression network is also probabilistic as we use MC dropout [16].

Inference procedure. We take T runs of M ϕ and compute the uncertainty as the mean ¯ δ 2 e = 1 T ∑ T t =1 ( ˆ δ 2 e ) t . We similarly obtain ¯ p ( e ) from ˆ p ( e ) . In Fig. 7, we illustrate the post-processing steps to obtain the structure-wise uncertainty heatmap. First, we obtain maps ¯ p = ∪ ¯ p e and ¯ δ 2 = ∪ ¯ δ 2 e having the same spatial resolution as the input x . We then binarize ¯ p , and overlay it onto the segmentation map obtained from F θ . We do this because Prob. DMT

assign via

shortest distance

skeletons from

from

+

final uncertainty

heatmap

overlaid map

gives us one-pixel wide skeleton structures but we need to recover the structure thickness. Next, we use shortest distance to assign uncertainty values from ¯ δ 2 to the pixels in the overlaid map. The shortest distance uses paths only along the foreground pixels. In Fig. 7 we show how we obtain the final uncertainty heatmap from the skeleton heatmap. We also note that the overlaid map is an additional output of our method: it is an improved segmentation map that can be used instead of the one obtained by F θ . We provide more details in Appendix D.

4 Experiments

Datasets. We evaluate our method on four datasets: DRIVE [65], ROSE [40], ROADS [44] and PARSE 2022 Grand Challenge [39, 70]. The DRIVE dataset contains 2D retinal vasculature; ROSE is a 2D retinal OCTA (Optical Coherence Tomography Angiography) segmentation dataset, ROADS is a large non-medical dataset containing aerial images, and PARSE contains 3D CT scans of pulmonary arteries. We further describe the datasets and the data splits in Appendix E.

Baselines. We broadly split our comparison baselines into three types: a) Standard vessel segmentation methods: UNet [57], DeepVesselNet [68], and CS 2 -Net [48]; b) Pixel-wise uncertainty estimation methods: Prob.-UNet [32], and PHiSeg [4]; c) Structure-wise uncertainty estimation method: Hu et al. [28]. Implementation details are provided in Appendix F.

Evaluation metrics. To evaluate the quality of uncertainty quantification, we use Expected Calibration Error (ECE) [50] and Reliability Diagrams (RD) [8]. Furthermore, we also evaluate the segmentation on metrics such as DICE [73], clDice [62], ARI [2], VOI [42], Betti Number error [27] and Betti Matching error [66]. We include clDice, ARI, VOI, Betti Number and Betti

Input image

UNet

PHiSeg

Ours

GT

UNet

Ours

GT

0.6

9.2

LO

0.8

0.6

0.8

0.6

ReliabilityDiagram(DRIVE)

Reliability Diagram (ROADS)

ReliabilityDiagram(PARSE)

1.0

0.8

Hu et. al.

Prob.UNet

PHiSeg

1.0

0.8

Prob.UNet

PHiSeg

Hu et. al.

0.8

1.0

Ours

Prob.UNet

PHiSeg

Hu et. al.

1.0

0.8

Prob.UNet

PHiSeg

Hu et. al.

0.6

Ours

0.6

Ideal

0.6

Ours

accura

Ideal

en

Ideal

ccura

Ideal

0.4

吊

0.4

0.4

e

0.4

0.2

0.2

0.2

0.2

0.0

0.0

0.0

0.0

0.0

0.2

0.4

confidence

0.6

0.8

1.0

0.0

0.2

0.4

confidence

0.6

0.8

1.0

0.0

0.2

0.4

confidence

0.6

0.8

1.0

0.0

0.2

0.4

confidence

0.6

0.8

1.0

Matching as they are topology-based metrics and hence are sensitive to the performance on thin structures. Detailed definitions are present in Appendix G.

4.1 Results

Tab. 1 shows the quantitative results against uncertainty methods, and Tab. 2 shows the quantitative results on different backbone architectures. We show the respective qualitative results in Fig. 8 and Fig. 10. In Fig. 9, we plot the Reliability Diagrams. We also perform the unpaired t-test [67] (95% confidence interval) to determine the statistical significance. Each table reports the mean and standard deviations for every metric, with statistically significant better performances in bold and numerically better (but not significant) performances in italics. For all the probabilistic methods, the average of five runs was used. For our method, we generated the structure-wise uncertainty estimates and the segmentation map by following the steps outlined in the 'Inference procedure' in Sec. 3.3. Due to space constraints, results on two metrics Betti Number and Betti Matching are reported in Appendix H. We discuss the remaining performances below.

Performance of uncertainty estimation. Tab. 1 shows that our method outperforms others on both ECE and segmentation metrics. Fig. 9 displays RDs, with our method following the ideal line much closely compared to others. This is because we explicitly model the distribution of the structures, thereby quantifying the uncertainty of the segmentation network. In Fig. 8, we also see that our method generates better fidelity structure-wise uncertainty maps compared to Hu et al. Our heatmaps

Input image

CS 2 -Net

CS 2 -Net+Ours Ours (uncertainty)

GT

UNet

UNet+Ours

Ours (uncertainty)

assign non-zero uncertainty to several false positives/negatives in the backbone UNet's outputs. This is because we reason about every structure while Hu et al. limits the structure space via pruning.

Performance over different backbones. Tab. 2 and Fig. 10 show that our method is backboneagnostic. It consistently improves the segmentation quality and produces high fidelity uncertainty maps for each of the underlying networks. This validates the practical applicability of our method.

Performance of proofreading. One of the motivations of this work is to streamline the proofreading process. Structure-wise uncertainty can be used as a guide, with a user having to simply accept/reject a structure. We conduct experiments on the ROSE dataset and simulate user interaction with our method and Hu et al.'s. The user is given each method's final segmentation map, and inspects structures in decreasing order of uncertainty (till 0 . 5 ). Each uncertain structure is subjected to a yes/no decision, which is denoted as one 'click'. The results are in Fig. 11. Our findings are consistent with the observation that Hu et al. assigns zero uncertainty to many structures; thus their margin of improvement is limited and saturates quickly.

Proofreading:clDicevs#clicks

Ours

0.72

e

n

0.70

cIDi

0.68

0.66

0.0

2.5

5.0

7.5

10.0

12.5

Numberofclicks

4.2 Ablation studies

To demonstrate the efficacy of the proposed method, we conduct ablation studies of the different components in our pipeline, as well as check the effect of changing hyperparameter values.

We also include ablation studies on the dimensionality of the input feature vector, and size of the crops/bounding boxes (which we report in Appendix I). All analyses are on the DRIVE dataset using UNet [57] as the backbone.

Ablation of different modules. Weconduct ablation studies on both parts of our framework: structure generation (DMT vs Prob. DMT), and regression network (GNN vs Multi-layer Perceptron (MLP)). The results are in Tab. 3. Prob. DMT results in a sharp improvement in ECE compared to the original DMT; this supports our assertion that Prob. DMT models intra-structural uncertainty. Similarly, using GNN over MLP results in improvement. The message-passing in GNNs accounts for inter-structural uncertainty, thus yielding higher fidelity uncertainty estimates.

Effect of hyperparameters. Our main hyperparameters are u, α, β, γ , with u used in the Bernoulli distribution, γ in the path-generation algorithm, and (shape α , scale β ) as prior hyperparameters of the Inverse Gamma distribution. We achieve the best ECE when u = 0 . 3 , γ = 0 . 2 , α = 2 . 0 , β = 0 . 01 , however, a reasonable range always yields improvement, thus demonstrating the robustness of the method. We provide results of testing different hyperparameter values in Appendix I.

5 Conclusion

In this work, we propose to quantify the structure-wise uncertainty of a given segmentation network. Our framework explicitly models structures as samples from a probability distribution, thus helping to estimate intra-structural uncertainty. Furthermore, we incorporate inter-structural uncertainty by jointly reasoning over the structures, resulting in better fidelity uncertainty estimates. This structurewise uncertainty quantification can streamline the proofreading process by reducing the time spent finding and correcting errors. Extensive experiments show the practical applicability of our method over different segmentation backbones and datasets. We further discuss the broader impact of our work and the limitations in Appendix J and Appendix K respectively.

Acknowledgements. This research was partially supported by grants NSF 2144901, NCI R21CA258493, and NIH 1R21CA258493-01A1. The content is solely the responsibility of the authors and does not necessarily represent the official views of the National Institutes of Health.

References

Abdar, M., Pourpanah, F., Hussain, S., Rezazadegan, D., Liu, L., Ghavamzadeh, M., Fieguth, P., Cao, X., Khosravi, A., Acharya, U.R., et al.: A review of uncertainty quantification in deep learning: Techniques, applications and challenges. Information Fusion (2021)

Arganda-Carreras, I., Turaga, S.C., Berger, D.R., Cire¸ san, D., Giusti, A., Gambardella, L.M., Schmidhuber, J., Laptev, D., Dwivedi, S., Buhmann, J.M., et al.: Crowdsourcing the creation of image segmentation algorithms for connectomics. Frontiers in neuroanatomy (2015)

Banerjee, S., Magee, L., Wang, D., Li, X., Huo, B.X., Jayakumar, J., Matho, K., Lin, M.K., Ram, K., Sivaprakasam, M., et al.: Semantic segmentation of microscopic neuroanatomical data by combining topological priors with encoder-decoder deep networks. Nature machine intelligence (2020)

Baumgartner, C.F., Tezcan, K.C., Chaitanya, K., Hötker, A.M., Muehlematter, U.J., Schawkat, K., Becker, A.S., Donati, O., Konukoglu, E.: Phiseg: Capturing uncertainty in medical image segmentation. In: MICCAI (2019)

Chen, L.C., Papandreou, G., Kokkinos, I., Murphy, K., Yuille, A.L.: Semantic image segmentation with deep convolutional nets and fully connected crfs. arXiv preprint arXiv:1412.7062 (2014)

Chen, L.C., Papandreou, G., Kokkinos, I., Murphy, K., Yuille, A.L.: Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs. TPAMI (2017)

Clough, J.R., Byrne, N., Oksuz, I., Zimmer, V.A., Schnabel, J.A., King, A.P.: A topological loss function for deep-learning based image segmentation using persistent homology. TPAMI (2020)

DeGroot, M.H., Fienberg, S.E.: The comparison and evaluation of forecasters. Journal of the Royal Statistical Society: Series D (The Statistician) (1983)

Delgado-Friedrichs, O., Robins, V., Sheppard, A.: Skeletonization and partitioning of digital images using discrete morse theory. TPAMI (2014)

Dey, T.K., Wang, J., Wang, Y.: Graph reconstruction by discrete morse theory. arXiv preprint arXiv:1803.05093 (2018)

Dey, T.K., Wang, J., Wang, Y.: Road network reconstruction from satellite images with machine learning supported by topological methods. In: Proceedings of the 27th ACM SIGSPATIAL International Conference on Advances in Geographic Information Systems (2019)

Eaton-Rosen, Z., Bragman, F., Bisdas, S., Ourselin, S., Cardoso, M.J.: Towards safe deep learning: accurately quantifying biomarker uncertainty in neural network predictions. In: MICCAI (2018)

Edelsbrunner, Letscher, Zomorodian: Topological persistence and simplification. Discrete & Computational Geometry (2002)

Forman, R.: A user's guide to discrete morse theory. Séminaire Lotharingien de Combinatoire [electronic only] (2002)

Fraz, M.M., Remagnino, P., Hoppe, A., Uyyanonvara, B., Rudnicka, A.R., Owen, C.G., Barman, S.A.: Blood vessel segmentation methodologies in retinal images-a survey. Computer methods and programs in biomedicine (2012)

Gal, Y., Ghahramani, Z.: Bayesian convolutional neural networks with bernoulli approximate variational inference. arXiv preprint arXiv:1506.02158 (2015)

Gal, Y., Ghahramani, Z.: Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In: ICML (2016)

Gal, Y., Islam, R., Ghahramani, Z.: Deep bayesian active learning with image data. In: ICML (2017)

Gal, Y., et al.: Uncertainty in deep learning (2016)

Gawlikowski, J., Tassi, C.R.N., Ali, M., Lee, J., Humt, M., Feng, J., Kruspe, A., Triebel, R., Jung, P., Roscher, R., et al.: A survey of uncertainty in deep neural networks. arXiv preprint arXiv:2107.03342 (2021)

González-Hidalgo, M.: A survey on curvilinear object segmentation in multiple applications. Pattern Recognition (2016)

Gupta, S., Hu, X., Kaan, J., Jin, M., Mpoy, M., Chung, K., Singh, G., Saltz, M., Kurc, T., Saltz, J., et al.: Learning topological interactions for multi-class medical image segmentation. In: ECCV (2022)

Haehn, D., Kaynig, V., Tompkin, J., Lichtman, J.W., Pfister, H.: Guided proofreading of automatic segmentations for connectomics. In: CVPR (2018)

Hazan, T., Jaakkola, T.: On the partition function and random maximum a-posteriori perturbations. In: ICML (2012)

He, K., Gkioxari, G., Dollár, P., Girshick, R.: Mask r-cnn. In: ICCV (2017)

Hu, X.: Structure-aware image segmentation with homotopy warping. NeurIPS (2022)

Hu, X., Li, F., Samaras, D., Chen, C.: Topology-preserving deep image segmentation. In: NeurIPS (2019)

Hu, X., Samaras, D., Chen, C.: Learning probabilistic topological representations using discrete morse theory. In: ICLR (2023)

Hu, X., Wang, Y., Fuxin, L., Samaras, D., Chen, C.: Topology-aware segmentation using discrete morse theory. In: ICLR (2021)

Kendall, A., Gal, Y.: What uncertainties do we need in bayesian deep learning for computer vision? In: NeurIPS (2017)

Kipf, T.N., Welling, M.: Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907 (2016)

Kohl, S., Romera-Paredes, B., Meyer, C., De Fauw, J., Ledsam, J.R., Maier-Hein, K., Eslami, S., Jimenez Rezende, D., Ronneberger, O.: A probabilistic u-net for segmentation of ambiguous images. In: NeurIPS (2018)

Kv, R., Prasad, K., Peralam Yegneswaran, P.: Segmentation and classification approaches of clinically relevant curvilinear structures: A review. Journal of Medical Systems (2023)

Lakshminarayanan, B., Pritzel, A., Blundell, C.: Simple and scalable predictive uncertainty estimation using deep ensembles. In: NeurIPS (2017)

Lazaro-Gredilla, M., Dedieu, A., George, D.: Perturb-and-max-product: Sampling and learning in discrete energy-based models. In: NeurIPS (2021)

Li, C., Hu, X., Chen, C.: Confidence estimation using unlabeled data. In: ICLR (2023)

Long, J., Shelhamer, E., Darrell, T.: Fully convolutional networks for semantic segmentation. In: CVPR (2015)

Lovász, L.: Random walks on graphs. Combinatorics, Paul erdos is eighty (1993)

Luo, G., Wang, K., Liu, J., Li, S., Liang, X., Li, X., Gan, S., Wang, W., Dong, S., Wang, W., et al.: Efficient automatic segmentation for multi-level pulmonary arteries: The parse challenge. arXiv preprint arXiv:2304.03708 (2023)

Ma, Y., Hao, H., Xie, J., Fu, H., Zhang, J., Yang, J., Wang, Z., Liu, J., Zheng, Y., Zhao, Y.: Rose: a retinal oct-angiography vessel segmentation dataset and new model. TMI (2020)

McClure, P., Rho, N., Lee, J.A., Kaczmarzyk, J.R., Zheng, C.Y., Ghosh, S.S., Nielson, D.M., Thomas, A.G., Bandettini, P., Pereira, F.: Knowing what you know in brain segmentation using bayesian deep neural networks. Frontiers in neuroinformatics (2019)

Meil˘ a, M.: Comparing clusterings-an information based distance. Journal of multivariate analysis (2007)

Milnor, J.W., Spivak, M., Wells, R., Wells, R.: Morse theory (1963)

Mnih, V.: Machine learning for aerial image labeling. University of Toronto (Canada) (2013)

Mnih, V., Hinton, G.E.: Learning to detect roads in high-resolution aerial images. In: ECCV (2010)

Mosinska, A., Marquez-Neila, P., Kozi´ nski, M., Fua, P.: Beyond the pixel-wise loss for topologyaware delineation. In: CVPR (2018)

Mou, L., Chen, L., Cheng, J., Gu, Z., Zhao, Y., Liu, J.: Dense dilated network with probability regularized walk for vessel detection. TMI (2019)

Mou, L., Zhao, Y., Fu, H., Liu, Y., Cheng, J., Zheng, Y., Su, P., Yang, J., Chen, L., Frangi, A.F., et al.: Cs2-net: Deep learning segmentation of curvilinear structures in medical imaging. MedIA (2021)

Murphy, K.P.: Conjugate bayesian analysis of the gaussian distribution. def (2 σ 2) (2007)

Naeini, M.P., Cooper, G., Hauskrecht, M.: Obtaining well calibrated probabilities using bayesian binning. In: AAAI (2015)

Nair, T., Precup, D., Arnold, D.L., Arbel, T.: Exploring uncertainty measures in deep networks for multiple sclerosis lesion detection and segmentation. MedIA (2020)

Papandreou, G., Yuille, A.L.: Perturb-and-map random fields: Using discrete optimization to learn and sample from energy models. In: ICCV (2011)

Peng, H., Long, F., Zhao, T., Myers, E.: Proof-editing is the bottleneck of 3d neuron reconstruction: the problem and solutions. Neuroinformatics (2011)

Rand, W.M.: Objective criteria for the evaluation of clustering methods. Journal of the American Statistical association (1971)

Ren, P., Xiao, Y., Chang, X., Huang, P.Y., Li, Z., Gupta, B.B., Chen, X., Wang, X.: A survey of deep active learning. ACM computing surveys (CSUR) (2021)

Robins, V., Wood, P.J., Sheppard, A.P.: Theory and algorithms for constructing discrete morse complexes from grayscale digital images. TPAMI (2011)

Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: MICCAI (2015)

Rupprecht, C., Laina, I., DiPietro, R., Baust, M., Tombari, F., Navab, N., Hager, G.D.: Learning in an uncertain world: Representing ambiguity through multiple hypotheses. In: ICCV (2017)

Scarselli, F., Gori, M., Tsoi, A.C., Hagenbuchner, M., Monfardini, G.: The graph neural network model. IEEE transactions on neural networks (2008)

Seeböck, P., Orlando, J.I., Schlegl, T., Waldstein, S.M., Bogunovi´ c, H., Klimscha, S., Langs, G., Schmidt-Erfurth, U.: Exploiting epistemic uncertainty of anatomy segmentation for anomaly detection in retinal oct. TMI (2019)

Settles, B.: Active learning literature survey (2009)

Shit, S., Paetzold, J.C., Sekuboyina, A., Ezhov, I., Unger, A., Zhylka, A., Pluim, J.P., Bauer, U., Menze, B.H.: cldice-a novel topology-preserving loss function for tubular structure segmentation. In: CVPR (2021)

Sohn, K., Lee, H., Yan, X.: Learning structured output representation using deep conditional generative models (2015)

Sousbie, T.: The persistent cosmic web and its filamentary structure-i. theory and implementation. Monthly Notices of the Royal Astronomical Society (2011)

Staal, J., Abràmoff, M.D., Niemeijer, M., Viergever, M.A., Van Ginneken, B.: Ridge-based vessel segmentation in color images of the retina. TMI (2004)

Stucki, N., Paetzold, J.C., Shit, S., Menze, B., Bauer, U.: Topologically faithful image segmentation via induced matching of persistence barcodes. In: ICML (2023)

Student: The probable error of a mean. Biometrika (1908)

Tetteh, G., Efremov, V., Forkert, N.D., Schneider, M., Kirschke, J., Weber, B., Zimmer, C., Piraud, M., Menze, B.H.: Deepvesselnet: Vessel segmentation, centerline prediction, and bifurcation detection in 3-d angiographic volumes. Frontiers in Neuroscience (2020)

Wang, G., Li, W., Aertsen, M., Deprest, J., Ourselin, S., Vercauteren, T.: Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks. Neurocomputing (2019)

Wang, K., Qiu, Z., Wang, W., Song, T., Cao, S., Zhao, Y., Liu, J., He, Y., Gan, S., Liang, X., Xu, M., Guo, Z.: Pulmonary artery segmentation challenge 2022 (Mar 2022). https://doi.org/10.5281/zenodo.6361906, https://doi.org/10.5281/zenodo.6361906

Wang, S., Wang, Y., Li, Y.: Efficient map reconstruction and augmentation via topological methods. In: Proceedings of the 23rd SIGSPATIAL international conference on advances in geographic information systems (2015)

Yang, J., Hu, X., Chen, C., Tsai, C.: A topological-attention convlstm network and its application to em images. In: MICCAI (2021)

Zou, K.H., Warfield, S.K., Bharatha, A., Tempany, C.M., Kaus, M.R., Haker, S.J., Wells III, W.M., Jolesz, F.A., Kikinis, R.: Statistical validation of image segmentation quality based on a spatial overlap index1: scientific reports. Academic radiology (2004)

Appendix

Appendix A provides a pseudocode of our proposed Probabilistic DMT module.

Appendix B provides a detailed discussion on the hyperparameters used in this work.

Appendix C reports the computational complexity of the method.

Appendix D expands on the 'Inference procedure' outlined in Sec. 3.3 of the main paper.

Appendix E provides the details of the datasets used.

Appendix F gives implementation details of our method and the baseline methods.

Appendix G provides definitions of the evaluation metrics used.

Appendix H provides results on topological metrics.

Appendix I provides results of ablation studies mentioned in Sec. 4.2 of the main paper.

Appendix J discusses the broader impact of this work.

Appendix K discusses the limitations of this work.

A Probabilistic DMT

In Algo. 1, we provide a pseudocode of our Probabilistic DMT module proposed in Sec. 3.1. We set max _ step as 50 in our implementation. The terminologies used are: f c denotes the likelihood map from F θ centered on structure e ; ( c s , c m ) are critical points between which the path e is generated: c s denotes the saddle point and c m denotes the maxima. IG denotes the Inverse Gamma distribution, and N denotes the Gaussian distribution.

Algorithm 1 Probabilistic DMT pseudocode

̸

B Discussion of hyperparameters

The main hyperparameters in this work are u, α, β, γ . We describe the importance of each below:

u : This is the parameter for the Bernoulli distribution. In our Prob. DMT module, for every structure, we have a choice to either retain the structure as obtained from DMT, or, generate a sample skeleton using the perturb-and-walk algorithm. We model this choice using the Bernoulli distribution. Essentially, in some runs we would like the original DMT structures to also interact with the others. Thus a low value of u works best. We found u = 0 . 3 to give the best performance, that is, for every structure there is a 30% chance that it's DMT form is used and a 70% chance that a sample variant is used. We find that 0 . 15 ≤ u ≤ 0 . 3 have comparable performance.

γ : This hyperparameter is used in the weighted combination of distance Q d and likelihood f n to obtain Q ( c ′ ) , which is used to determine the next pixel location. It maintains a tradeoff between the distance regularizer Q d and the perturbed likelihood f n . The higher the value of γ , the greater the distance regularizer, and consequently the generated path will become closer to that of a straight line. This is not desirable, as a straight line would lose the original composition of the structure. Additionally, because of the perturbation in the likelihood, we do not want the path to go astray. To ensure path completeness, we require γ to be non-zero. Through experiments, we obtain the best performance when γ = 0 . 2 .

α, β : These are prior hyperparameters of the Inverse Gamma (IG) distribution. We perturb the likelihood using a Gaussian model. As the variance of the Gaussian model is unknown, we use Bayesian probability theory to sample the variance from the IG distribution (its conjugate prior). And so, α is the shape parameter and β is the scale parameter of this IG distribution. Ideally we would like a small perturbation of the likelihood and not a strong one. This is because a strong perturbation would corrupt wholly and we would not be able to sample a reasonable skeleton. At the same time, the perturbation should not be too small, otherwise we will not obtain a significant variant. The mean of the IG distribution is β α -1 (when α > 1 , β > 0 ), which on average is the value of the sampled variance for the Gaussian distribution. We achieve the best performance when α = 2 . 0 and β = 0 . 01 . The resulting sampled variance for the Gaussian model thus generates reasonable perturbation.

C Computational complexity

We report the inference time for 5 runs on a 256 × 256 input image patch as follows: Prob.-UNet: 0 . 196 sec; PHiSeg: 1 . 811 sec; Hu et al.: 5 . 485 sec; Ours: 7 . 433 sec.

The module which takes the most time is the DMT / Prob.DMT computation. Presently, this is the most optimized version as we have implemented it as an external module in C++. We will work towards porting the code to run on GPU to bring down the runtime even more.

Following [10], the computational complexity of DMT is O ( n log n ) , where n is the number of pixels in the image. Since Prob. DMT additionally computes structure variants, the complexity is O ( n log n + m ) where m is approximately the number of foreground pixels, and typically m<<n for curvilinear structure datasets. The linear term is added as we traverse each foreground pixel only once when generating the sample skeleton.

D Inference procedure

We illustrate the inference procedure in Fig. 12. There are two outcomes of our framework M ϕ , namely, the structure-wise uncertainty heatmap as well as an improved discrete segmentation map that can be used instead of the one obtained by F θ .

For each structure e , we obtain ˆ p ( e ) (the regression output) and ˆ δ 2 e (the uncertainty) from M ϕ . We take T runs of M ϕ and then for each structure e , we compute the mean across T runs as ¯ δ 2 e = 1 T ∑ T t =1 ( ˆ δ 2 e ) t , and, ¯ p ( e ) = 1 T ∑ T t =1 ˆ p ( e ) t . Next, we consider only those structures e for which ¯ p ( e ) ≥ 0 . 5 , i.e., e has a minimum probability of 50% of being a true positive. This is the threshold step in Fig. 12. We do this so as to consider only the true positive, false positive, and false negative

b)

Morse skeleton

Graph

construction

Joint

reasoning

(GNN)

threshold (0.5)

a)

Seg. map

+

e)

our discrete

segmentation map

f)

our uncertainty

heatmap

assign via shortest distance

Prob.

DMT

Segmentation

network

(run T times)

c)

d)

assign via

shortest distance

Fig. 12 (d)

Fig. 12 (e)

Fig. 12 (f)

structures in the final outcomes. We use these structures to create a skeletal discrete segmentation map (see Fig. 12(c)) which has the same spatial resolution as Fig. 12(a). As we want to recover the thickness of each structure, we overlay the two maps to get the final discrete segmentation map (see Fig. 12(e)).

The uncertainty heatmap that we obtain from M ϕ is also skeletal (see Fig. 12(d)). We recover the structure thickness to get the final uncertainty heatmap (see Fig. 12(f)). We use shortest distance to do this. Shortest distance is used to assign uncertainty values from Fig. 12(d) to the pixels in Fig. 12(e). The shortest distance uses paths only along the foreground pixels and not along the background ones. This ensures that pixels within a structure are not assigned uncertainty values from other nearby structures. We provide a zoomed-in view in Fig. 13.

We note that generating the Morse complex is computationally heavy, however, it needs to be computed only once across the T runs. As described in Sec. 3.1, the sampled structures are between ( c s , c m ) , and so the Morse complex is generated only in the first run.

E Dataset details

In this paper, we validate our results on four segmentation datasets: DRIVE [65], ROSE [40], ROADS [44], and PARSE 2022 Grand Challenge [39, 70].

DRIVE [65] The DRIVE dataset is a 2D retinal vessel dataset with 40 images. Each image has a resolution of 584 × 565 . We use the dataset's predetermined split of 20 training images and 20 test images. For training, we keep aside four randomly-chosen samples as validation, and train on the remaining 16 samples.

ROSE [40] The ROSE dataset is a 2D retinal OCTA (Optical Coherence Tomography Angiography) segmentation dataset. We use ROSE-1 (SVC) in this work. It has a predetermined split of 30 train and 9 test samples, with each sample having a resolution of 304 × 304 . For training, we keep aside four randomly-chosen samples as validation, and train on the remaining 26 samples.

ROADS [44] The ROADS dataset is a large, non-medical dataset containing 1171 aerial images ( 1108 train, 14 validation, and 49 test), each of 1500 × 1500 resolution. It is a challenging dataset due to obstruction from nearby trees, shadows, varying texture/color of roads, road class imbalance, and so on.

PARSE [39, 70] The PARSE dataset is a 3D CT dataset containing pulmonary artery segmentations. The dataset contains 100 volumes and their sizes vary from 512 × 512 × 228 to 512 × 512 × 376 . As there is no predetermined train/test split, we use 4 -fold cross-validation and report the average performance.

F Implementation details

Weuse the PyTorch framework, a single NVIDIA Tesla V100-SXM2 GPU (32G Memory), and a Dual Intel Xeon Silver 4216 CPU@2.1Ghz (16 cores) for all the experiments. The training hyperparameters for our method as well as the baselines are as tabulated in Tab. 4. Note that although PARSE is a 3D dataset, all the segmentation networks (backbones) F θ are 2D, that is, the networks are trained on 2D slices of the dataset. This was done to maintain a fair comparison across all baselines, as methods such as Prob.-UNet and PHiSeg only had 2D implementations available.

F.1 Baselines

We use the publicly available codes for the baselines:

UNet [57]: https://github.com/johschmidt42/PyTorch-2D-3D-UNet-Tutorial

Prob.-UNet [32]: https://github.com/stefanknegt/ Probabilistic-Unet-Pytorch

PHiSeg [4]: https://github.com/gigantenbein/UNet-Zoo

Hu et al. [28]: https://github.com/HuXiaoling/Structural_Uncertainty

DeepVesselNet [68]: https://github.com/dhavalshah18/deepvesselnet

CS 2 -Net [48]: https://github.com/iMED-Lab/CS-Net

F.2 Our method

Our code is available at https://github.com/Saumya-Gupta-26/struct-uncertainty . For reproducibility, we provide the architecture details as follows. The 'Joint reasoning' module described in Sec. 3.2 of our framework is a Graph Neural Network (GNN) [59], specifically a Graph Convolution Network (GCN) [31]. As per Fig. 6, the input feature vector for each graph node is constructed by passing [ x c , f c , m ] through the following architecture: C (3 , 24) → ReLU → D (0 . 2) → C (3 , 32) → ReLU → D (0 . 2) → MaxPool → Concat ( pers ) , where, C ( a, b ) denotes a convolution layer having kernel size a and number of output channels b ; D ( p ) denotes a Dropout layer with probability p ; MaxPool denotes the adaptive maxpool layer 4 returning a 1 × 1 output for each channel; and pers is a scalar value denoting the persistence of the structure. Furthermore, the bounding box size of [ x c , f c , m ] is 32 × 32 centered at each structure for 2D (for 3D, it is 32 × 32 × 32 ).

The aforementioned layers generate the input feature vector for each graph node. They are then passed through the GNN which contains the following layers: GCN (32) → ReLU → D (0 . 2) →

GCN (64) → ReLU → D (0 . 2) → GCN (32) → ReLU → D (0 . 2) , where GCN ( a ) denotes a GCNConv layer 5 having a number of output channels. The output from this sequence of layers is then fed to two separate GCN (1) layers to output the regression likelihood ˆ p ( e ) and the uncertainty ˆ δ 2 e . As per GNN fashion, the weights of the layers are shared across all the nodes.

G Evaluation metrics

We use both segmentation and uncertainty metrics to evaluate our method. We describe the metrics in detail below.

G.1 Uncertainty metrics

We use Reliability diagrams [8] and Expected calibration error (ECE) [50] to evaluate the quality of uncertainty. As both the metrics were originally designed for classification, we adapt from the classification task to semantic segmentation by treating each pixel as an independent sample. For both metrics, we first divide the probability interval [0 , 1] into N equal-sized probability intervals (each of size 1 N ). We use N = 20 bins in this work. We then calculate the accuracy and confidence of each bin.

Reliability diagrams (RD) [8] Reliability diagrams are a visual representation of model calibration by plotting the expected accuracy as a function of confidence ( confidence = 1 -uncertainty ). Perfect calibration corresponds to an identity function in the RD, i.e., the model is not over/underconfident. Consider the set of pixels/structures whose predicted probabilities fall into the bin B i . The accuracy and confidence are given by:

where, Y is the discrete segmentation ground truth (GT), and ˆ Y is the discrete segmentation map outputted by the model. In our method, ˆ Y is as shown in Fig. 12(c). Additionally, ˆ P is the pixelwise probability (likelihood) outputted by the model, whereas in our case, it is the structure-wise uncertainty ¯ δ 2 (Fig.12(d)). For our method and Hu et al., the x ∈ B i denotes structures, while in the other methods, it denotes pixels.

Expected calibration error (ECE) [50] RDs are only a visual cue, and so we also use ECE: a scalar to summarize the calibration performance. RDs do not take into account the number of pixels/structures in each bin. Thus, to account for such variations of the number of samples in a bin, we use ECE. It is given by:

where n = ∑ N i | B i | is the total number of pixels/structures. The difference between acc and conf for a given bin represents the calibration gap. When there is perfect calibration, ECE is zero.

The definition of acc and conf remains the same as defined for RDs.

G.2 Segmentation metrics

DICE [73] DICE score is a popular metric which measures the area/volumetric overlap between the predicted and ground truth discrete masks. It overcomes the class imbalance problem in the pixel-wise accuracy metric by considering only the foreground classes for measuring the overlap. The higher the DICE, the better the segmentation.

clDice [62] clDice is derived from DICE, however, clDice uses the skeleton of the predictions. This makes it sensitive to the performance of thin structures like vessels which is important in curvilinear segmentation. The higher the value, the better the segmentation.

ARI [54] The Rand index computes similarity between two clusterings. This raw score is 'adjusted for chance" to get ARI (Adjusted Rand Index) [2]. The ARI takes into account the fact that some agreement between the two clusterings can occur by chance, and it adjusts the Rand index to account for this possibility. The higher the value, the better the segmentation.

VOI [42] The VOI metric is defined as the sum of the conditional entropies between two segmentations. A lower VOI value indicates better segmentation.

Betti Number error [27] This metric computes the difference between the Betti numbers of the prediction and the ground truth. We separately compute this metric for 0-dimension and 1-dimension structures.

Betti Matching error [66] This metric is similar to Betti Number, however, in Betti Matching, the spatial location of the 0-dimension and 1-dimension components is also taken into account. Hence this metric is stricter compared to Betti Number error.

H Quantitative results on topological metrics

In Tab. 5, we provide results on two topologically important metrics, namely, Betti Number error [27] and Betti Matching error [66]. Our method consistently improves the segmentation result in terms of topology. This is consistent with our results in Tab. 1 of the main paper where our method outperforms the baselines on other topology-based metrics like clDice [62], ARI [2] and VOI [42]. Note that for the 3D PARSE dataset, we were unable to provide Betti Matching error results as its official implementation handles only 2D inputs.

I Ablation study

We check the effect of the different hyperparameters in our work by conducting experiments on the DRIVE dataset using UNet as the backbone. Our main hyperparameters are u, α, β, γ , with u used in the Bernoulli distribution, γ in the path-generation algorithm, and ( α , β ) 8 as prior hyperparameters of the Inverse Gamma distribution. We test different values and report the ECE (the lower the better) in Fig. 14. For all the experiments, we set u = 0 . 3 . We achieve the best ECE when γ = 0 . 2 , α = 2 . 0 , β = 0 . 01 , however, a reasonable range always yields improvement (notice how non-zero γ results in a sharp improvement). This demonstrates the robustness of our proposed method.

We also provide additional ablation studies in Tab. 6. We now include ablation studies on the dimensionality of the input feature vector, and size of the crops/bounding boxes, and report this in Tab. 6. We obtain the best results when the input feature vector size is 32 and the bounding box is 32 × 32 . For lower values ( 16 and 16 × 16 ), the performance reduces, while for higher values ( 64 and 64 × 64 ) we did not observe any statistically significant improvement. Thus to maintain the tradeoff between complexity and performance, we respectively use 32 and 32 × 32 for these hyperparameters.

ECE vs.(α,β)

β = 1.0e-2

4.8

β = 1.6e-2

β = 2.0e-2

(%)

4.6

E

4.4

4.2

1.5

2.0

2.5

3.0

ECE vS.Y

4.8

4.7

4.6

4.5

汇

4.4

E

4.3

4.2

0.0

0.1

0.2

0.3

0.4

0.5

Y

J Broader impact

In this work, we aim to capture structure-wise uncertainty of a given network, where a structure is defined to be a coherent set of pixels a user can intuitively understand, e.g., small vessels/branches, short stretches of road etc. Fine-scale structures such as vessels, neurons, and membranes often consist of interconnected branches or structures that form a cohesive entity. Thus structure-wise uncertainty maps can highlight uncertain instances or branches as a whole, providing a more accurate indication of regions where the segmentation may be inaccurate or uncertain. This is beneficial for proofreading or error-correction tasks as they can direct the focus of human annotators to uncertain structures that require further attention. This can save time and effort compared to pixel-wise uncertainty maps that highlight numerous pixels as uncertain, many of which do not require correction. Thus structure-wise uncertainty can provide more interpretable estimates and is a desirable approach for improving segmentation accuracy and supporting downstream analysis tasks. This can go a long way as the benefit of proofreading is twofold: it improves segmentation quality, and it also helps expand the body of labeled data that can be further used to train automatic segmentation methods. Our work is thus a useful tool in streamlining the process of scalable annotation. At the present stage, we do not foresee any potential negative societal impacts.

K Limitations

Our method currently fits in the context of curvilinear segmentation. In general, large object segmentation could also benefit from structure-wise uncertainty (structures in this case would be smaller patches/volumes). Discrete Morse theory can be used in this setting, however, we would need to make use of topological features other than the stable manifold. In its present form, our proposed solution is currently not applicable in a setting beyond curvilinear segmentation.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a framework that utilizes a probabilistic approach to extract structure-wise uncertainty estimates. This is achieved by extending the DMT to a probabilistic setting that models each structure as a sample from a probability distribution, thus capturing the intra-structural uncertainty. The proposed method then incorporates inter-structural uncertainty through a regression network which jointly reasons over the structures, using a Graph Neural Network (GNN). Additionally, a specialized inference procedure and post-processing steps are used to generate a structure-wise uncertainty heatmap, which can improve segmentation and quantify uncertainty more effectively.\n\nSTRENGTHS:\n- The authors demonstrate the versatility of their method by applying it to multiple segmentation network backbones and datasets. The proposed method was found to improve the quality of segmentation and produce high fidelity uncertainty maps for each network, making it backbone-agnostic.\n- The introduction of a probabilistic DMT and a GNN to reason about inter-structural uncertainty is a significant advancement in this field.\n- The proposed method for quantifying structure-wise uncertainty in segmentation networks has significant implications for medical image analysis. The ability to better understand and represent the uncertainty can potentially lead to improved segmentation quality, and could be crucial in medical applications where decision-making is based on these segmentation outputs.\n\nWEAKNESSES:\n- Small dataset size: DRIVE dataset 40 images. ROSE 39 images. PARSE 100 volumes. Larger dataset should be considered to evaluate the proposed method.\n- Limited Explanation of Hyperparameter: while the authors have considered a range of hyperparameters in main body and supplementary, the paper could benefit from a more detailed discussion of the influence of these hyperparameters on the method's performance and how the optimal values were determined.\n\nQUESTIONS:\nPlease add larger dataset and evaluate the proposed method. Also, explain more about the hyperparameters.\n\nLIMITATIONS:\nPlease see the 'questions' part.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

comments: SUMMARY:
This paper proposed a topology-aware uncertainty estimation method to segment curvilinear objects. The main contribution focuses on the application of discrete Morse theory (DMT). On several public datasets, the proposed method achieves SOTA performance and the visual results demonstrate the connectivity of vessels or other objects can be enhanced.

STRENGTHS:
1. Clear and well-organized paper.
2. The object connectivity is improved by the proposed sound framework.

WEAKNESSES:
1. This paper is mainly based on [24]. The technical contribution is a bit marginal here. Please clearly state the main differences and the take-home insights.
2. Some non-deep methods are also good at achieving better topologies. See:
[1] Liu, Siqi, et al. "Rivulet: 3D neuron morphology tracing with iterative back-tracking." Neuroinformatics 14 (2016): 387-401.
Please include them for discussion and comparison.
3. The title is a bit over-claimed. "Topology-Aware Uncertainty for Curvilinear Object Segmentation" might be more suitable.
4. Using graph models to capture the topology information is not new. See:
[2] Shin, Seung Yeon, et al. "Deep vessel segmentation by learning graphical connectivity." Medical image analysis 58 (2019): 101556.
More discussions and comparison should be added.

QUESTIONS:
See the above weaknesses.

LIMITATIONS:
See the above weaknesses.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel method for the estimation of uncertainty of the structures in the segmentation results from existing methods/models, in order to facilitate the subsequent proofreading process. To this end, it models the intra-structure and inter-structure uncertainties in two modules. While the former considers the geometry, contrast and model’s confidence, the latter considers neighbours and thus context. The former is implemented using the Probabilistic discrete Morse theory (DMT), which samples the Morse skeletons using the inverse Gamma distribution. The latter uses graph neural network (GNN) for the prediction of the uncertainty of all the structures jointly through regression using the attenuation loss. The proposed method has been validated over three publicly accessible datasets and compared with several state-of-the-art. Relatively better results have been obtained in different metrics. A number of ablation studies have also been carried out on variants of the DMT and GNN and some hyperparameters. Some insights have been accumulated into how the proposed method behaves under different configurations.\n\nSTRENGTHS:\n1.\tThe topic of the paper is interesting and important for the subsequent validation of the results for image segmentation and find many applications in the read world such as medial image segmentation and analysis, object classification and recognition, industrial quality assurance, etc.\n2.\tA novel method has been proposed to post-process the segmentation results produced by some existing models for their proofreading and validation. It includes two main modules: intra-structure uncertainty and inter-structure uncertainty estimation. While the former is modelled using the probabilistic DMT, the latter is modelling using the graph neural network. The method is well motivated and supported by solid theory.    \n3.\tThe proposed method has been validated over three publicly accessible datasets and compared with several state-of-the-art. Relatively better results have been obtained in different metrics. The t-test has also been performed whether the improvement is significant.\n4.\tA number of ablation studies have been carried out on the variants of DMT and GNN and some hyperparameters in the process of sampling of Morse skeletons. Some insights have been accumulated into how the proposed method behaves under different conditions, that would instruct how it can be applied in the real world.\n5.\tThe proposed uncertainty estimation has been applied to re-calibrate the segmentation results obtained. The experimental results have shown that it does help identify the false positives and false negatives.\n\nWEAKNESSES:\n1.\tWhile the paper targets the proofreading of the segmentation results by experts, its necessity could be emphasised: the automatic algorithms cannot guarantee  the correctness of the segmentation, especially for medical imaging where the anatomy and structures may vary from one subject to another, and their prior knowledge is not always available  This is also the process for the relevant researchers/experts to learn and accumulate insights about the variation of the structures of the subjects for individualised diagnosis, treatment and medicine.\n2.\tSome details, elaborations, clarifications and discussions are missing from place to place. For example, while the construction of the graph is described later, its main idea could be summarised in Introduction. “structure” is widely used throughout the paper but its count has ever been discussed. While the ground truth is again used to train the GNN, some discussions could be made about its special requirement: GT plays a crucial role in the validation and guiding the estimation of the uncertainty for false positives and false negatives and thus, may have special requirement about quality and reliability. It is not clear how to guarantee this, especially when the datasets are large and include subtle structures. On the other hand, is that possible to use the salient structures to guide the search for faint structures in the spirit of Canny edge detection?  “Persistence value” in L239-240: how to calculate it? where is it from? No details or references are given.    \n3.\tThe key steps for the re-calculation of the Morse skeletons in probabilistic DMT may require further investigation. The next point selection criteria should be stated explicitly: maximize the sum of the inverse distance and the likelihood from the candidates and discuss why this criterion is feasible, especially when the distance is not really comparable directly with the likelihood. For the definition of Q(c’), this definition may require further investigation: what is the rationale for this definition? while the first term is inverse distance, is that comparable to the likelihood in the second term? Is this optimal? Is there any other alternative? Is that possible to draw conclusion that the first term will always lie in the unit interval [0, 1]? Overall, this is still a heuristic, which may not hold in some cases. “This process is done separately and in parallel for every structure.” In L222-223: how many Morse skeletons were sampled for each structure and can they be directly used to calculate as uncertainty their means and variances? The final results may be affected by the number of runs and their combinations made.\n4.\tThe computational complexity and time have not been analysed and reported. Thus, it is not clear how efficient the proposed method is and how much time it requires to process a set of given images. \n5.\tFurther analysis of the experimental results would help. For example, while the proposed method is effective in improving the results of the existing methods in ECE, clDice, ARI and VOI, but it is not always in Dice. It is not clear why. Any further insights and explanations would really help.  \n6.\tMore ablation studies could be carried out on other parameters and components such as the dimensionality of input feature vector,  crops/bounding boxes on the structure, alternatives to the shortest distance for structure inference. \n7.\tThe claim that the proposed method may be applicable to non-medical applications: civil engineering, road network and railway track segmentation, has not been validated. \n8.\tIt is not clear how the prior knowledge can be used to guide the image segmentation, uncertainty estimation, and result re-calibration, rather than just based on some heuristics in skeleton-recalculation and structure inference, while such heuristics may not hold in some case. Also, it is not clear how such heuristics contribute to the final errors. Some further investigation would be encouraged.\n\nQUESTIONS:\nPlease see the comments/concerns in weaknesses,\n\nMore detailed comments are as follows:\n1.\tL41: accept or reject/correct structural proposals efficiently: relative to what: prior knowledge or experience?\n2.\tL132, each structure: some elaborations would help: what that is: line segment, blob, junction, keypoint, or anything else.\n3.\tL140-14, At each training iteration, it takes one sample skeleton for each structure: Does this mean that the number of iterations will be determined by the number of the total structures and the number of sampled skeletons for each structure? How to count the number of the structures? any details?\n4.\tL170, a perturb-and-walk algorithm: refs?\n5.\tL192-194, At every step, we always walk to the neighboring pixel with the  highest likelihood value: This is a heuristic only, which may not hold in some cases.  \n6.\tL224, The output of Prob. DMT is effectively one sample skeleton: is the same number of sample skeletons as the number of runs? Can all the sample skeletons be directly used to calculate the uncertainties? How to count the number of structures?\n7.\tL227-228, network that takes as input each structure: This is confusing: takes as input each structure each time or all the structures together? If the former is the case, do you have to run the model many times? what is the motivation for this? Can they be combined together and run the model once to directly estimate the uncertainties?\n8.\tL237, are smaller crops/bounding boxes: what are the sizes of the boxes? Any details?\n9.\tSe 3.3: shortest distance: This is heuristic only, which may not hold in some cases.\n10.\tL299, others -> the others\n\nLIMITATIONS:\nThe limitations of the proposed work are discussed in the supplementary materials.\n\nThe potential negative social impact are not highly related and thus are not discussed.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work aims to contribute to proofreading by proposing uncertain structures in a topological sense. The work proposes a method to quantify a form of structure-wise uncertainty from segmentations, where the framework explicitly models structures as samples from a probability distribution. First, the structures are extracted via discrete Morse theory (DMT). Next, the uncertainty is modeled via a joint prediction model that estimates the uncertainty of a structure in consideration of the surrounding structures. Furthermore, the authors propose a novel probabilistic DMT concept to model intra-structure uncertainty. The method is then successfully experimentally validated.\n\nSTRENGTHS:\n-  The work presents a real generalization of the work on DMT for segmentation [25]. E.g., if I understand it correctly, if one chooses perturbation 0, the result of the method will be exactly the DMT result. This is a nice property and a strong extension of prior work. \n\n- The motivation is clear and interesting. \n\n- The method is well described and formalized.\n\nWEAKNESSES:\n1) **Experimentation:** The authors state to propose a \"topology-aware\" method. However, the authors do not evaluate their results on popular topology-related metrics. E.g., Betti number errors in dimensions 0,1, and 2 (for the 3D dataset). Evaluation of the performance with respect to these metrics will improve the Experimentation. Comparison to the recently published Betti matching error (1), which considers the spatial agreement of the topological structures, would further increase the interpretability of results. \n\n\n2) **Method** If I understand the method correctly, the DMT calculation at the end could be seen as a post-processing step or additional network to improve connectivity on the results; it has been shown for curvilinear structure segmentation that such an additional step improves segmentation performance. Clearly, such a concept makes this a multi-step procedure which has limitations compared to the cited method by Hu et al. \n\n3) The definition of \"topological structures\" is not very clear, and to me, it appears to not align with some definitions in algebraic topology, especially in dimension-1. Intuitively I would expect a \"topological structure\" to be (e.g., in dimension-1) a closed loop. This appears not to be the case here. If I misunderstand this, could the authors clarify how they represent the cycles, and how this is different from their representation of features in dimension-0 and provide more explanations? \n\n\n\n**References**: \n\n[] are references from manuscript \n\n(1) Stucki, N. , et al. \"Topologically Faithful Image Segmentation via Induced Matching of Persistence Barcodes.\" ICML (2023).\n\nQUESTIONS:\nThis is a minor question regarding the utility of the method. The authors motivate by and mention that their contribution is a way to simplify and improve proofreading. Do you have practical experiments in a proofreading setting that you can share? Is there an experiment with human experts, e.g., readers in ophthalmology?\n\nLIMITATIONS:\nThe authors do not provide a dedicated limitations section. I would like to learn more about the limitations of their method in the context of stricter definitions in algebraic topology.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}