QCResUNet: Joint Subject-level and Voxel-level Segmentation Quality Prediction

In this study, we used three datasets for evaluation of the QC performance on brain tumor segmentation. First, we used pre-operative multimodal MRI scans with gliomas of all grades (WHO Central Nervous System grades 2-4) grades from the BraTS 2021 challenge training dataset ($n=1251$). The BraTS dataset is a heterogeneous dataset consisting of cases from 23 different sites with various levels of quality and protocols. The BraTS dataset was used for training, validation, and internal testing. Additionally, we used two datasets that are not included in the BraTS 2021 dataset (i.e., the BraTS-SSA dataset and the WUSM dataset) for external testing, allowing for an unbiased assessment of the generalizability of the proposed method. The BraTS-SSA dataset ($n=40$) (Adewole et al., 2023) is an extension of the original BraTS 2021 dataset with patients from Sub-Saharan Africa, which includes lower-quality MRI scans (e.g., poor image contrast and resolution) as well as unique characteristics of gliomas (i.e., suspected higher rates of gliomatosis cerebri). The WUSM dataset ($n=175$) was obtained from the retrospective health records of the Washington University School of Medicine (WUSM), with a waiver of consent, in accordance with the Health Insurance Portability and Accountability Act, as approved by the Institutional Review Board (IRB) of WUSM (IRB no. PA18-1113). Each subject in all datasets comprised four modalities viz. pre-contrast T1-weighted (T1), T2-weighted (T2), post-contrast T1-contrast (T1c), and Fluid attenuated inversion recovery (FLAIR). In addition, multi-class tumor segmentation masks annotated by experts were also available. Segmentation masks delineated enhancing tumor (ET), non-enhancing tumor core (NCR), and edema (ED) classes. Following standard BraTS procedures, we combined the binary ET, NCR and ED segmentation masks to delineate the whole tumor (WT), tumor core (TC), and enhancing tumor. The WT mask consists of all tumor tissue classes (i.e., ET, NCR, and ED), while the TC mask comprises ET and NCR tissue classes.

Scans from the BraTS training and BraTS-SSA datasets were already registered to the SRI24 anatomical atlas (Rohlfing et al., 2010), resampled to 1-mm³ isotropic resolution and skull-stripped. For consistency, raw MRI scans from WUSM were pre-processed following the same protocol using the Integrative Imaging Informatics for Cancer Research: Workflow Automation for Neuro-oncology (Chakrabarty et al., 2022) framework. Subsequently, we z-scored all the skull-stripped scans in the BraTS datasets and the WUSM dataset on a per-scan basis. Finally, scans from the entire dataset were cropped to exclude background regions and then were zero-padded to a common dimension of $160\times 192\times 160$ using the nnUNet preprocessing pipeline (Isensee et al., 2018b).

In evaluating the QC performance on the cardiac segmentation task, we used the Automated Cardiac Diagnosis Challenge (ACDC) dataset (Bernard et al., 2018), which consists of 100 subjects (200 volumes). Each volume in the ACDC dataset is associated with a multi-class segmentation mask delineating the left ventricle (LV), myocardium (Myo), and right ventricle (RV). Lastly, each volume was cropped and zero-padded to a common dimension of $16\times 16\times 160$ using the nnUNet pipeline (Isensee et al., 2018b).

The majority of previous regression-based methods were trained and validated on datasets limited by both sample size and heterogeneity. To address this limitation, we constructed an extensive dataset consisting of segmentation results of diverse quality levels. This enabled us to capture the variability of segmentations and obtain a reliable performance assessment of the proposed method on a wide range of segmentations. For this purpose, we adopted a four-step approach to produce segmentation results using different methods with various combinations of imaging modalities as input.

As far as the brain tumor segmentation QC task is concerned, we first employed both CNN-based (i.e., nnUNet (Isensee et al., 2018b)) and Transformer-based (i.e., nnFormer (Zhou et al., 2021)) segmentation UNets and trained them independently seven times by taking different modalities as input (namely T1-only, T1c-only, T2-only, FLAIR-only, T1-FLAIR, T1c-T2, and all four modalities). Since each modality is sensitive to certain tissue types (e.g., T1c is effective in detecting enhancing tumors and FLAIR is advantageous in identifying edema), this approach enabled us to produce segmentations with varying levels of quality. Second, we sampled the segmentation along the training routines at various iterations to generate segmentation results from both fully trained and inadequately trained models. For this purpose, a small learning rate ($1~{}\times~{}10^{-6}$) was used during training to slower their convergence, allowing the segmentation to improve gradually from low quality to high quality. Third, we employed a data augmentation strategy based on image transformations to produce diverse segmentation results, termed SegGen. A series of image transformations, including rotation, translation, scaling, and elastic deformation, were randomly applied to each ground-truth segmentation three times with a probability of $0.5$ (see Table 1). Fourth, apart from the segmentations produced by the aforementioned three methods, we also created out-of-sample segmentation samples from a different segmentation method to assess the generalizability of our proposed model in the testing phase. To accomplish this, we employed the DeepMedic framework (Kamnitsas et al., 2017) to train seven segmentation models using the same diverse input modalities as we did in training the nnUNet and nnFormer. After applying the aforementioned four steps, we were able to create a set of erroneous segmentations for the cases included in the internal and external datasets. The erroneous segmentations along with the ground truth segmentations were subsequently used to estimate the DSC, NSD, and the ground truth segmentation error map (SEM) that were used to train the proposed network. The SEM was computed as the difference between a query segmentation (i.e., one of the erroneous segmentations we produced) and the corresponding ground truth. Examples of the generated segmentations can be found in Figure 1(b).

For the cardiac segmentation task, we employed a procedure similar to the one used in generating segmentation for the brain tumor task. However, unlike brain tumor MRI, cardiac MRI involves only a single modality. Consequently, we trained one nnUNet model and one nnFormer model, sampling the segmentation along the training routines at various iterations. We also generated segmentation results using SegGen during model training. However, due to the technical difficulty of training DeepMedic on the ACDC dataset (i.e., one axis in the ACDC dataset has a significantly smaller input size), we could not include any segmentation produced by DeepMedic for the cardiac segmentation QC task.

Despite our efforts to sample across different convergence stages by adopting a small learning rate, we observed that the deep learning models could successfully segment the majority of the cases. As a consequence, the resulting dataset was distributed unevenly across the different levels of quality, comprising mostly of higher quality segmentation results (see Figure 1(a)). Following the resampling strategy in Robinson et al. (2018), we randomly selected a subset of samples from each bin of the DSC histogram. This ensured that the number of segmentations from each bin was equal to the lowest count-per-bin value ($n_{s}$) across the distribution for all the segmentation datasets (refer to Figure 1(a)). For this purpose, we divided DSC values into 10 evenly spanned bins ranging from 0 to 1. We sampled cases from these bins such that we ensured that all bins contained an equal number of cases. For the training set, we did not manually select the DSC samples. Instead, at each training iteration, we randomly sampled segmentations that were evenly distributed across 10 bins to avoid imbalanced DSC values in each training batch. We kept this process stochastic to let the model see as many samples as possible without skewing its exposure to either low DSC or high DSC samples. This means we randomly sampled $n_{s}$ segmentations for bins that contained more than $n_{s}$ samples at each training iteration. Conversely, during the evaluation phase, we opted for a deterministic resampling approach, which enabled unbiased evaluations across various quality levels. In order to account for the variability introduced by the resampling process, we implemented a three-fold cross-validation strategy for all our experiments (see Section 3).

The proposed 3D U-shaped QCResUNet (Figure 2(a)) aims to automatically evaluate the quality of a query multi-class segmentation mask by predicting subject-level segmentation quality measures as well as identifying voxel-level segmentation error maps for each tissue class (SEMs). Without loss of generality, QCResUNet takes as input $m$ imaging modalities represented by $\{M_{1},M_{2},\cdots,M_{m}\}$ (e.g., $m=4$ the brain tumor segmentation QC case) and the query multi-class segmentation mask ($S_{query}$). The outputs of QCResUNet include two complementary subject-level segmentation quality measures (DSC and NSD) and a collection of binary voxel-level SEMs, each delineating the segmentation error corresponding to a specific tissue class in $S_{query}$. We would like to point out that we followed the standard definition of DSC for calculating the subject-level DSC for multi-class segmentation masks (see A).

We performed three-fold cross-validation to validate the effectiveness of the proposed method for the brain tumor segmentation QC task. For this purpose, we held out an internal testing set with 251 subjects. In each iteration of the three-fold cross-validation, the remaining BraTS dataset was randomly partitioned into training and validation subsets with 667 and 333 subjects, respectively. The application of the previously described four-step approach and resampling (see Section 2.2) resulted in 98,048 training segmentation samples, 28,800 validation segmentation samples, as well as 39,063 internal testing segmentations (36,144 from nnUNet and nnFormer, 2,919 from DeepMedic). We only used the segmentations generated by the DeepMedic as the testing set. This is because segmentations generated by DeepMedic were not seen during training and validation. The BraTS-SSA and WUSM datasets were used as external testing datasets to validate the generalizability of the proposed method. By employing the previously described four-step approach (see Section 2.2), we extended the original BraTS-SSA, and WUSM datasets to include 6,040 and 26,425 segmentations, respectively. After applying the resampling in Section 2.2, the internal BraTS testing, BraTS-SSA, and WUSM datasets were resampled to include 4,753, 1,204, and 2,693 segmentations, respectively.

We performed hyperparameter tuning to determine two critical hyperparameters in the proposed method: the initial learning rate and the loss balance parameter $\lambda$. Hyperparameter tuning was carried out by performing a Bandit-based Bayesian optimization on the training and validation dataset implemented using the Raytune framework (Liaw et al., 2018). Specifically, we sampled the learning rate from a range of $[1\times 10^{-5},2\times 10^{-4}]$ following a log-uniform distribution. Additionally, the loss balance parameter $\lambda$ was chosen from the following set of values $\{0.1,1,2\}$. The optimal balance parameter $\lambda$ for QCResUNet was determined to be 1. However, we did not observe a significant difference when varying $\lambda$ from 0.1 to 1 and 2 (see Figure S1). Although the optimal learning rate varied from model to model (refer to Table 2), it generally fell within the scale of $1\times 10^{-4}$ as we used a relatively small batch size of 4. We kindly direct the readers to C for the detailed learning curve in hyperparameter tuning using RayTune.

Results reported hereafter were obtained using the above optimal hyper-parameters. For both internal (BraTS) and external (BraTS-SSA and WUSM) testing sets, we reported the results after performing a model ensemble technique for the proposed method and all baseline methods. This involved averaging the results from three models for DSC and NSD prediction and performing majority voting for the SEM prediction. The mean ($\pm$ standard deviation) of MAEs and DSC${}_{\text{SEM}}$ were calculated over all subjects.

In the cardiac MRI segmentation QC task, we use the ACDC dataset for training, validation, and testing. The ACDC dataset was partitioned into training, validation, and testing subsets with a ratio of 7:1:2 by following the partitioning rule outlined in (Zhou et al., 2021; Chen et al., 2021; Cao et al., 2022) to avoid overlapping of subjects. This results in a training set consisting of 140 volumes, a validation set with 20 volumes, and a testing set containing 40 volumes. By employing the same four-step approach as in the brain tumor QC task, the original ACDC training, validation, and testing subsets were extended to include 4,900, 640, and 1,280 segmentations. The testing set was then resampled to include 1,011 segmentations. We used the same hyperparameters from the brain tumor segmentation QC task to train the model for the cardiac segmentation QC task. Similar to the brain tumor segmentation QC task, we also reported the mean ($\pm$ standard deviation) of MAEs and DSC${}_{\text{SEM}}$ over all subjects for the cardiac segmentation QC task.

To evaluate the effectiveness of the proposed model, we compared its performance with five baseline models: (i) RCA method (Robinson et al., 2017); (ii) uncertainty estimation (UE) based method (Jungo et al., 2020); and three regression-based methods: (iii) a UNet model (Ronneberger et al., 2015); (iv) a ResNet-34 model (He et al., 2016); and (v) a ResNet-50 (He et al., 2016) model (Robinson et al., 2018).

The selection of the three regression-based methods (UNet, ResNet-34, ResNet-50) as baselines was based on their architectural resemblance to the proposed approach. To ensure a fair comparison, the residual blocks in ResNet-34 and ResNet-50 were identical to those used in the proposed QCResUNet. As the UNet model typically outputs a segmentation mask of the same dimensions as the input, we added an average pooling layer after its final feature map followed by a fully connected layer to enable the prediction of a single DSC value. The RCA and UE-based methods are considered state-of-the-art in the segmentation QC literature.

RCA for brain tumor segmentation QC was implemented using as segmentation method the DeepMedic framework, which was adapted following the protocol described by Valindria et al. (2017). Specifically, we reduced the number of feature maps in each layer by one third compared to the default setting of DeepMedic. We also reduced the feature maps in the last fully connected layer from 150 to 45. We opted for this approach instead of the atlas-based segmentation approach because establishing spatial correspondences between query and reference images in tumorous areas can be challenging due to the high spatial and phenotypic heterogeneity of brain tumors. In contrast, RCA for cardiac MRI segmentation QC was implemented using the atlas-based segmentation method by following the protocols in (Robinson et al., 2019). This is because registration correspondences between query and reference are easier to establish in healthy autonomy. For the selection of the reference data set, we randomly selected 100 samples from the training dataset following a previously well-validated protocol (Robinson et al., 2019).

The implementation of the UE-based method followed the protocol in (Jungo et al., 2020), where the uncertainty map was calibrated by evaluating the uncertainty-error (UE) overlap. The UE overlap measures the overlap between the binarized uncertainty map and the corresponding SEM (see B.1 for details). We chose the UE calibration proposed by Jungo et al. (2020) instead of other calibration methods (Wen et al., 2020; Ashukha et al., 2020; Mehta et al., 2022) because it produces a binary mask by thresholding the uncertainty map. We can directly compare the resulting binary mask to the SEM generated by the proposed QCResUNet. The subject-level DSC was computed by applying a random forest with 102 extracted radiomics features (see B.2).

We evaluated the performance of QC segmentation at both the subject level and the voxel level. At the subject level, we evaluated the precision of the prediction of the segmentation quality metric based on the MAE and the Pearson correlation coefficient $r$. MAE measures the amount of error in the prediction. $r$ measures the linear correlation between the predicted DSC/NSD value (DSC_pred/NSD_pred) and the ground truth DSC/NSD (DSC_gt/NSD_gt). It ranges from $-1$ (strongly uncorrelated) to $+1$ (strongly correlated).

At the voxel level, the performance was evaluated based on the DSC (DSC${}_{\text{SEM}}^{(c)}$) between the predicted SEM (SEM${}_{pred}^{(c)}$) and ground-truth SEM (SEM${}_{gt}^{(c)}$) for each tissue class $c$. Here, we report the average DSC across all $C$ tissue classes. We performed a paired t-test to compare the DSC predicted by the QCResUNet with those predicted by the corresponding baselines. The resulting p-value ($P$) was reported as a measure to determine if there was a statistically significant difference between the performance of the proposed method and the baseline methods. The significance level was set a priori to $P<0.05$.

Training procedures were the same for both tasks by using segmentation results were produced by nnUNet, nnFormer, and SegGen. Training was carried out using Adam optimizer (Kingma and Ba, 2014) with a batch size of 4 for 100 epochs. Training was started with an optimal initial learning rate determined by the hyperparameter search and was exponentially decayed by a factor of 0.9 for each epoch until it reached $1\times 10^{-6}$. We applied an $L_{2}$ weight decay of $1\times 10^{-4}$ in all trainable layers. To prevent overfitting, various data augmentations, including random rotation, scaling, mirroring, Gaussian noise, and gamma intensity correction, were applied during training. All model training was performed on NVIDIA Tesla A100 GPUs. By default, we used a mix of single-precision (FP32) and half-precision (FP16) tensors during training to reduce training time and memory consumption. The proposed method was implemented using PyTorch v1.12.1. All the code and pre-trained models developed for this work will be available at https://github.com/sotiraslab/QCResUNet.

The proposed QCResUNet achieved good voxel-level segmentation error localization in terms of DSC${}_{\text{SEM}}$ for both brain tumor and cardiac segmentation QC tasks (see Table 5). Specifically, QCResUNet achieved an average DSC${}_{\text{SEM}}$ of 0.769 on the BraTS internal testing set, 0.702 on the BraTS-SSA dataset, and 0.684 on the WUSM dataset in the brain tumor segmentation QC task. Despite a slight performance drop on external datasets, the results remain high quality. Inter-rater agreement studies in glioma segmentation (Visser et al., 2019) indicate that an overlap measure above 0.7 signifies a good segmentation result. Our out-of-distribution DSC${}_{\text{SEM}}$ results were on average around 0.7, despite the increased difficulty of the task. In addition, QCResUNet achieved an average DSC${}_{\text{SEM}}$ of 0.703 in the cardiac segmentation QC task. The detailed distribution of tissue-level DSC${}_{\text{SEM}}$ (i.e., [DSC${}_{\text{SEM}}^{\text{WT}}$, DSC${}_{\text{SEM}}^{\text{TC}}$, DSC${}_{\text{SEM}}^{\text{ET}}$] in brain tumor segmentation and [DSC${}_{\text{SEM}}^{\text{LV}}$, DSC${}_{\text{SEM}}^{\text{Myo}}$, DSC${}_{\text{SEM}}^{\text{RV}}$] in cardiac segmentation) is shown in Figure 5.

Compared to the UE-based QC method, QCResUNet improved error localization significantly by an average of 127.5% in terms of average DSC${}_{\text{SEM}}$ on the brain tumor segmentation QC task (see Table 5). Similarly, QCResUNet significantly outperformed the UE-based QC method by 52.8% in the cardiac segmentation QC task for segmentation error localization (Table 5).

We performed ablation studies to validate the effectiveness of the proposed multi-task learning strategy and attention-based SEM aggregation mechanism. Although sharing the same network structure for subject-level segmentation quality prediction, QCResUNet outperformed ResNet-34 in both brain tumor and cardiac segmentation tasks (Table 3 and Table 4). In addition, we found that the performance of voxel-level segmentation error prediction deteriorated when removing the subject-level prediction task (see Table 5). This demonstrates the effectiveness of the proposed multi-task learning framework for segmentation QC tasks.

In the voxel-level segmentation error localization task, incorporating the proposed attention-based SEM aggregation mechanism improved performance by an average of 13.9% for the brain tumor segmentation QC task and 14.7% for the cardiac segmentation QC task. We conjectured the improvement introduced by the proposed attention-based SEM aggregation could be attributed to the fact that the features in the last layer of the decoder contribute differently to the prediction of each tissue class error segmentation mask.

We hypothesized that the performance gain of the proposed QCResUnet compared to the baselines of UNet, ResNet-34, and ResNet-50 is due to joint training under both subject-level and voxel-level supervision. One potential explanation could be that the joint training allowed the proposed QCResUNet to effectively localize on segmentation errors, as suggested by the attention maps produced by Gradient-weighted Class Activation Mapping (Grad-CAM) (Selvaraju et al., 2017) (see Figure 7). We observed that the CAM obtained from ResNet-34 and ResNet-50 did not focus on the areas of segmentation errors. Although ResNet-50 had more parameters and depth, it did not demonstrate significant improvement over ResNet-34, indicating that these factors did not improve the QC performance. In contrast, the UNet CAM showed a coarse localization on the difference between the ground truth and the predicted segmentation. In contrast to the baseline methods, the QCResUNet CAM was well localized on segmentation errors.