Anatomical Priors for Image Segmentation via Post-processing with Denoising Autoencoders

1. Author

Agostina J. Larrazabal, Cesar Martinez, and Enzo Ferrante
Research Institute for Signals, Systems and Computational Intelligence, sinc(i),
FICH-UNL/CONICET, Santa Fe, Argentina

2. Abstract

In this work we propose Post-DAE, a post-processing method based on denoising autoencoders (DAE) trained using only segmentation masks.

3. Introduction

One of the tricks that enables the use of CNNs in large images (by reducing the number of learned parameters) is known as parameter sharing scheme.
This trick is especially useful for tasks like image classification, where invariance to translation is a desired property since objects may appear in any location.

3.1 Contributions

1. we show, for the first time, that DAE can be used as independent post-processing step to correct problematic and non-anatomically plausible masks produced by arbitrary segmentation methods.
2. we design a method that can be trained using segmentation-only datasets or anatomical masks coming from arbitrary image modalities, since the DAE is trained using only segmentation masks, and no intensity information is required during learning.
3. we validate Post-DAE in the context of lung segmentation in X-ray images, bench-marking with other classical post-processing method and showing its robustness by improving segmentation masks coming from both, CNN and RF-based classifiers.

3.2 Related Works

One popular strategy to incorporate prior knowledge about shape and topology into medical image segmentation is to modify the loss used to train the model.

Our post-processing method makes use of a DAE to impose shape priors, transforming any segmentation mask into an anatomically plausible one.

4 Anatomical Priors for Image Segmentation via Post-processing with DAE

4.1 Problem Statement

Given a dataset of unpaired anatomical segmentation masks ${\mathcal{D}}_{\mathcal{A}}={\left\{S_i^A\right\}}_{0\le i\le |{\mathcal{D}}_{\mathcal{A}}|}$
we aim at learning a model that can bring segmentations ${\mathcal{D}}_{\mathcal{P}}={\left\{S_i^P\right\}}_{0\le i\le |{\mathcal{D}}_{\mathcal{P}}|}$ predicted by arbitrary classifiers $P$ into an anatomically feasible space.

4.2 Denoising Autoencoders

DAEs are neural networks designed to reconstruct a clean input from a corrupted version of it.

4.3 Mask Degradation Strategy

The masks used to train the DAE were artificially degraded during training to simulate erroneous segmentations.
To this end, we randomly apply the following degradation functions $\varphi \left(S_i\right)$ to the ground truth masks $S_i$:

1. addition and removal of random geometric shapes (circles, ellipses, lines and rectangles) to simulate over and under segmentations
2. morphological operations (e.g. erosion, dilation, etc.) with variable kernels to perform more subtle mask modifications
3. random swapping of foreground-background labels in the pixels close to the mask borders

4.4 Post-processing with DAEs

We learn such low-dimensional anatomically plausible manifold using the aforementioned DAE.
Then, given a segmentation mask $S_i^P$ obtained with an arbitrary predictor $P$ (e.g. $CNN$ or $RF$), we project it into that manifold using $f_{enc}$ and reconstruct the corresponding anatomically feasible mask with $f_{dec}$.