TEDS-Net: Enforcing Diffeomorphisms in Spatial Transformers to Guarantee Topology Preservation in Segmentations

MICCAI 2021

Madeleine K. Wyburd1
Nicola Dinsdale2
Ana I.L. Namburete1
Mark Jenkinson2,3,4

1Ultrasound NeuroImage Analysis Group, University of Oxford
2Wellcome Centre for Integrative Neuroimaging, FMRIB,University of Oxford
3Australian Institute for Machine Learning (AIML), Department of Computer Science, University of Adelaide
4South Australian Health and Medical Research Institute (SAHMRI)

[Paper]
[GitHub]




Abstract

Accurate topology is key when performing meaningful anatomical segmentations, however, it is often overlooked in traditional deep learning methods. In this work we propose TEDS-Net: a novel segmentation method that guarantees accurate topology. Our method is built upon a continuous diffeomorphic framework, which enforces topology preservation. However, in practice, diffeomorphic fields are represented using a finite number of parameters and sampled using methods such as linear interpolation, violating the theoretical guarantees. We therefore introduce additional modifications to more strictly enforce it. Our network learns how to warp a binary prior, with the desired topological characteristics, to complete the segmentation task. We tested our method on myocardium segmentation from an open-source 2D heart dataset. TEDS-Net preserved topology in 100% of the cases, compared to 90% from the U-Net, without sacrificing on Hausdorff Distance or Dice performance.


Pipeline

Schematic of TEDS-Net architecture. A CNN learns two initial fields, u, at different resolutions, from an input image X. The fields are enforced to be diffeomorphic using an activation function amplified through composition layers and "super" upsampled to 2x the resolution of the input. The bulk displacement samples a binary prior, P, generating YBulk, which is then sampled by the fine tuning field ΦFT. The asterisks show the elements removed during ablation studies.


Qualitative Results

Examples of myocardium manual annotations (green), compared to U-Net and TEDS-Net segmentations (pink) from the ACDC dataset


Parameter Searching

The effect that the number of integration layers (a) and the radius of the binary prior (b) had on segmentation performance and the diffeomorphic nature of the generated fields. Due to the image dimensions, 50 was the maximum radius used.


Acknowledgements

MW and ND is supported by the Engineering and Physical Sciences Research Council (EPSRC) and Medical Research Council (MRC) [grant number EP/L016052/1]. MJ is supported by the National Institute for Health Research (NIHR) Oxford Biomedical Research Centre (BRC), and this research was funded by the Well- come Trust [215573/Z/19/Z]. The Wellcome Centre for Integrative Neuroimaging is supported by core funding from the Wellcome Trust [203139/Z/16/Z]. AN is grateful for support from the UK Royal Academy of Engineering under the Engineering for Development Research Fellowships scheme.

This template was originally made by Phillip Isola and Richard Zhang for a colorful ECCV project; the code can be found here.