Skeleton Recall Loss for Connectivity Conserving and Resource Efficient Segmentation of Thin Tubular Structures

TL;DR

This work tackles the challenge of connectivity-preserving segmentation for thin tubular structures, where standard overlap losses fail to preserve topology. It introduces Skeleton Recall Loss, which avoids differentiable skeletons by computing a tubed ground-truth skeleton on the CPU and applying a soft recall against network predictions via $\mathcal{L}_{SkelRecall} = - \frac{1}{|C|}\sum_{c\in C}\frac{\sum_i Y_{\mathrm{skel},i,c} \cdot \hat{Y}_{i,c}}{\sum_i Y_{\mathrm{skel},i,c}}$, enabling efficient, multi-class segmentation. The method demonstrates state-of-the-art performance on five public datasets spanning 2D and 3D, binary and multi-class tasks, while reducing overheads by more than 90% relative to differentiable-skeleton losses and maintaining compatibility with various architectures. This approach offers practical implications for topology-aware segmentation in resource-constrained settings and broad applicability across domains, with code freely available for implementation and extension.

Abstract

Accurately segmenting thin tubular structures, such as vessels, nerves, roads or concrete cracks, is a crucial task in computer vision. Standard deep learning-based segmentation loss functions, such as Dice or Cross-Entropy, focus on volumetric overlap, often at the expense of preserving structural connectivity or topology. This can lead to segmentation errors that adversely affect downstream tasks, including flow calculation, navigation, and structural inspection. Although current topology-focused losses mark an improvement, they introduce significant computational and memory overheads. This is particularly relevant for 3D data, rendering these losses infeasible for larger volumes as well as increasingly important multi-class segmentation problems. To mitigate this, we propose a novel Skeleton Recall Loss, which effectively addresses these challenges by circumventing intensive GPU-based calculations with inexpensive CPU operations. It demonstrates overall superior performance to current state-of-the-art approaches on five public datasets for topology-preserving segmentation, while substantially reducing computational overheads by more than 90%. In doing so, we introduce the first multi-class capable loss function for thin structure segmentation, excelling in both efficiency and efficacy for topology-preservation.

Figures (8)

• Figure 1: Diversity of thin structures. Segmentation of thin structures is a challenging task in engineering and medical imaging. This is highlighted in 5 diverse datasets used in this work to incorporate the segmentation of: a) Roads in satellite imagery, b) Retinal vessels, c) Cracks in concrete structures, d) Inferior alveolar canal in facial CTs, and e) Circle of Willis arterial vessel components.
• Figure 2: Comparison of state-of-the-art loss functions on the task of thin structure segmentation.Top: Our Skeleton Recall Loss efficiently addresses connectivity conservation, unlike standard dice loss, without the overhead of clDice Loss, making it ideal for multi-class problems as well. Bottom: Qualitative results on the TopCoWtopcow dataset. Due to computational cost clDice Loss can not be used for multi-class segmentation.
• Figure 3: The challenges of Differentiable Skeletons. Visual comparison of (c) the soft skeleton used for the calculation of the clDice Loss clDice and (d) the proposed tubed skeleton used for Skeleton Recall Loss for (a) an image and the corresponding (b) ground truth segmentation, originating from the TopCoW dataset topcow.
• Figure 4: Overview of our method in comparison to differentiable skeleton based approaches. Initially, a segmentation network (green) predicts a segmentation mask. Our proposed Skeleton Recall Loss (blue) calculates the soft recall of the prediction on the precomputed Tubed Skeleton of the ground truth. In doing so, we mitigate the massive overheads introduced by differentiable skeleton based methods (red).
• Figure 5: Connectivity conservation in qualitative results on 5 datasets. nnUNet with conventional segmentation losses performs well in adequately delineating general structures, particularly thicker ones. However, challenges arise in accurately capturing thin structures and maintaining connectivity within the segmentation. This is demonstrated on examples from (top to bottom) Roads, DRIVE, Cracks, Toothfairy and TopCoW datasets. Augmenting the model with clDice Loss yields some improvement but falls short in addressing connectivity issues. In contrast, our proposed Skeleton Recall Loss demonstrates enhanced preservation of topology and improved connectivity in segmentation outputs.