NexToU: Efficient Topology-Aware U-Net for Medical Image Segmentation

TL;DR

NexToU introduces a hybrid GNN-based architecture that effectively captures topological relationships in medical images, improving segmentation accuracy across diverse datasets by integrating global and local topological features.

Contribution

The paper presents NexToU, a novel hybrid architecture combining Pool GNN and Swin GNN modules with a topological interaction module, enhancing topological feature learning in medical image segmentation.

Findings

Outperforms state-of-the-art methods on three diverse datasets

Effectively encodes topological constraints among anatomical structures

Reduces computational costs while capturing global and local topology

Abstract

Convolutional neural networks (CNN) and Transformer variants have emerged as the leading medical image segmentation backbones. Nonetheless, due to their limitations in either preserving global image context or efficiently processing irregular shapes in visual objects, these backbones struggle to effectively integrate information from diverse anatomical regions and reduce inter-individual variability, particularly for the vasculature. Motivated by the successful breakthroughs of graph neural networks (GNN) in capturing topological properties and non-Euclidean relationships across various fields, we propose NexToU, a novel hybrid architecture for medical image segmentation. NexToU comprises improved Pool GNN and Swin GNN modules from Vision GNN (ViG) for learning both global and local topological representations while minimizing computational costs. To address the containment and exclusion relationships among various anatomical structures, we reformulate the topological interaction (TI) module based on the nature of binary trees, rapidly encoding the topological constraints into NexToU. Extensive experiments conducted on three datasets (including distinct imaging dimensions, disease types, and imaging modalities) demonstrate that our method consistently outperforms other state-of-the-art (SOTA) architectures. All the code is publicly available at https://github.com/PengchengShi1220/NexToU.