A Learning-based Framework for Topology-Preserving Segmentation using Quasiconformal Mappings

TL;DR

TPSN presents a learning-based framework that preserves topology during image segmentation by deforming a template mask through a bijective quasiconformal map. It combines a Deformation Estimation Network (DEN) with a Beltrami Adjustment Module (BAM) to enforce bijectivity via the Beltrami coefficient, supplemented by a Linear Beltrami Solver for reconstruction; the framework supports supervised and unsupervised training and extends to multi-level and multi-object scenarios. Key innovations include the ReLU-Jacobian regularizer, the truncation of Beltrami coefficients to ensure $|\\tilde{\\mu}|<1$, the BS-Net-based efficient mapping, and the Fill First Dig Second strategy to handle complex topologies. Empirical results across 2D and 3D medical datasets demonstrate topology-preserving segmentation with competitive Dice scores and superior geometric/topological metrics, along with robustness to corrupted data and unsupervised capabilities.

Abstract

We propose the Topology-Preserving Segmentation Network, a deformation-based model that can extract objects in an image while maintaining their topological properties. This network generates segmentation masks that have the same topology as the template mask, even when trained with limited data. The network consists of two components: the Deformation Estimation Network, which produces a deformation map that warps the template mask to enclose the region of interest, and the Beltrami Adjustment Module, which ensures the bijectivity of the deformation map by truncating the associated Beltrami coefficient based on Quasiconformal theories. The proposed network can also be trained in an unsupervised manner, eliminating the need for labeled training data. This is achieved by incorporating an unsupervised segmentation loss. Our experimental results on various image datasets show that TPSN achieves better segmentation accuracy than state-of-the-art models with correct topology. Furthermore, we demonstrate TPSN's ability to handle multiple object segmentation.

Figures (13)

• Figure 1: The TPSN network architecture consists of two components: the Deformation Estimation Network (represented by the yellow box) and the Beltrami Adjustment Module (represented by the blue box). For a segmentation task for $q$ classes of objects, the Deformation Estimation Network utilizes an encoder-decoder architecture that takes in a prior $M_{temp}$ and an image $I$ with $C$ channels, and produces $q$ mappings for each class. This mapping is then used to deform prior masks to enclose regions of interest $M_{pred}$. The component is designed with ReLU-Jacobian regularization, which promotes the output of a topology-preserving mapping. The Beltrami Adjustment Module consists of two convolutional layers activated by a $Tanh$. This component ensures that the mapping is topology-preserving by truncating $\mu$ into $\Tilde{\mu}$. The Beltrami Solver Network chen2021deep reconstructs the quasiconformal mapping $\Tilde{f}$ that corresponds to the truncated $\Tilde{\mu}$. The corresponding prediction mask $\widetilde{M}_{pred}$ is then obtained.
• Figure 2: Quasi-conformal maps infinitesimal circles to ellipses. The Beltrami coefficient measures the distortion or dilation of the ellipse under the QC map.
• Figure 3: Illustration for Projection construction.
• Figure 4: Illustration of the multi-level TPSN. The original image is first downsampled to a low-resolution image, which is then inputted into a TPSN with a prior of the same size. The predicted mask is subsequently upsampled and used as the prior mask in the next layer of the TPSN.
• Figure 5: Illustration of Fill First, Dig Second strategy. A disk without holes is deformed to enclose the filled label mask. A hole will be dug out and further optimized to locate the inner boundary.

Theorems & Definitions (2)

• Theorem 1
• proof