[Paper Review] Cascaded multitask U-Net using topological loss for vessel segmentation and centerline extraction

Pierre Rougé, Nicolas Passat+1 more|arXiv (Cornell University)|Jul 21, 2023

Cerebrovascular and Carotid Artery Diseases6 citations

TL;DR

This paper proposes a cascaded multitask U-Net that replaces the differentiable soft-skeleton algorithm with a learnable U-Net-based skeletonizer to compute the clDice loss, enabling topological constraint learning for 3D vessel segmentation. The method achieves significantly improved topological accuracy, particularly in preserving vessel connectivity and centerline topology, outperforming baselines on 3D TOF MRA data.

ABSTRACT

Vessel segmentation and centerline extraction are two crucial preliminary tasks for many computer-aided diagnosis tools dealing with vascular diseases. Recently, deep-learning based methods have been widely applied to these tasks. However, classic deep-learning approaches struggle to capture the complex geometry and specific topology of vascular networks, which is of the utmost importance in most applications. To overcome these limitations, the clDice loss, a topological loss that focuses on the vessel centerlines, has been recently proposed. This loss requires computing, with a proposed soft-skeleton algorithm, the skeletons of both the ground truth and the predicted segmentation. However, the soft-skeleton algorithm provides suboptimal results on 3D images, which makes the clDice hardly suitable on 3D images. In this paper, we propose to replace the soft-skeleton algorithm by a U-Net which computes the vascular skeleton directly from the segmentation. We show that our method provides more accurate skeletons than the soft-skeleton algorithm. We then build upon this network a cascaded U-Net trained with the clDice loss to embed topological constraints during the segmentation. The resulting model is able to predict both the vessel segmentation and centerlines with a more accurate topology.

Motivation & Objective

• Address the challenge of topologically accurate 3D vessel segmentation in medical imaging, especially for low-resolution, noisy human vascular data.
• Overcome the limitations of the soft-skeleton algorithm in 3D, which produces suboptimal skeletons and hinders effective clDice loss computation.
• Improve vessel topology preservation in segmentation by learning a differentiable skeletonization network as part of a multitask U-Net architecture.
• Enable end-to-end training of a segmentation model that jointly predicts vessel volumes and centerlines with enhanced topological fidelity via clDice loss.

Proposed method

• Replace the differentiable soft-skeleton algorithm with a dedicated U-Net to learn skeletonization from predicted segmentations, improving skeleton accuracy in 3D.
• Train the skeletonization U-Net independently on ground truth segmentations before integrating it into a cascaded multitask U-Net architecture.
• Construct a cascaded U-Net where the segmentation head uses the predicted skeleton from the first network as input to compute the clDice loss.
• Optimize the joint training with a weighted combination of clDice and Dice losses, using hyperparameters λ₁ = 0.25 and λ₂ = 0.25, with skeletonization network weights frozen.
• Use the clDice metric, which focuses on centerline similarity, as the primary loss to enforce topological correctness during training.
• Apply the method to 3D time-of-flight MRA data, evaluating performance on DSC, clDice, precision, sensitivity, and topological error metrics.

Figure 1: The baseline U-Net architecture used in the proposed approach (see Section 2.2 ). The output of this network is either the vascular segmentation (as shown here) or the vascular skeleton, depending on the chosen task.

Experimental results

Research questions

• RQ1 Can a learnable U-Net-based skeletonizer produce more accurate skeletons than the soft-skeleton algorithm in 3D medical images?
• RQ2 Does replacing the soft-skeleton with a learned skeletonizer improve the topological quality of vessel segmentation in 3D TOF MRA?
• RQ3 Can a cascaded multitask U-Net trained with clDice loss using a learned skeletonizer achieve better segmentation and centerline accuracy than baselines?
• RQ4 How does the proposed method compare to standard U-Net and U-Net with soft-skeleton-based clDice in terms of topological error and connectivity preservation?

Key findings

• The learned U-Net skeletonizer produced significantly more accurate skeletons than the soft-skeleton algorithm, with p-values < 0.001 for all topological metrics (χ error, β₀, β₁, and skeletonization error).
• The cascaded multitask U-Net achieved a 24.3% relative reduction in χ error (0.379 ± 0.186) compared to the U-Net (Dice) baseline (0.495 ± 0.187), with p < 0.001.
• The method showed a significant improvement in β₁ error (0.170 ± 0.104) compared to the U-Net (Dice+clDice) baseline (0.212 ± 0.137), with p = 0.02.
• Despite similar DSC and clDice scores, the proposed method demonstrated better sensitivity (0.746 ± 0.040) and reduced false negatives, indicating improved detection of connected vessel structures.
• Qualitative results in Figure 3 show that the proposed method better reconnects distal vessel branches, preserving topology more effectively than baselines.
• The method outperformed the U-Net (Dice+clDice) baseline in topological metrics, even though it had slightly lower precision, indicating a trade-off favoring connectivity over strict boundary accuracy.

Figure 2: Architecture of the proposed cascaded U-Net (see Section 2.2 ).