Context-aware virtual adversarial training for anatomically-plausible segmentation

TL;DR

This paper introduces CaVAT, a novel semi-supervised segmentation method that enforces anatomical plausibility by incorporating topological constraints through adversarial training and reinforcement learning.

Contribution

It presents a generic framework combining adversarial training and reinforcement learning to enforce complex anatomical constraints in segmentation networks.

Findings

Produces anatomically plausible segmentations with high accuracy

Effectively enforces topological constraints like connectivity

Applicable to various segmentation models

Abstract

Despite their outstanding accuracy, semi-supervised segmentation methods based on deep neural networks can still yield predictions that are considered anatomically impossible by clinicians, for instance, containing holes or disconnected regions. To solve this problem, we present a Context-aware Virtual Adversarial Training (CaVAT) method for generating anatomically plausible segmentation. Unlike approaches focusing solely on accuracy, our method also considers complex topological constraints like connectivity which cannot be easily modeled in a differentiable loss function. We use adversarial training to generate examples violating the constraints, so the network can learn to avoid making such incorrect predictions on new examples, and employ the Reinforce algorithm to handle non-differentiable segmentation constraints. The proposed method offers a generic and efficient way to add any constraint on top of any segmentation network. Experiments on two clinically-relevant datasets show our method to produce segmentations that are both accurate and anatomically-plausible in terms of region connectivity.