Automated ensemble method for pediatric brain tumor segmentation

Abstract

Brain tumors remain a critical global health challenge, necessitating advancements in diagnostic techniques and treatment methodologies. A tumor or its recurrence often needs to be identified in imaging studies and differentiated from normal brain tissue. In response to the growing need for age-specific segmentation models, particularly for pediatric patients, this study explores the deployment of deep learning techniques using magnetic resonance imaging (MRI) modalities. By introducing a novel ensemble approach using ONet and modified versions of UNet, coupled with innovative loss functions, this study achieves a precise segmentation model for the BraTS-PEDs 2023 Challenge. Data augmentation, including both single and composite transformations, ensures model robustness and accuracy across different scanning protocols. The ensemble strategy, integrating the ONet and UNet models, shows greater effectiveness in capturing specific features and modeling diverse aspects of the MRI images which result in lesion wise Dice scores of 0.52, 0.72 and 0.78 on unseen validation data and scores of 0.55, 0.70, 0.79 on final testing data for the ”enhancing tumor”, ”tumor core” and ”whole tumor” labels respectively. Visual comparisons further confirm the superiority of the ensemble method in accurate tumor region coverage. The results indicate that this advanced ensemble approach, building upon the unique strengths of individual models, offers promising prospects for enhanced diagnostic accuracy and effective treatment planning and monitoring for brain tumors in pediatric brains.

* contributed equally

1 Introduction

Brain tumor has been a major global health challenge, impacting not only adults but also children and adolescents [1]. In 2022, the United States alone reported an estimated 40,594 individuals, from infancy to 19 years of age, diagnosed with primary brain or other central nervous system (CNS) tumors. Among these, pilocytic astrocytoma was the predominant histopathologic group, accounting for 8,264 cases. Given that survival rates following a diagnosis are particularly low among infants, there is a huge need to develop age-specific segmentation models [2]. These models could enable precise and automated detection of tumor regions within pediatric brains, thereby facilitating more efficient and fast diagnosis and treatment.

With the increasing use of deep learning techniques in conjunction with different modalities of MRI, different models, especially the U-shaped architectures, have demonstrated precise and accurate performance across various medical image segmentation tasks [3, 4, 5]. Various models have been employed for whole brain segmentation, yet certain architectural designs have been observed to perform better in segmenting specific regions within the brain [6]. This phenomenon can be attributed to the intricate relationship between the complexities inherent in both the architecture and the brain’s structure. Just as different modalities provide complementary information about anomalous regions in the brain, different architectures may be more adept at handling particular areas of the brain[7].

In this paper, we are focused on using the CBTN-CONNECT-DIPGr-ASNR-MICCAI BraTS-PEDs 2023 Challenge data which constitutes four different modalities (native T1, post-contrast T1-weighted (T1Gd), T2-weighted (T2), and T2 Fluid Attenuated Inversion Recovery (T2-FLAIR)) [8]. We introduce a novel ensemble approach using ONet and various modified versions of UNet along with modified loss functions which yielded a more precise segmentation for the pediatric tumors.

2 Methodology

Figure 1: Comprehensive workflow visualizing the end-to-end process with model training and ensemble approaches.

2.1 Data Augmentation

Due to the inherent variability and noise that can occur during acquisition using different scanners, the integration of various physics-inspired augmentation techniques is essential. By populating the training data with these techniques, the model becomes robust and accurate across diverse scanning protocols.[9]

In this study, we used two different approaches of augmentation: single and composite transformation. For the single transformation, we used techniques like flip, affine transformation, elastic deformation, noise, rescale Intensity, and random bias field. Additionally, in composite transformation, a combination of some or all of the aforementioned techniques was applied to further enrich the dataset.[10, 11, 12, 13]

2.2 Model Architectures

2.2.1 UNet3D Family

In this work, we implemented eight unique variations using two fundamental base architectures, each paired with different loss functions (further explained in Section 2.3) and tailored hyperparameters. Figure 3 illustrates the architecture for one of the variants of the UNet3D configurations.

Figure 2: Core UNet3D Architecture.

The UNet variations are as follows:

• 3D UNet: This is the standard configuration utilizing an input channel size of 4, output class size of 3, and a channel size of 32.

• 3D UNet GELU: A variant of the standard 3D UNet where the ReLU activation function is substituted with GELU (Gaussian Error Linear Unit) to potentially improve non-linear learning capabilities.

• 3D UNet SingleConv: This version modifies the typical double convolution process in each upscaling and downscaling step, replacing it with a single convolution, which might affect the model’s ability to capture complex features.

• 3D UNet Attention: Incorporates an attention mechanism into the 3D UNet structure, aiming to enhance the model’s focus on relevant features for improved segmentation.

• 3D UNet Dropout: Introduces dropout layers into the 3D UNet architecture to prevent overfitting and promote generalization.

2.2.2 ONet3D Family

We also explore variations within the ONet3D family, which differs from the UNet3D by concatenating the encoder-decoder sections before the output convolution layer:

• 3D ONet SingleConv Kernel_1: Utilizes a single convolution with a kernel size of 1 in the ONet3D architecture.

• 3D ONet SingleConv Kernel_5: Adopts a single convolution with a kernel size of 5, offering a broader receptive field compared to Kernel_1.

• 3D ONet DoubleConv Kernel_1: Features double convolution operations with a kernel size of 1, potentially enhancing the model’s ability to extract fine-grained features.

2.3 Loss Functions

In our segmentation framework, we explored different loss functions to optimize the model performance. Specifically, we employed two distinct loss combinations: Binary Cross-Entropy (BCE)[14] plus Dice Loss, and Generalized Dice Loss[15] plus BCE. The choice of these loss functions is motivated by their abilities to tackle the class imbalance problem often encountered in medical image segmentation.

2.3.1 Binary Cross-Entropy combined with Dice Loss

The Binary Cross-Entropy (BCE) loss is commonly used for binary classification tasks and computes the cross-entropy between the true labels and predicted probabilities. To enhance the model’s sensitivity to the region of interest, we combined BCE with Dice Loss. The Dice Loss computes the overlap between the predicted segmentation and the ground truth, providing a spatially aware metric that is sensitive to the shape of the segmented structures.

The BCE loss is defined as:

$$BCE(y,\hat{y})=-\frac{1}{N}\sum_{i=1}^{N}\left(y_{i}\log(\hat{y}_{i})+(1-y_{i})\log(1-\hat{y}_{i})\right)$$ (1)

and the Dice Loss is defined as:

$$Dice(y,\hat{y})=1-\frac{2\sum_{i=1}^{N}y_{i}\hat{y}_{i}+\varepsilon}{\sum_{i=1}^{N}y_{i}+\sum_{i=1}^{N}\hat{y}_{i}+\varepsilon}$$ (2)

where $\varepsilon$ is a small constant to avoid division by zero. The combined loss function is:

$$Loss(y,\hat{y})=\alpha\cdot BCE(y,\hat{y})+\beta\cdot Dice(y,\hat{y})$$ (3)

where $\alpha$ and $\beta$ are weights that balance the contribution of each term.

2.3.2 Generalized Dice Loss combined with BCE

Recognizing the limitations of standard Dice Loss in handling cases with varying object sizes, we also experimented with Generalized Dice Loss (GDL). GDL extends the traditional Dice Loss by incorporating class-wise weights, thus accommodating the imbalanced class distribution. This loss function is especially well-suited to medical image segmentation where certain classes or labels may be underrepresented.

The GDL is defined as:

$$GDL(y,\hat{y})=1-\frac{2\sum_{c=1}^{C}w_{c}\sum_{i=1}^{N}y_{ci}\hat{y}_{ci}+\varepsilon}{\sum_{c=1}^{C}w_{c}(\sum_{i=1}^{N}y_{ci}+\sum_{i=1}^{N}\hat{y}_{ci})+\varepsilon}$$ (4)

where $C$ is the number of classes, $w_{c}$ is the weight for class $c$, defined as:

$$w_{c}=\frac{1}{(\sum_{i=1}^{N}y_{ci})^{2}}$$ (5)

Here, $\varepsilon$ is again a small constant to avoid division by zero.

2.4 Ensemble Strategy

The ensemble strategy integrates the predictive capabilities of ONet and UNet models, capitalizing on their complementary strengths. Initial predictions from the models are combined through a summation of logits, followed by thresholding. A majority vote across the models for each voxel finalizes the ensemble prediction, aiming to capture specific features and model diverse aspects of the MRI images.

2.5 Post Processing

Post-processing techniques are applied to the combined predictions to refine segmentation quality:

• Size Filtering Based on Voxel Volumes: This technique removes small isolated regions, eliminating noise.

• Morphological Reconstruction: A more advanced method for interpolation of voxels violating constraints and smoothing boundaries for all labels.

3 Results

3.1 Evaluation Metrics

The BraTS challenge evaluates submitted models using two primary metrics: the lesion-wise (LW) Dice score and the 95th percentile lesion-wise Hausdorff distance (HD95). These metrics are used to assess segmentations across three distinct tumor sub-regions: the whole tumor (WT), tumor core (TC) , and enhancing tumor (ET).

The Lesion-wise Dice Score and 95th Percentile Lesion-wise Hausdorff Distance (HD95) are key metrics for evaluating segmentation models in medical imaging. The Dice score, ranging from 0 (no overlap) to 1 (perfect overlap), measures the accuracy of predicted segmentations against the true segmentations on a lesion-by-lesion basis, penalizing False Positives and False Negatives by assigning a 0 score. The HD95 metric quantifies the maximum deviation between predicted and actual segmentations for each lesion, with False Positives and False Negatives receiving a fixed penalty value of 374. Mean scores for both metrics are calculated across case IDs. These evaluations offer a detailed insight into the model’s segmentation performance, highlighting areas for potential enhancement

3.2 Quantitative Performance

Table 1: Comparison of evaluation metrics for Enhancing Tumor on validation data

Models	LW Dice $\uparrow$	Dice $\uparrow$	LW HD95 $\downarrow$	HD95 $\downarrow$
UNet3D	0.43	0.44	168.68	114.03
ONet3D	0.52	0.48	131.28	121.74
Ensemble (BCE, Dice)	0.38	0.38	186.73	141.84
Ensemble (BCE, GD)	0.52	0.49	127.11	105.46

Table 2: Comparison of evaluation metrics for Tumor Core on validation data

Models	LW Dice $\uparrow$	Dice $\uparrow$	LW HD95 $\downarrow$	HD95 $\downarrow$
UNet3D	0.70	0.75	44.68	12.19
ONet3D	0.71	0.74	33.26	19.82
Ensemble (BCE, Dice)	0.64	0.73	60.00	17.68
Ensemble (BCE, GD)	0.72	0.74	23.84	11.48

Table 3: Comparison of evaluation metrics for Whole Tumor on validation data

Models	LW Dice $\uparrow$	Dice $\uparrow$	LW HD95 $\downarrow$	HD95 $\downarrow$
UNet3D	0.75	0.82	45.51	11.85
ONet3D	0.76	0.80	32.88	19.31
Ensemble (BCE, Dice)	0.70	0.82	61.74	15.97
Ensemble (BCE, GD)	0.78	0.82	25.88	10.89

In our assessment of various models, we submitted the results to the synapse portal for testing the results on the validation. The findings revealed that the ensemble training approach demonstrates better effectiveness in comparison to the single model training approach, across almost all of the validation cases.

Furthermore, an individual evaluation of predictions indicated that the single-model training approach of the ONet3D model matches the ensemble training approach for lesion-wise dice as outlined in Table 1. This indicates that there can be an ensemble model which might be robust and be generalizable for a larger dataset. However for specific types of segmentation and evaluation metrics, the approaches might differ.

Table 4: Evaluation metrics for sub-regions on Final Test Data

Sub-region	LW Dice $\uparrow$	Dice $\uparrow$	LW HD95 $\downarrow$	HD95 $\downarrow$
ET	0.55	65.23	0.55	0.99
TC	0.71	31.61	0.61	0.99
WT	0.79	22.36	0.70	0.99

In the evaluation of our brain tumor segmentation model on the BraTS test set, the results demonstrate varying degrees of performance across different tumor sub-regions, indicating the model’s proficiency in distinguishing between tumor core (TC), enhancing tumor (ET), and whole tumor (WT) regions. Specifically, the model achieved Dice scores of 0.5519, 0.7054, and 0.7938 for ET, TC, and WT, respectively, suggesting a higher accuracy in segmenting the whole tumor region compared to the more challenging enhancing tumor and tumor core regions. Similarly, the Hausdorff distance (95th percentile) metrics, which assess the model’s precision in outlining tumor boundaries, show values of 65.23, 31.61, and 22.36 for ET, TC, and WT, respectively, indicating the model’s increased boundary precision in WT segmentation. The sensitivity scores—0.5544 for ET, 0.6096 for TC, and 0.7010 for WT—further highlight the model’s ability to correctly identify tumor pixels, with WT regions being most accurately detected. In contrast, the specificity scores, which are consistently high across all regions (0.9997 for ET, 0.9998 for TC, and 0.9999 for WT), underscore the model’s effectiveness in correctly classifying non-tumor pixels. These results collectively underscore the model’s robustness in brain tumor segmentation, with notable strengths in whole tumor delineation and high specificity across tumor sub-regions. In the end it’s evident from the validation and test results that the model generalizes well.

3.3 Visual Comparison

Figure 3 illustrates a side-by-side comparison of prediction results derived from two separate methodologies for modeling: individual model training and the ensemble approach.

The comparison reveals a noticeable difference between the two approaches. The predictions stemming from the single model training method, as evidenced by UNet3D and ONet3D, demonstrate to capture less core tumor region and region affected by the tumor (enhancing tumor). This limitation is contrasted by the predictions generated through the ensemble training method.

Figure 3: A, B, and C depict three slice planes (sagittal, coronal, axial) from the same CaseID, showcasing predictions generated by three different models—UNet3D, ONet3D, and the ensemble model respectively.

The ensemble approach adds the strengths of multiple individual models, potentially leading to more accurate and robust predictions. In the context of the figure, it becomes evident that employing the ensemble method results in more effective coverage of tumor regions. This might suggest that the ensemble training not only compensates for the limitations observed in the single model training but possibly enhances the overall predictive capability.

4 Discussion and Conclusion

The integration of ONet and UNet models through an ensemble technique, together with the novel proposed post-processing strategy, creates an advanced approach to medical image segmentation. This method builds upon the unique strengths of both ONet and UNet models, combining them in an interdependent manner that leverages their individual capabilities. By doing so, the approach not only amplifies the robustness of the segmentation but also adds a level of precision that might be unattainable with a single-model training strategy. The clinical significance of this methodology lies in its potential to offer enhanced diagnostic accuracy and effective treatment planning.