Topology-Aware Loss for Aorta and Great Vessel Segmentation in Computed Tomography Images

TL;DR

This work tackles the limitation of pixel-wise segmentation losses by enforcing global geometric invariants in CT vessel segmentation. It introduces a topology-aware loss built on persistent homology, using the Vietoris-Rips filtration to compute persistence diagrams for both ground-truth and predicted vessel maps and measuring their dissimilarity with the Wasserstein distance. The principal formulation, TopLoss = ∑_I ω_I L_I with $ω_I = 1 + α_I d_0(Π_{S_I}, Π_{ ilde S_I}) + β_I d_1(Π_{S_I}, Π_{ ilde S_I})$, guides training of a UNet where the topological signal complements the base loss; a 25-epoch warm-up precedes topology-driven optimization. Evaluated on 4327 CT slices from 24 subjects, the topology-aware loss improves pixel- and vessel-level metrics over baselines including a likelihood-filtration method and Fourier descriptors, demonstrating concurrent modeling of the aortic shape and the geometry between great vessels. This topology-preserving approach offers robust segmentation under limited data and has potential to enhance quantitative vascular analyses in clinical workflows.

Abstract

Segmentation networks are not explicitly imposed to learn global invariants of an image, such as the shape of an object and the geometry between multiple objects, when they are trained with a standard loss function. On the other hand, incorporating such invariants into network training may help improve performance for various segmentation tasks when they are the intrinsic characteristics of the objects to be segmented. One example is segmentation of aorta and great vessels in computed tomography (CT) images where vessels are found in a particular geometry in the body due to the human anatomy and they mostly seem as round objects on a 2D CT image. This paper addresses this issue by introducing a new topology-aware loss function that penalizes topology dissimilarities between the ground truth and prediction through persistent homology. Different from the previously suggested segmentation network designs, which apply the threshold filtration on a likelihood function of the prediction map and the Betti numbers of the ground truth, this paper proposes to apply the Vietoris-Rips filtration to obtain persistence diagrams of both ground truth and prediction maps and calculate the dissimilarity with the Wasserstein distance between the corresponding persistence diagrams. The use of this filtration has advantage of modeling shape and geometry at the same time, which may not happen when the threshold filtration is applied. Our experiments on 4327 CT images of 24 subjects reveal that the proposed topology-aware loss function leads to better results than its counterparts, indicating the effectiveness of this use.

Figures (7)

• Figure 1: (a) Anatomic formation of the aortic arch and the large arteries. (b) Manual annotations of the aortic arch and the large arteries on three exemplary axial slices. (c) CT scans for these annotations. Note that the annotations only for the green, orange, and blue squares are illustrated for better visualization. All pixels outside the rectangles are annotated as background.
• Figure 2: Illustration of persistent homology. In this pedagogical example, we start with a point cloud consisting of the 19 black dots above. As time progresses, we gradually include more points in the filtration. At time $t,$ the associated level of the filtration consists of the blue colored region in diagram $t$ above. The rank of the 0-dimensional (resp. 1-dimensional) homology, which computes the number of connected components (resp. holes) in the filtration at time $t,$ is the number of intersections of the line $x=t$ with the set of blue lines (resp. green lines) in the last diagram. The homology groups at each level of the filtration can be computed via the corresponding Vietoris-Rips simplicial complexes as indicated above. The blue and green lines, which are called the barcodes for the persistent homology, give a summary of the rough shape of the configuration of points. The above times $t_1, \cdots, t_6$ are chosen to represent the precise levels where the ranks of the homology groups change, and hence are precisely the levels where the blue and green lines are born or die.
• Figure 3: (a), (c) Boundaries of two homotopy equivalent objects, and (b), (d) 1-dimensional persistent homologies of the objects shown in (a) and (c), respectively. Even though the standard homology groups of (a) and (c) are the same, since they are homotopy equivalent, their persistent homologies are not. The latter takes into account the shape of the object in this case. Namely, the smaller circular bump in the upper right part of (c) is responsible for the smaller bar in (d).
• Figure 4: (a), (c) Boundaries of two pairs of homotopy equivalent objects, and (b), (d) 0-dimensional persistent homologies of the objects shown in (a) and (c), represented by two bars. The long bar in each case represents the connected component which persists forever. The smaller bar in each case corresponds to the distance between the two connected components. This bar is smaller in (b) than in (d), consistent with the fact that the components are closer to each other in (a).
• Figure 5: (a) Boundaries of the objects found in the ground truth map ${\cal S}_I$ and (b) its persistence diagram $\Pi_{{\cal S}_I}$ for the homology group 1. (c) Boundaries of the objects found in the prediction map ${\widehat{\cal S}}_I$ and (d) its persistence diagram $\Pi_{{\widehat{\cal S}}_I}$ for the homology group 1. (e) Illustration of the best matching of these two diagrams. The Wasserstein distance is the sum of the powers of the distances between the matched points.