Cascaded multitask U-Net using topological loss for vessel segmentation and centerline extraction

Pierre Rougé,Nicolas Passat,Odyssée Merveille

2024-02-22

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.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The problem that this paper attempts to solve is the accuracy of blood vessel segmentation and centerline extraction, especially when dealing with complex vascular networks. Specifically:

1. **Challenges of complex geometries and topological structures**: Traditional deep - learning methods have difficulty in capturing the complex geometric shapes and specific topological structures of vascular networks, which are crucial for many applications (such as blood flow modeling).

2. **Problems in calculating differentiable skeletons**: In order to use the clDice loss function, it is necessary to calculate the skeleton of the predicted segmentation in a differentiable way. Current differentiable skeletonization algorithms are either inaccurate or too computationally expensive.

To solve these problems, the authors propose a cascaded multi - task U - Net architecture, which improves blood vessel segmentation and centerline extraction in the following ways:

- **Introducing topological loss**: Using the clDice loss function, which focuses on the blood vessel centerline, so as to better maintain the topological correctness of the segmentation results.

- **Efficient skeleton calculation**: A U - Net - based skeletonization method is proposed, which achieves a good balance between accuracy and computation time and is naturally differentiable.

- **Multi - task learning**: Through the cascaded U - Net architecture, segmentation is carried out first, and then the segmentation result and the original image are used for the skeletonization task. This architecture not only improves the accuracy of segmentation but also reduces error propagation.

### Specific problem summary

- **Problem 1**: How to improve the accuracy of blood vessel segmentation and centerline extraction, especially when dealing with complex vascular networks?
- **Solution**: Introduce the clDice loss function, focus on the blood vessel centerline, and ensure the topological correctness of the segmentation results.

- **Problem 2**: How to efficiently calculate a differentiable skeleton for training a neural network?
- **Solution**: Use a U - Net - based skeletonization method that is both accurate and computationally efficient.

- **Problem 3**: How to reduce error propagation and improve the overall performance of the model?
- **Solution**: Through a multi - task learning framework, jointly train the segmentation and skeletonization tasks, so that the segmentation task benefits from the inductive bias of the skeletonization task.

Through these methods, the cascaded multi - task U - Net proposed by the authors can provide more accurate topological structure segmentation results while maintaining a relatively low computation time.