Shape Equivariant Learning for Robust MRI Segmentation

Ainkaran Santhirasekaram, Mathias Winkler, Andrea G. Rockall, Ben Glocker

Keywords: Shape Equivariance, Robustness, Segmentation, MRI

TL;DR: We propose learning a discrete shape equivariant embedding space for robust segmentation.

Abstract: The reliability of deep learning based segmentation models is essential to the safe translation of these models into clinical practise. Unfortunately, these models are sensitive to distributional shifts. This is particularly notable in MRI, where there is a large variation of acquisition protocols across different domains leading to varying textural profiles. We hypothesise that the constrained anatomical variability across subjects can be leveraged to discretize the latent space to a dictionary of shape components. We achieve this by using multiple MRI sequences to learn texture invariant and shape equivariant features which are used to construct a shape dictionary using vector quantisation. This dictionary is then sampled to compose the segmentation output. Our method achieves SOTA performance in the task of single domain generalisation (SDG) for prostate zonal segmentation.