Large Scale Supervised Pretraining For Traumatic Brain Injury Segmentation

TL;DR

This work tackles the challenging task of segmenting lesions in Moderate to Severe Traumatic Brain Injury on T1-weighted MRI by leveraging a large-scale, MultiTalent-inspired supervised pretraining across diverse datasets to form a robust foundation model. The foundation is a ResencL U-Net implemented within the nnU-Net framework, which is subsequently fine-tuned on msTBI data with three learning-rate strategies. Development results show the pretrained model outperforms the baseline by up to 2 Dice points, and test performance could have improved further if a numerical precision issue during ensembling had been resolved.Overall, the approach demonstrates the value of cross-dataset pretraining for neuroimaging segmentation and highlights practical considerations for model deployment and fairness in clinical datasets.

Abstract

The segmentation of lesions in Moderate to Severe Traumatic Brain Injury (msTBI) presents a significant challenge in neuroimaging due to the diverse characteristics of these lesions, which vary in size, shape, and distribution across brain regions and tissue types. This heterogeneity complicates traditional image processing techniques, resulting in critical errors in tasks such as image registration and brain parcellation. To address these challenges, the AIMS-TBI Segmentation Challenge 2024 aims to advance innovative segmentation algorithms specifically designed for T1-weighted MRI data, the most widely utilized imaging modality in clinical practice. Our proposed solution leverages a large-scale multi-dataset supervised pretraining approach inspired by the MultiTalent method. We train a Resenc L network on a comprehensive collection of datasets covering various anatomical and pathological structures, which equips the model with a robust understanding of brain anatomy and pathology. Following this, the model is fine-tuned on msTBI-specific data to optimize its performance for the unique characteristics of T1-weighted MRI scans and outperforms the baseline without pretraining up to 2 Dice points.

Figures (3)

• Figure 1: Prediction performance by age group (right) and the number of patients per age group in the dataset (left). We excluded all cases where no ground truth foreground was annotated, and our method made no predictions.
• Figure 2: Prediction performance by Time Since Injury (TSI) group (right) and the number of patients per TSI group in the dataset (left). We excluded all cases where no ground truth foreground was annotated, and our method made no predictions.
• Figure 3: Qualitative Segmentation Results: The segmentation results for the first two patients closely align with the ground truth annotations. In contrast, the last image was chosen due to a significant discrepancy between the prediction and the ground truth.