Topology-aware loss for segmentation in computed tomography images

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 trained with a standard loss function. Incorporating these invariants into network training can help regularize training, and thus improve segmentation performance, particularly when annotated data are limited. To this end, this paper presents a topology aware loss function that introduces a new term penalizing topology dissimilarities between the ground truth and prediction through persistent homology. Unlike previously suggested segmentation network designs that apply the threshold filtration on the likelihood function of the prediction map and the Betti numbers of the ground truth, the proposed model employs the Vietoris-Rips filtration to obtain persistence diagrams of both ground truth and prediction maps and measures their dissimilarity with the Wasserstein distance between the corresponding diagrams. The proposed loss term is integrated into an encoder-decoder network with UNet backbone, implemented in PyTorch, with topological computations performed using Ripser++ and Gudhi libraries. Experiments were conducted on 4327 computed tomography (CT) images of 24 subjects for aorta and great vessel segmentation. The results show that the proposed topology-aware loss function improved the performance of its counterparts. Furthermore, compared to the base loss functions that do not include the proposed topological term, our model achieved statistically significant 1%-3% improvement in vessel-level f-scores.