CAT: Constrained Adversarial Training for Anatomically-Plausible Semi-Supervised Segmentation.

Abstract

Deep learning models for semi-supervised medical image segmentation have achieved unprecedented performance for a wide range of tasks. Despite their high accuracy, these models may however yield predictions that are considered anatomically impossible by clinicians. Moreover, incorporating complex anatomical constraints into standard deep learning frameworks remains challenging due to their non-differentiable nature. To address these limitations, we propose a Constrained Adversarial Training (CAT) method that learns how to produce anatomically plausible segmentations. Unlike approaches focusing solely on accuracy measures like Dice, our method considers complex anatomical constraints like connectivity, convexity, and symmetry which cannot be easily modeled in a loss function. The problem of non-differentiable constraints is solved using a Reinforce algorithm which enables to obtain a gradient for violated constraints. To generate constraint-violating examples on the fly, and thereby obtain useful gradients, our method adopts an adversarial training strategy which modifies training images to maximize the constraint loss, and then updates the network to be robust to these adversarial examples. The proposed method offers a generic and efficient way to add complex segmentation constraints on top of any segmentation network. Experiments on synthetic data and four clinically-relevant datasets demonstrate the effectiveness of our method in terms of segmentation accuracy and anatomical plausibility.