Universal Topology Refinement for Medical Image Segmentation with Polynomial Feature Synthesis

TL;DR

This paper introduces a versatile, model-agnostic topology refinement method for medical image segmentation that uses synthetic polynomial-based masks to correct topological errors, improving accuracy without retraining models.

Contribution

The authors propose a universal, plug-and-play topology refinement network trained with synthetic polynomial-based masks, outperforming existing methods and compatible with various segmentation pipelines.

Findings

Outperforms existing topology correction methods

Compatible with multiple polynomial bases

Enhances segmentation accuracy when combined with learning models

Abstract

Although existing medical image segmentation methods provide impressive pixel-wise accuracy, they often neglect topological correctness, making their segmentations unusable for many downstream tasks. One option is to retrain such models whilst including a topology-driven loss component. However, this is computationally expensive and often impractical. A better solution would be to have a versatile plug-and-play topology refinement method that is compatible with any domain-specific segmentation pipeline. Directly training a post-processing model to mitigate topological errors often fails as such models tend to be biased towards the topological errors of a target segmentation network. The diversity of these errors is confined to the information provided by a labelled training set, which is especially problematic for small datasets. Our method solves this problem by training a model-agnostic topology refinement network with synthetic segmentations that cover a wide variety of topological errors. Inspired by the Stone-Weierstrass theorem, we synthesize topology-perturbation masks with randomly sampled coefficients of orthogonal polynomial bases, which ensures a complete and unbiased representation. Practically, we verified the efficiency and effectiveness of our methods as being compatible with multiple families of polynomial bases, and show evidence that our universal plug-and-play topology refinement network outperforms both existing topology-driven learning-based and post-processing methods. We also show that combining our method with learning-based models provides an effortless add-on, which can further improve the performance of existing approaches.