Automated Label Merging and Splitting for Memory-Efficient Whole Brain Parcellation

Core Concepts

A method for significantly reducing the number of labels required for whole brain parcellation while maintaining segmentation accuracy by automatically merging spatially separate labels and recovering the original labels during inference.

Abstract

The proposed method for whole brain parcellation consists of the following key steps:

Preprocessing: Brain images are aligned to the MNI152 template and resampled.

Label Merging:

A distance matrix is calculated between original labels based on their spatial locations in the training set.
An adjacency matrix is constructed by applying thresholds on the distance and volume ratio between labels.
A greedy graph coloring algorithm is used to automatically group and merge multiple spatially separate labels that satisfy the distance and volume ratio constraints.
The merged labels are assigned unique IDs.

CNN Training:

A 3D U-Net model is trained to predict the merged labels, reducing the effective number of labels and memory requirements.

Label Splitting:

A fuzzy prior map is created from the training set label volumes and enhanced with a Euclidean Distance Transform.
Influence region maps are generated for each merged label based on the fuzzy prior.
At inference, the predicted merged labels are split back into the original labels using the influence region maps.

The experiments on three public datasets (Mindboggle101, AOMIC, IXI) demonstrate that the proposed merge-and-split approach can reduce the number of labels by up to 68% while achieving segmentation accuracy comparable to the baseline model without label merging and splitting. This leads to significant reductions in training and inference time (43-49%) as well as GPU memory requirements (50% during training, 21% during inference).

Stats

The number of labels was reduced by a factor of at least three for both the AOMIC and IXI training sets.
GPU memory required during training was halved and reduced by 21% during inference.
Epoch times during training were reduced by more than 43% and inference time was halved.

Quotes

"The proposed method can be applied to all semantic segmentation tasks with a large number of spatially separate classes within an atlas-based prior."
"While the results presented here were obtained with the popular U-Net architecture, the method can in principle be combined with any segmentation approach to reduce the effective number of labels."

Key Insights Distilled From

Label merge-and-split: A graph-colouring approach for memory-efficient brain parcellation

by Aaron Kujawa... at arxiv.org 04-17-2024

https://arxiv.org/pdf/2404.10572.pdf

Deeper Inquiries

How could the label merging and splitting approach be extended to work without the need for affine registration to a common space

To extend the label merging and splitting approach without relying on affine registration to a common space, one could consider utilizing registration-free techniques such as non-rigid image registration. Instead of aligning all images to a common template, the method could incorporate spatial transformations directly into the label merging and splitting process. By calculating spatial distances and volume ratios in the native space of each image, the approach can still determine which labels should be merged while accounting for individual anatomical variabilities. This way, the method can operate without the need for a prior affine registration step, making it more flexible and applicable to a wider range of datasets.

What other constraints or criteria could be introduced in the graph construction and coloring process to control which labels are merged, beyond just spatial distance and volume ratio

In the graph construction and coloring process for label merging, additional constraints or criteria can be introduced to further refine the merging decisions beyond just spatial distance and volume ratio. One possible constraint could involve incorporating anatomical priors or knowledge about the relationships between different brain structures. For example, certain labels that are known to be functionally or anatomically related could be grouped together or prevented from merging with unrelated labels. By integrating domain-specific information into the graph construction, the method can make more informed decisions about which labels should be merged based on their biological relevance and spatial characteristics.

How could the proposed method be adapted to handle combined tasks, such as whole brain parcellation and lesion segmentation, where certain labels should not be merged with others

To adapt the proposed method for handling combined tasks like whole brain parcellation and lesion segmentation, where specific labels should not be merged with others, a tailored approach can be implemented. One strategy could involve introducing label-specific constraints during the graph construction phase. For instance, labels representing lesions could be assigned unique properties or connectivity rules in the graph to prevent them from being merged with healthy brain structures. By incorporating task-specific constraints into the label merging and splitting process, the method can ensure that certain labels remain distinct and unaffected by the merging operations, enabling accurate segmentation for combined tasks without compromising the integrity of individual structures.