Whole Brain Segmentation with Full Volume Neural Network

TL;DR

This paper introduces a full volume neural network approach for whole brain segmentation, leveraging complete 3D MRI data to improve accuracy over slice-based methods, using HRNet architecture and mixed precision training.

Contribution

It proposes a novel full volume segmentation framework with HRNet architecture, enhancing spatial detail and performance in brain MRI segmentation tasks.

Findings

Outperforms existing methods on benchmark datasets

Uses full volume input for better contextual understanding

Achieves state-of-the-art segmentation accuracy

Abstract

Whole brain segmentation is an important neuroimaging task that segments the whole brain volume into anatomically labeled regions-of-interest. Convolutional neural networks have demonstrated good performance in this task. Existing solutions, usually segment the brain image by classifying the voxels, or labeling the slices or the sub-volumes separately. Their representation learning is based on parts of the whole volume whereas their labeling result is produced by aggregation of partial segmentation. Learning and inference with incomplete information could lead to sub-optimal final segmentation result. To address these issues, we propose to adopt a full volume framework, which feeds the full volume brain image into the segmentation network and directly outputs the segmentation result for the whole brain volume. The framework makes use of complete information in each volume and can be implemented easily. An effective instance in this framework is given subsequently. We adopt the $3$D high-resolution network (HRNet) for learning spatially fine-grained representations and the mixed precision training scheme for memory-efficient training. Extensive experiment results on a publicly available $3$D MRI brain dataset show that our proposed model advances the state-of-the-art methods in terms of segmentation performance. Source code is publicly available at https://github.com/microsoft/VoxHRNet.