{'fraction_non_alphanumeric': 0.029298470876619586, 'fraction_numerical': 0.009746702462939185, 'mean_word_length': 4.814387512724805, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Collaborative Boundary-aware Context Encoding Networks for Error Map Prediction

Medical image segmentation is usually regarded as one of the most important intermediate steps in clinical situations and medical imaging research. Thus, accurately assessing the segmentation quality of the automatically generated predictions is essential for guaranteeing the reliability of the results of the computer-assisted diagnosis (CAD). Many researchers apply neural networks to train segmentation quality regression models to estimate the segmentation quality of a new data cohort without labeled ground truth. Recently, a novel idea is proposed that transforming the segmentation quality assessment (SQA) problem intothe pixel-wise error map prediction task in the form of segmentation. However, the simple application of vanilla segmentation structures in medical image fails to detect some small and thin error regions of the auto-generated masks with complex anatomical structures. In this paper, we propose collaborative boundaryaware context encoding networks called AEP-Net for error prediction task. Specifically, we propose a collaborative feature transformation branch for better feature fusion between images and masks, and precise localization of error regions. Further, we propose a context encoding module to utilize the global predictor from the error map to enhance the feature representation and regularize the networks. We perform experiments on IBSR v2.0 dataset and ACDC dataset. The AEP-Net achieves an average DSC of 0.8358, 0.8164 for error prediction task,and shows a high Pearson correlation coefficient of 0.9873 between the actual segmentation accuracy and the predicted accuracy inferred from the predicted error map on IBSR v2.0 dataset, which verifies the efficacy of our AEP-Net.


Introduction
Image segmentation plays a fundamental role in medical imaging research and clinical situations such as radiotherapy and image-guided interventions. In recent years, deep learning methods have achieved improvements in various medical image segmentation tasks. However, it is challenging to obtain promising segmentation results when lacking enough manual annotations. Thus, it is of vital importance to employ an accurate segmentation quality assessment of per case for reducing errors in subsequent analysis or downstream procedures, and it is also helpful to the clinicians. However, it's unrealistic to get all manually labeled ground truths (GTs) of testing scans. Therefore, many efforts have been taken to SQA methods without ground truth. The reverse classification accuracy (RCA) method [1] trained a reverse classifier with the generated mask, and utilized the segmentation performance on a set of reference images as the proxy assessment. The main drawback of this method is the high cost of computational time. Some methods [2,3] trained the regression networks to predict non-reference quality metrics. For example, [2] regressed the similarity between a degraded vessel tree and a manual segmentation. [3] directly regressed the Dice Similarity Coefficient (DSC) of the generated masks. These methods need a large number of training samples, and they are vulnerable to the adversarial attacks. The uncertainty measures between multiple candidate segmentation masks have become another popular method to assess segmentation quality automatically, such as variation or DSC agreements between Monte Carlo samples in [4]. However, they usually need multiple forward propagations of models to get candidate masks, increasing the computational overhead and reducing efficiency.

Error Map Prediction
Recently, Zhang R, et al. [5] proposed an error map prediction network to predict the misclassified pixels in the generated masks, which transferred the quality assessment task into a segmentation task. The high potential of this method is that it can provide clinicians with a visual guide by pointing out the mislabeled pixels, and the quality metric also can be inferred from the predicted error map. However, this task is different from the general segmentation frameworks. It cannot achieve satisfactory results on the dataset which has complex anatomical structures when just using the vanilla structure, such as 3D U-Net [9], VoxResNet [10], and directly concatenating the original images with masks as network input. Specifically, there is a big semantic gap [6] between the original images and the segmentation masks. The low-level features from the original images are too noisy to provide more useful information. Secondly, the mislabeled pixels are prone to appear on the semantic boundary, the vanilla structure always fails to highlight the information of semantic boundary. Last but not least, the error map itself persists the global contextual supervision information representing segmentation quality. Understanding and utilizing this information is essential for this task.
Contribution: To address these problems, we propose an accurate error prediction framework called AEP-Net, which consists of three parts: the main error-prediction (MEP) branch, the collaborative boundary-aware feature transformation (CBFT) branch, the context encoding unit (CEU) [7]. Specifically, the CBFT integrates class boundary information to enhance the feature fusion and provide guidance for error detection. In addition, the CEU imports the error map predicator as the global information constraint to boost the feature representation of the networks. Under the cooperation of these three parts, we achieve higher average DSCs for error map segmentation on IBSR v2.0 and ACDC datasets, and observe a strong positive correlation between the predicted segmentation accuracy and the actual segmentation accuracy. Fig. 2. The proposed framework, which consists of the MEP branch, the CBFT branch and the CEU module. Fig. 2 illustrates the architecture of our AEP-Net, which is based on the MEP branch, the CBFT branch and the CEU. We leverage the idea of feature sharing [8] and the attention mechanism to build the collaborative relationships between the MEP and the CBFT. Maintaining the class boundary in the CBFT can capture more useful features near the class boundaries, which are also error-prone regions. Thus, this collaborative branch can provide more valuable transformed features and fine-grained attention for the MEP branch. The MEP takes the one-hot masks as input, and it obtains the upgraded features from the CBFT branch during the forward process. Both two branches are based on encoder-decoder structures. Additionally, the global quality information derived from the actual error map is utilized to regularize the networks through the CEU.

Main Error Prediction Branch
We employ a u-net structure to predict error map. The two-class Generalized Dice Loss [12] is used as the error map segmentation loss, since it performs better for class imbalance problem. Let be the reference error map (one-hot ground truth) with values ( ) for class c, and be the predicted two-channel error map, ( ) be the predicted probability of voxel for class . It can be formulated as: For better feature fusion of the two branches at different semantic levels, we employ the different feature sharing methods. Specifically, at the encoder stage, the relatively low-level local features from the CBFT branch are concatenated with the features from the MEP branch. While at the decoder stage, we boost the feature fusion by introducing the boundary-aware attention module. Let the and the denote the features from the MEP branch and the CBFT branch respectively, the 1×1×1 denote the 1×1×1 convolution, the process of attention module is described in Eq.2.

Boundary-aware Feature Transformation Branch
As stated in section 1.1, there is a giant semantic gap between the original data and the segmentation result. The direct fusion of high-level segmentation result with the pure low-level original data helps little. Thus, we propose a soft-boundary detection network to transform the features of the original data gradually. We also employ a u-net which keeps pace with the MEP branch. The enhanced boundary, which reflects the semantic class changes between different tissues and reveals the error-prone region, is generated by applying a 3D Sobel layer to the generated mask and a simple boundary enhancement method. Assume that the total number of voxels is , let denote the original data, (·) represent the network parameters of the CBFT branch, denote the Sobel gradient of . In all positions where the gradient is greater than 0, we add a constant of

Context Encoding Unit
Conv+GN+ReLU Pooling the max gradient of , which ensures the values of the predicted boundary are all greater than 0.5 after the normalization. Thus, the enhanced boundary can be formulated as: The MSE loss is used in the CBFT module. Thus, the 1 can be defined as: This branch has two major benefits. Firstly, it guarantees that the shared features persist progressively enhanced semantic information. Secondly, it promotes the MEP to focus on the error-prone semantic boundary by providing attention maps as shown in Fig. 1 (e).

Context Encoding Unit
In order to exploit the global contextual information from the real error map, we define the real error rate ( ), which is inferred from the real error map, as the proportion of the misclassified voxels to the total voxels. Then we employ a Context Encoding Unit (CEU) to enhance the feature representation of the MEP by regressing the . The CEU is attached to the end encoder layer of the MEP. The CEU begins with 3×3×3 convolution layers with group normalization and ReLU, and reduces the spatial dimension by the 3D max-pooling, and the global average pooling operations, ends with two fully connected layers and a sigmoid activation function. We define the error rate predicted by the CEU as .The loss of this module can be defined as: Thus, the total loss of our proposed AEP-Net is defined as: = α · 1 + β · 2 + γ ·

Datasets and Mask Generation
The AEP-Net is evaluated on two datasets. One is the brain anatomical segmentation dataset IBSR v2.0 (https://www.nitrc.org/projects/ibsr), which has 18 T1-weighted MR images of 256 ×128×256 voxels, and manually labeled segmentation for three tissues, white matter, grey matter, cerebrospinal fluid. Another one is the training set of the Automated Cardiac Diagnosis Challenge (ACDC) dataset [11], which consists of 100 4D short-axis cine-CMR scans, and has manual reference images of the right ventricular cavity (RV), the myocardium (Myo), and the left ventricular cavity (LV) in ED (enddiastole) and ES (end-systole) phases, respectively. The in-plane dimension of all scans has a wide range, from 154×224 to 428×512.
The candidate segmentation masks are generated based on 3D U-Net and VoxResNet, two classical segmentation models for volumetric medical images. In IBSR v2.0 dataset, we implement three 3D U-Nets with different number of initial feature maps: 16, 24, and 32, and other operations are same. Then we collect segmentation masks from three training epochs. For VoxResNet, one side output by the deep supervision loss and the final output are generated as two segmentation masks. Further, we also collect masks from two training epochs. Thus, each scan has 13 generated masks. In ACDC dataset, we collect four groups of 800 masks from two 3D U-Nets and VoxResNet. The weighted Dice score histograms of two datasets are shown in supplementary materials. Then we produce the actual error maps by comparing the generated masks with the manual ground truths.

Implementation details
The evaluation is implemented with three-fold cross validation on both two datasets. The generated masks and the real error maps of the training sets are further used to train the AEP-Net. For mask generation, we utilize the cross-entropy loss and multi-class Dice loss to train the segmentation networks in a standard manner. For the error prediction task, we utilize the proposed to train the AEP-Net in an end-to-end way. For data processing and augmentation, we apply the z-score standardization and normalize the voxel intensity to [0,1] on all scans. During the training, we use a random crop size 128×32×128 for brain images, and a size 200×200×8 for cardiac images. Since the dimensions of some scans in ACDC are smaller than this crop size, we utilize zero-padding to make these scans meet the minimum input requirements. Further, we apply the random mirror flipping (along , axes for brain and , axes for cardiac). We implemented our networks on PyTorch and trained them on a single NVIDIA 1080 Ti GPU with 11GB RAM. The Adam optimizer with the poly learning rate policy was used in our experiments, where the initial learning rate is 1 −3 , the poly power is 0.9. It took about 36 hours to train the AEP-Net on both two datasets. The time cost of predicting error map by forward propagation is 14 seconds for one brain MRI scan, and 3.2 seconds for one cardiac scan on average.

Evaluation Metrics
To evaluate the efficacy of the proposed AEP-Net, we explain the indicators we use as follows. The . and . means the dice similarity coefficient and accuracy between the generated masks and the ground truths, while the , , ( ), and ( ) are calculated between the real error map and the predicted error map.

Results and analysis
To validate the performance of the proposed AEP-Net, we show the error prediction results of the different kinds of auto-generated masks in IBSR v2.0 dataset in Table 1.
We only show the overall error prediction performance of the ACDC dataset due to the space. In IBSR v2.0 dataset, the masks are divided into five groups in terms of which range the . belongs to. We also list the number of masks in the second column, and the last two rows show the overall performance on all masks. It is noticed that the larger error regions in the low-quality masks are easier to be found, while the small, thin error regions near the class boundaries in the good quality masks are harder to be detected in the absence of ground truths [5]. The proposed AEP-Net performs well in both five groups. Specifically, it can be observed in the fifth row that even the masks with good quality, our proposed AEP-Net still performs well for detecting the small, thin error regions. We achieve a DSC of 0.8358 on the IBSR v2.0 dataset and a DSC of 0.8164 on the ACDC dataset. More experimental results of the ACDC dataset will be shown in supplementary materials.  Ablation study. In order to prove the effectiveness of the network design. We compare the AEP-Net with a traditional 3D U-Net structure, which has the similar number of parameters as the AEP-Net, and it takes the concatenation of the original images and the segmentation masks as input. We also remove the CEU from the AEP-Net for ablation experiments. All comparative experiments use the same data division. As shown in Table 2, the higher average produced by the AEP-Net illustrates the framework design is more suitable for the error prediction task.

Segmentation Quality Assessment.
To evaluate the effectiveness of the segmentation quality assessment inferred from the predicted error map, we define the as the percentage of the correctly classified voxels to the total voxels inferred from the predicted error map, which denotes the predicted segmentation accuracy. Based on the error prediction results of the IBSR v2.0 dataset, we calculate the Pearson correlation coefficient ( ) and the mean absolute error ( ) between the . and the , and the between the . and the . We also draw scatter plots of these two groups of metrics. The high and , low shows the segmentation quality can be accurately assessed according to the predicted error map.

Conclusion
To address the need for fine-grain segmentation quality assessment of per case in image-based automated medical diagnosis system or medical research, we propose a collaborative boundary-aware context encoding framework for accurate error prediction in medical image segmentation. The proposed framework shows good performance on the auto-generated masks of different quality, especially for detecting the small and thin error regions, which ensures that the quality measurement derived from the predicted error map reflects the segmentation quality accurately and has a high reference value.
In the future, we will study multi-class error map prediction and how to improve the segmentation accuracy by using error map prediction adaptatively. As shown in Fig. 2, we randomly selected 400 masks from ACDC dataset to draw scatter plots. The . sometimes fails to reflect the overall accuracy of per case due to the small target region and the wide range of the in-plane dimension. However, it can be observed from (b) and (c) that the distribution of two scatters is very similar, which illustrates the segmentation accuracy inferred from the error map still performs well. Additionally, we calculate the Pearson correlation coefficient ( ) and the mean absolute error ( ) between the . and , and the between the . and . The high and low also verify the efficacy of our proposed AEP-Net.