Continuous and complete liver vessel segmentation with graph-attention guided diffusion

TL;DR

The paper addresses the challenge of accurate liver vessel segmentation from CT, where continuity and small-vessel visibility are difficult and annotations vary. It presents a diffusion-based segmentation framework conditioned on both image slices and a multiscale vascular graph, using vanilla, dynamic, and graph-attention conditioning with a joint loss $L_{total}=L_{den}+L_{graph}$. The reverse diffusion is guided by embeddings $\mathbf{f}^t=\mathbf{f}_{\mathbf{c}}^t+\mathbf{f}_{\mathbf{v}}^t$, with $p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)=\mathcal{N}(\mathbf{x}_{t-1};\mu_\theta(\mathbf{x}_t|\mathbf{f}^t,t),\Sigma_\theta(\mathbf{x}_t|\mathbf{f}^t,t))$, and the graph component uses GATv2 and LIIF to capture multiscale geometry. Empirical results on 3D-ircadb-01 and LiVS show superior Dice similarity coefficient and sensitivity, along with improved vessel connectivity (clDice and Con), demonstrating enhanced continuity and completeness over state-of-the-art baselines. The work has practical impact for preoperative planning by providing more complete and connected liver vessel trees, even under annotation variability, and releases code for reproducibility.

Abstract

Improving connectivity and completeness are the most challenging aspects of liver vessel segmentation, especially for small vessels. These challenges require both learning the continuous vessel geometry, and focusing on small vessel detection. However, current methods do not explicitly address these two aspects and cannot generalize well when constrained by inconsistent annotations. Here, we take advantage of the generalization of the diffusion model and explicitly integrate connectivity and completeness in our diffusion-based segmentation model. Specifically, we use a graph-attention module that adds knowledge about vessel geometry, and thus adds continuity. Additionally, we perform the graph-attention at multiple-scales, thus focusing on small liver vessels. Our method outperforms eight state-of-the-art medical segmentation methods on two public datasets: 3D-ircadb-01 and LiVS. Our code is available at https://github.com/ZhangXiaotong015/GATSegDiff.

Figures (11)

• Figure 1: (a) Vessel trees of nnUNetb27, Swin UNETRb28, our proposed model and the ground truth. (b) Vessel trees predicted by our method across different steps ($t{=}0$ is the final diffusion iteration). Rescaling is the process of recovering the resolution of resized CT images to that of the original CT images. Our method predicts complete and continuous vessel trees.
• Figure 2: Overview of our model: (A) A vanilla diffusion model starting from noisy images $\mathbf{x}_T$, and predicting segmentation masks $\mathbf{x}_0$ (in orange); (B) A dynamic conditioning model, conditioned on three CT slices $\mathbf{c}$ (in gray); and (C) A multiscale graph-attention conditioning model, starting from a graph structure $\mathcal{G}$ (in blue).
• Figure 3: Our network architecture: (A) vanilla conditioning diffusion (orange); (B) dynamic conditioning diffusion (gray); (C) multiscale graph-attention conditioning (blue). These components interact through the vertical dashed/ solid arrows. The dashed upwards arrows adapt the conditioning over time. The solid downwards arrows add the conditioning features: CT slices $\mathbf{c}$, CT slice embeddings $\mathbf{f}_\mathbf{c}^t$, and graph embeddings $\mathbf{f}_\mathbf{v}^t$.
• Figure 4: Visualizations on the 3D-ircadb-01b40 dataset. (a) The liver vessel tree segmented by our proposed model; (b), (c), (d), (e), (f), (g), (h) and (i) are the liver vessel tree segmented by the baselines: HiDiffb63, MERITb62, TransUNetb68, MedSegDiffb20, EnsemDiffb19, Swin UNETRb28, nnUNetb27 and b41 respectively; (j) The ground truth liver vessel tree. The yellow and white boxes compare the completeness and continuity between our proposed model and the baselines. We also show the boxes enlarged in the top/ bottom right corners. Our method (1$st$ column) achieves the most similar predictions to the ground truth (last column) in both fine vessel segmentation, and connectivity.
• Figure 5: Results on the LiVS b41 dataset.Left: The performance of our method compared with HiDiffb63, MERITb62, TransUNetb68, MedSegDiffb20, EnsemDiffb19, Swin UNETRb28, nnUNetb27 and Gao et al.b41. clDice and Con are not applicable for the LiVS dataset with discontinuous annotations. Our method outperforms others in terms of Sen scores. Spe of our model is slightly lower than EnsemDiff, which is a trade-off between completeness (Sen) and accuracy (Spe) of vessel segmentation. Right: Our DSC score is comparable with nnUNet is because that the high-contrast tumors in the liver will cause more outliers (shown in the box plot) for our method, and negatively affect the averaged DSC score over all data. Deterministic methods are more efficient than generative methods in inference.