Universal Topology Refinement for Medical Image Segmentation with Polynomial Feature Synthesis

TL;DR

Evidence is shown that the universal plug-and-play topology refinement network outperforms both existing topology-driven learning-based and post-processing methods and that combining the method with learning-based models provides an effortless add-on, which can further improve the performance of existing approaches.

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.

Figures (4)

• Figure 1: Top row: Conventional topology-refinement post-processing networks are trained on, and biased towards the topological errors of a specific upstream segmentation network $g_\phi$, leading to errors when applied to a different $g_\phi$. Re-training is required when $g_\phi$ changes. Bottom row: To train an unbiased topology-refinement network $f_{\theta}$, our method synthesizes training samples from an unbiased distribution, as illustrated in the blue ellipses. This covers plausible real segmentations with different topological errors, generated by complete and orthogonal polynomials. During inference, segmentations $\mathbf{M}$ from $g_\phi$ can be refined by a universally trained $f_{\theta}$ network.
• Figure 2: Workflow for generating arbitrary 2D topology error maps from orthogonal polynomials (Chebyshev polynomials as examples).
• Figure 3: Qualitative results for topology preservation/refinement. Our approach successfully revises disconnected structures that should be connected.
• Figure 4: Ablation study for different types of polynomials. Left: Betti error across models trained with different polynomials. Right: Perturbation masks generated from Legendre, Hermite-Gaussian, and Chebyshev polynomials at order 4, 6, 8, and 10.