PI-Att: Topology Attention for Segmentation Networks through Adaptive Persistence Image Representation

TL;DR

This work tackles topology-aware segmentation in medical imaging under limited training data. It introduces PI-Att, a loss that computes topological dissimilarity via persistence image representations of ground-truth and predicted maps, avoiding unstable point-matching of persistence diagrams. An adaptive scheduler adjusts the persistence image weighting with a function $\omega(y,\gamma)$ and update rule $\gamma^{t+1}= \gamma^{t} (1 - \lambda \cdot {\rm CE}^{t} \cdot {\rm TD}^{t})$ (with $\gamma^0=2$ and $\lambda=0.0005$) to emphasize topology outline early and topology details later. Across two CT vessel datasets and architectures, PI-Att yields superior segmentation accuracy and topology-consistency metrics compared with baselines and prior topology-aware losses, highlighting its potential as a data-efficient regularizer.

Abstract

Segmenting multiple objects (e.g., organs) in medical images often requires an understanding of their topology, which simultaneously quantifies the shape of the objects and their positions relative to each other. This understanding is important for segmentation networks to generalize better with limited training data, which is common in medical image analysis. However, many popular networks were trained to optimize only pixel-wise performance, ignoring the topological correctness of the segmentation. In this paper, we introduce a new topology-aware loss function, which we call PI-Att, that explicitly forces the network to minimize the topological dissimilarity between the ground truth and prediction maps. We quantify the topology of each map by the persistence image representation, for the first time in the context of a segmentation network loss. Besides, we propose a new mechanism to adaptively calculate the persistence image at the end of each epoch based on the network's performance. This adaptive calculation enables the network to learn topology outline in the first epochs, and then topology details towards the end of training. The effectiveness of the proposed PI-Att loss is demonstrated on two different datasets for aorta and great vessel segmentation in computed tomography images.

Figures (8)

• Figure 1: (a) 3D visualization of the aortic arch and great vessels (large arteries and veins). (b) Segmentation maps of axial slice examples. Vessels in an axial slice are not randomly distributed in the body but found in a particular topology, which differs according to position of the axial slice with respect to the heart.
• Figure 2: Illustration of 1-dimensional persistent homology and the process of constructing a persistence diagram using a distance filtration on contour points. This filtration was selected only for illustration, as it is easier to see how persistent homology simultaneously models shape and geometry. Our model will use another filtration also defined on the contours but associated to a density function as it gives better topological summary. The bottom indicates the filtration process with the barcodes and depicts the resulting persistence diagram. The top shows the evolved contour points (blue) and holes (red) at six time indices, from (a) to (f) and with dashed lines. (a) At time 0, five holes are born inside the initial contours of the objects. (b) Two holes die as their corresponding objects are smaller than the others. (c) Another hole is born inside the object at the upper-right corner. The shape of this object (figure-eight-like shape) is different from the others. (d) All holes inside the objects die but another hole between three objects is born. Since these objects are closer to each other than the remaining two, a hole between them is born before the final hole illustrated at time (e). The barcodes of the last two holes are associated with the objects' geometry. (f) All holes die at the end.
• Figure 3: Schematic overview of the proposed PI-Att loss.
• Figure 4: Illustration of the proposed adaptive scheduler mechanism for three different steps of network training. At the top, the existing ground truth of an example image and its persistence images calculated in different steps are given. At the bottom, the segmentation map predicted in a given step and its persistence image are given. All persistence images are calculated using $\gamma^t$ of the corresponding step. The scheduler blocks used in this mechanism are shown as green and red boxes, and their operations are depicted on the right.
• Figure 5: For the in-house dataset, visual results on exemplary CT images. (a) Ground truths. Visual results obtained by (b) Baseline, (c) ActiveContourLoss chen2019learning, (d) HausdorffDistanceLoss karimi2019reducing, (e) TopologicalLossLikelihoodFiltration hu2019topologypreservingdeep, (f) TopologicalLossDistanceFiltration ozcelik2023topologyawareloss, (g) clDiceTopologyPreservingLoss shit2021cldicenovel, (h) NonAdaptivePersistenceImageLoss, and (j) proposed PI-Att:PersistenceImageLoss. Vessel pixels are shown in green on the CT images.