SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Saumya Gupta $^{1*}$  Yikai Zhang $^{2}$  Xiaoling Hu $^{1,3}$  Prateek Prasanna $^{1}$  Chao Chen $^{1}$
$^{1}$ Stony Brook University, NY, USA  $^{2}$ Morgan Stanley, NY, USA
3 Athinoula A. Martinos Center for Biomedical Imaging,
Massachusetts General Hospital and Harvard Medical School, MA, USA
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
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 fine-scale 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,
*Email: saumgupta@cs.stonybrook.edu
Figure 1: Motivating examples for structure-wise uncertainty. In the segmentation result (c), orange highlights a false positive structure, and pink highlights a false negative. Methods (d)-(f) are uncertainty estimates of the prediction in (c). PHiSeg [4] assigns pixels along boundaries as uncertain. Hu et al. [28] captures uncertainty at a structural level, but produces overconfident maps (assigns zero uncertainty to many structures). Ours produces better structure-wise uncertainty estimates: both the highlighted false positive/negative structures have high uncertainty.
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 error-prone 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
ii) Morse skeleton
a) Intra-structural uncertainty
iii) DMT vs. GT
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.
iv) Prob. DMT
vi) Likelihood map
b) Inter-structural uncertainty
vii) w/o inter
viii) w/ inter
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.
1. We propose a novel method to estimate the uncertainty of a given segmentation network at a structural level.
2. We propose Probabilistic DMT, a probabilistic method to generate structural variations and to capture intra-structural uncertainty.
3. 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.
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
Figure 3: An overview of the proposed method  $M_{\phi}$ . The given segmentation network  $F_{\theta}$  has frozen weights. Probabilistic DMT decomposes the likelihood into structures, and samples skeleton representations of each. A graph is then constructed over the structures to perform joint reasoning of their uncertainty. The training is supervised by comparing with the GT (via the loss  $L_{UQ}$ , red arrow).
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. Structure-wise 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.
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_{\theta}$  denote the trained segmentation network, and  $M_{\phi}$  denote our proposed external uncertainty quantification framework.  $M_{\phi}$  takes as input the likelihood map of  $F_{\theta}$  and the input image. It generates a set of structures, and estimates an uncertainty value for each of them. During training,  $M_{\phi}$  is trained by comparing with the ground truth (GT) annotation.
$M_{\phi}$  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_{\theta}(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_{\theta}$ . 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  $(\nabla 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 one-dimensional 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.
Figure 4: Orange indicates FN structures; (c) shows saddle points (red) and maximas (blue), and omits minimas; (d) shows the union of the stable manifolds of the saddle points.
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
the likelihood map, the sample skeletons will be less variant, resulting in a low uncertainty.
Likelihood  $f$  DMT: e Prob. DMT:  $\hat{e} (\# 1)$  Prob. DMT:  $\hat{e}$  (#2)
Figure 5: Structures (#1,#2) sampled from the distribution. Green arrow is path chosen using  $Q(c')$ ; red arrow is next step w/o considering  $Q_{d}(c')$ .
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  $\hat{e}$ , we first draw a likelihood  $f_n$  from a distribution centered on  $f$
as  $f_{n} \sim 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  $\hat{e}$  from  $f_{n}$ . Our walk algorithm continuously grows  $\hat{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'' = \operatorname{argmax}(Q(c'))$ , where,  $c' \in \operatorname{neighborhood}(c)^2$  and  $Q(c') = \gamma Q_d(c') + (1 - \gamma)f_n(c')$ , and,  $Q_d(c') = \frac{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  $\hat{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  $\gamma = 0$ . In practice, with some probability, we consider the original structure  $e$  from DMT over generating its variant  $\hat{e}$ . Specifically, following a Bernoulli distribution, with a small probability  $u$  we retain  $e$ , while with probability  $1 - u$  we sample its variant  $\hat{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 \in E$ , and outputs whether it is a true positive and the uncertainty of  $F_{\theta}$  in predicting it.