An End-to-End learnable Flow Regularized Model for Brain Tumor Segmentation

TL;DR

This paper introduces an end-to-end trainable neural network integrated with energy-based optimization for brain tumor segmentation, improving accuracy by learning features directly from data and solving segmentation via primal-dual methods.

Contribution

It combines deep neural network features with energy minimization models in an end-to-end framework for biomedical image segmentation, which is a novel integration.

Findings

Enhanced segmentation accuracy demonstrated on brain tumor datasets.

Segmentation contours evolve actively through iterations, aiding diagnosis.

The method balances sensitivity and boundary smoothness effectively.

Abstract

Many segmentation tasks for biomedical images can be modeled as the minimization of an energy function and solved by a class of max-flow and min-cut optimization algorithms. However, the segmentation accuracy is sensitive to the contrasting of semantic features of different segmenting objects, as the traditional energy function usually uses hand-crafted features in their energy functions. To address these limitations, we propose to incorporate end-to-end trainable neural network features into the energy functions. Our deep neural network features are extracted from the down-sampling and up-sampling layers with skip-connections of a U-net. In the inference stage, the learned features are fed into the energy functions. And the segmentations are solved in a primal-dual form by ADMM solvers. In the training stage, we train our neural networks by optimizing the energy function in the primal form with regularizations on the min-cut and flow-conservation functions, which are derived from the optimal conditions in the dual form. We evaluate our methods, both qualitatively and quantitatively, in a brain tumor segmentation task. As the energy minimization model achieves a balance on sensitivity and smooth boundaries, we would show how our segmentation contours evolve actively through iterations as ensemble references for doctor diagnosis.