InfiNet: Fully Convolutional Networks for Infant Brain MRI Segmentation

TL;DR

InfiNet is a novel, efficient fully convolutional neural network designed for fast, accurate voxel-wise segmentation of infant brain MRI images, effectively handling multi-modal data and class imbalance.

Contribution

The paper introduces InfiNet, a new architecture with a unique decoder upsampling method that improves segmentation performance and efficiency for infant brain MRI images.

Findings

Achieves whole-volume segmentation in under 50 seconds.

Demonstrates competitive performance against state-of-the-art architectures.

Effectively handles multi-modal MRI data and class imbalance.

Abstract

We present a novel, parameter-efficient and practical fully convolutional neural network architecture, termed InfiNet, aimed at voxel-wise semantic segmentation of infant brain MRI images at iso-intense stage, which can be easily extended for other segmentation tasks involving multi-modalities. InfiNet consists of double encoder arms for T1 and T2 input scans that feed into a joint-decoder arm that terminates in the classification layer. The novelty of InfiNet lies in the manner in which the decoder upsamples lower resolution input feature map(s) from multiple encoder arms. Specifically, the pooled indices computed in the max-pooling layers of each of the encoder blocks are related to the corresponding decoder block to perform non-linear learning-free upsampling. The sparse maps are concatenated with intermediate encoder representations (skip connections) and convolved with trainable filters to produce dense feature maps. InfiNet is trained end-to-end to optimize for the Generalized Dice Loss, which is well-suited for high class imbalance. InfiNet achieves the whole-volume segmentation in under 50 seconds and we demonstrate competitive performance against multiple state-of-the art deep architectures and their multi-modal variants.