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

TL;DR

This paper introduces PI-Att, a topology-aware loss function for segmentation networks that uses persistence image representations to improve topological correctness in medical image segmentation, adapting during training for better results.

Contribution

The paper presents a novel topology-aware loss function using persistence images and an adaptive mechanism to enhance segmentation accuracy in medical images.

Findings

Improved topological accuracy in segmentation results.

Effective in aorta and vessel segmentation tasks.

Adaptive persistence image calculation enhances learning.

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.