Disconnect to Connect: A Data Augmentation Method for Improving Topology Accuracy in Image Segmentation

TL;DR

This work addresses the difficulty of topology-aware image segmentation for thin tubular structures by introducing CoLeTra, a data augmentation technique that erases parts of tubular structures during training while preserving ground-truth labels to teach the model that visually disconnected segments are actually connected. CoLeTra operates with two intuitive hyper-parameters and can be combined with existing topology losses without requiring extra GPU memory. Extensive experiments across four datasets and two architectures show consistent reductions in Betti errors and often improvements in Dice and HD95, with robust performance across settings. The method also contributes a Narwhal dataset to promote topology-focused evaluation and provides code for easy integration into existing pipelines, offering a practical tool to enhance topology accuracy in segmentation tasks.

Abstract

Accurate segmentation of thin, tubular structures (e.g., blood vessels) is challenging for deep neural networks. These networks classify individual pixels, and even minor misclassifications can break the thin connections within these structures. Existing methods for improving topology accuracy, such as topology loss functions, rely on very precise, topologically-accurate training labels, which are difficult to obtain. This is because annotating images, especially 3D images, is extremely laborious and time-consuming. Low image resolution and contrast further complicates the annotation by causing tubular structures to appear disconnected. We present CoLeTra, a data augmentation strategy that integrates to the models the prior knowledge that structures that appear broken are actually connected. This is achieved by creating images with the appearance of disconnected structures while maintaining the original labels. Our extensive experiments, involving different architectures, loss functions, and datasets, demonstrate that CoLeTra leads to segmentations topologically more accurate while often improving the Dice coefficient and Hausdorff distance. CoLeTra's hyper-parameters are intuitive to tune, and our sensitivity analysis shows that CoLeTra is robust to changes in these hyper-parameters. We also release a dataset specifically suited for image segmentation methods with a focus on topology accuracy. CoLetra's code can be found at https://github.com/jmlipman/CoLeTra.

Figures (7)

• Figure 1: a) Labels missing important areas, compromising the structure's true topology. b) CoLeTra yields segmentations topologically more accurate.
• Figure 2: Overview of CoLeTra data augmentation.
• Figure 3: CoLeTra incorporates prior knowledge about structures' connectedness. Top: training and validation set images from the synthetic dataset. Bottom: Results in the validation set, where top/bottom is without/with CoLeTra.
• Figure 4: Example segmentations on four datasets and the best-performing loss functions, with and without CoLeTra.
• Figure 5: Betti error and number of small ($<$ 30 pixels) connected components achieved with CoLeTra on the DRIVE dataset with different number of patches (a,b) and patch sizes (c,d). (DA---data augmentation)