CAPE: Connectivity-Aware Path Enforcement Loss for Curvilinear Structure Delineation

Abstract

Promoting the connectivity of curvilinear structures, such as neuronal processes in biomedical scans and blood vessels in CT images, remains a key challenge in semantic segmentation. Traditional pixel-wise loss functions, including cross-entropy and Dice losses, often fail to capture high-level topological connectivity, resulting in topological mistakes in graphs obtained from prediction maps. In this paper, we propose CAPE (Connectivity-Aware Path Enforcement), a novel loss function designed to enforce connectivity in graphs obtained from segmentation maps by optimizing a graph connectivity metric. CAPE uses the graph representation of the ground truth to select node pairs and determine their corresponding paths within the predicted segmentation through a shortest-path algorithm. Using this, we penalize both disconnections and false positive connections, effectively promoting the model to preserve topological correctness. Experiments on 2D and 3D datasets, including neuron and blood vessel tracing demonstrate that CAPE significantly improves topology-aware metrics and outperforms state-of-the-art methods.

Computation of L_CAPE

The ground-truth mask is skeletonised and converted to a graph. In every iteration we sample a connected vertex pair and find its shortest path with Dijkstra’s algorithm; this path has zero cost because the label is perfectly connected. The two vertices are then shifted to the lowest-cost pixels within a small window on the network’s predicted distance map, providing tolerance to minor mis-alignments. Fixing those endpoints, we apply Dijkstra again on the prediction. The L_CAPE for this pair is the sum of squared distance values along that predicted path—close to zero when connectivity is correct and larger when a gap is present. Averaging such costs over many pairs and adding them to a standard pixel-wise loss (e.g. MSE) yields the total objective used for back-propagation.

Masking Strategy

We convert the ground-truth path into a binary mask \(M\) and dilate it by 10 pixels to obtain \(M_{\text{dilated}}\). Multiplying the predicted distance map \(\hat{y}\) by this dilated mask assigns an infinite cost to every pixel outside the mask. Dijkstra’s algorithm is then run between the projected vertices \(v'_1\) and \(v'_2\) on the masked map. Because the search is confined to the dilated region, it cannot bypass a disconnection through a loop, and the number of pixels examined is smaller, which reduces the computation time of the loss.

Vertex Shifting

Shifted

Not shifted

Directly mapping a graph vertex to the corresponding pixel in the predicted distance map \(\hat{y}\) can penalise the model for errors introduced by noisy center-line annotations. To mitigate this, each vertex \(v\) is relocated to the pixel within a radius of \(r = 5\) px that minimises \(\hat{y}(v)\). Formally, \[ v' \;=\; \arg\min_{u \in \mathcal{N}_r(v)} \hat{y}(u), \] where \(\mathcal{N}_r(v)\) is the \((2r{+}1)\times(2r{+}1)\) neighbourhood (or cubic analogue). This keeps endpoints aligned with the predicted center line, giving a more stable start for the masked shortest-path search and smoother gradients during optimisation.

Algorithm

Overall algorithm. We iteratively samples vertex pairs until every edge of the ground-truth graph has been processed. For each pair the shortest path is extracted, the corresponding mask is generated and then it is dilated. The vertices are shifted on the prediction to compensate for the noisy center-line annotations. Then, shortest path within the masked region is computed. The squared distance values accumulated along all predicted paths define L_CAPE.

Usage / CAPE

The CAPE loss module compares shortest paths in predicted and ground-truth distance maps. It samples random node pairs from the ground-truth graph (or skeletonized mask) and penalizes deviations in path cost via Dijkstra’s algorithm.

class CAPE(nn.Module):
    def __init__(
        self, window_size=128, three_dimensional=False, dilation_radius=10, shifting_radius=5, is_binary=False, 
        distance_threshold=20, single_edge=False
    )

Parameters

window_size : int (default=128): Size of each square patch (2D) or cubic volume (3D) to process at a time.
three_dimensional : bool (default=False): If True, operate on 3D volumes (D×H×W); otherwise, on 2D images (H×W).
dilation_radius : int (default=10): Radius (in pixels/voxels) to dilate ground-truth paths for masking.
shifting_radius : int (default=5): Radius (in pixels/voxels) to refine start/end points to lowest-cost nearby pixels.
is_binary : bool (default=False): If True, operate on binary maps rather than distance maps.
distance_threshold : float (default=20): Maximum value used for clipping ground-truth distance maps.
single_edge : bool (default=False): If True, sample a single edge at a time; otherwise, sample a path.

Returns

loss : torch.Tensor: A scalar tensor representing the average CAPE loss over the batch.

Notes

Predictions must be a torch.Tensor of shape (batch, H, W) for 2D or (batch, D, H, W) for 3D.
Ground truths can be a list of graphs in networkx.Graph format, or images (np.ndarray or torch.Tensor) of the same shape as prediction.

Usage / Graph extraction

The utils folder contains our implementation of two key functions—graph_from_skeleton_2D and graph_from_skeleton_3D which CAPE uses. Each one turns an skeleton mask into an undirected networkx.Graph. We also include helper functions for cropping these graphs into smaller patches under the same directory.

For large datasets it is faster to build graphs once and cache them rather than regenerate them at every training step. Our helper extract_graph.py wraps the graph_from_skeleton_2D and graph_from_skeleton_3D functions: each binary .npy mask is converted to a networkx.Graph and the result is saved as .gpickle.

Saving examples:

# 2-D masks → graphs (saved to data_as_graphs)
python extract_graph.py npy_images

# 3-D volumes → graphs (custom output directory)
python extract_graph.py brain_vols --dim 3 --out_dir brain_graphs

--dim {2 | 3}: Choose the 2-D or 3-D skeleton-to-graph routine · default 2
--threshold <T>: Binarise masks whose values are not already 0/1 · default 0.5
--out_dir <DIR>: Destination folder for the generated .gpickle files · default data_as_graph

Loading example:

from extract_graph import read_gpickle

gpickle_path = "data_as_graph/example_graph.gpickle"
G = read_gpickle(gpickle_path)

Results

	Input	Label	Perc	clDice	MALIS	CAPE
CREMI
DRIVE
Brain

Qualitative comparison of the test results in 2D and 3D datasets. The arrows highlight regions that remain disconnected in alternative methods but are successfully connected using CAPE. The connectivity improves significantly when our approach is used.

BibTeX

@misc{esmaeilzadeh2025CAPE,
        title={CAPE: Connectivity-Aware Path Enforcement Loss for Curvilinear Structure Delineation}, 
        author={Elyar Esmaeilzadeh and Ehsan Garaaghaji and Farzad Hallaji Azad and Doruk Oner},
        year={2025},
        eprint={2504.00753},
        archivePrefix={arXiv},
        primaryClass={cs.CV},
        url={https://arxiv.org/abs/2504.00753}, 
  }