BRAIN-CATS: Brain Tumour Reliability-Aware Imaging with Neural Networks Using Calibration-Aware Training and Segmentation
摘要
Accurate and reliable brain tumour segmentation from MRI remains a clinical challenge, particularly in low-resource settings such as Sub-Saharan Africa (SSA). We present BRAIN-CATS, a segmentation framework that combines the Attention U-Net architecture with calibration-aware training to improve both accuracy and model reliability. Our approach is specifically optimized for low-resolution, multi-modal MRI data typical in under-resourced environments. We trained the model using 5-fold cross-validation on n = 60 patients from the BraTS-Africa 2023 dataset, employing advanced preprocessing techniques, data augmentation, and a composite loss function that includes Dice, Binary Cross-Entropy, Focal Loss, and the marginal L1 Average Calibration Error (mL1-ACE). The calibration-aware component penalizes miscalibrated predictions, improving confidence estimates across tumour boundaries. The model achieved an average Dice score of 90.38% for edema (ED), 81.94% for enhancing tumour (ET) and 78.88% for the tumour core (TC) across the folds. On external validation using 35 held-out cases, mean Dice scores were 65.30% (ET), 66.50% (TC) and 50.80% for whole tumour (WT), confirming the model’s ability to generalize to unseen data despite limited training resources. Early stopping between epochs 36–39 prevented overfitting, and stochastic weight averaging further improved generalization. Ensemble inference using all fold checkpoints yielded smooth, consistent predictions on validation data. Visual and quantitative evaluations confirm BRAIN-CATS as a practical and resource-efficient solution for glioma segmentation in challenging imaging contexts.