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

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.

💡 Research Summary

The paper addresses a fundamental problem in biomedical image segmentation: while deep convolutional networks have dramatically improved pixel‑wise accuracy for curvilinear structures such as neuronal processes and blood vessels, conventional loss functions (cross‑entropy, Dice, MSE) do not enforce the global topological correctness required for downstream analysis. Existing topology‑aware losses—based on persistent homology, Rand index, or soft‑skeletonization—either optimize indirect proxies of connectivity or produce sparse gradients that make training unstable, especially in the presence of loops.

To directly optimize connectivity, the authors propose CAPE (Connectivity‑Aware Path Enforcement), a differentiable loss that operates on graph representations of the ground‑truth and the network’s predicted distance map. The method proceeds as follows: (1) From the ground‑truth graph G = (V,E), a pair of connected vertices (v₁, v₂) is randomly sampled. (2) Dijkstra’s algorithm computes the true shortest path path_G between them. (3) The predicted network outputs a per‑pixel distance map ŷ, where each pixel encodes its Euclidean distance to the nearest foreground structure; this formulation yields linearly increasing values for larger gaps, making the loss more sensitive to disconnections. (4) The sampled vertices are projected onto ŷ by locating local minima within a 7 × 7 window, producing refined positions v₁′ and v₂′. (5) A binary mask M is rendered from path_G, dilated by 10 pixels to form M_dilated, and element‑wise multiplied with ŷ. (6) Dijkstra’s algorithm is run again on the masked distance map to obtain the predicted shortest path path_ŷ between v₁′ and v₂′. (7) The cost of path_ŷ is defined as the sum of squared distance values along the path: cost(path_ŷ) = ∑_{n∈path_ŷ} ŷ(n)². Because the ground‑truth path is assumed perfectly connected, its cost is set to zero, and the CAPE loss for this pair reduces to cost(path_ŷ).

The algorithm iterates until all edges of G have been processed, removing each edge after its associated pair has been evaluated to avoid duplicate counting. The total CAPE loss is the sum of costs over all sampled paths. To provide dense supervision across the whole image, CAPE is combined with a per‑pixel mean‑squared‑error (MSE) term: L_total = L_MSE + α L_CAPE, where α balances the two objectives.

Experiments were conducted on three datasets: (i) CREMI – 2‑D electron microscopy images of Drosophila neuronal arbors, (ii) DRIVE – 2‑D retinal fundus images of blood vessels, and (iii) Brain – 3‑D light‑microscopy volumes of mouse brain vasculature. For 2‑D tasks a standard U‑Net was used; for 3‑D a 3‑D U‑Net with identical architecture was employed. Models were trained for 10 k (2‑D) or 50 k (3‑D) iterations using Adam (lr = 1e‑3, weight decay = 1e‑3). Baselines included plain MSE, a perceptual loss using a pretrained VGG‑19, clDice (soft‑skeletonization), and inverse MALIS (both 2‑D and projection‑based 3‑D variants).

Evaluation metrics comprised pixel‑wise CCQ (Correctness, Completeness, Quality) and Dice, as well as topology‑aware metrics: APLS (Average Path Length Similarity) and TL‑TS (Too‑Long‑Too‑Short). CAPE consistently outperformed all baselines on APLS and TL‑TS across the three datasets, achieving improvements of 3–7 percentage points. Pixel‑wise scores remained comparable, with CAPE often matching or slightly exceeding the best baseline. Qualitative visualizations showed that structures broken in other methods were correctly re‑connected when CAPE was used, confirming that the loss effectively penalizes gaps while tolerating minor spatial shifts thanks to the masking strategy.

Key contributions are: (1) a fully differentiable loss that directly optimizes a graph‑based connectivity metric, (2) the use of a distance‑map representation to provide a smooth penalty for disconnections, (3) a masking technique that forces the predicted path to stay near the ground‑truth skeleton, avoiding shortcuts through loops, and (4) extensive validation on both 2‑D and 3‑D biomedical datasets demonstrating superior topological fidelity.

Limitations include dependence on the random sampling of vertex pairs (requiring sufficient samples and an appropriate α to balance gradients) and sensitivity to the size of the dilation mask and distance‑map scaling, which may need dataset‑specific tuning. Future work will explore adaptive mask generation, multi‑scale path sampling, and extensions to other domains such as road network extraction or surgical planning, aiming to make the connectivity‑aware loss more universally applicable.