SGEResU-Net for brain tumor segmentation

1.   Introduction

Brain tumors are abnormal cells that grow in the human brain, and they can be divided into benign and malignant tumors ^[1]. Nowadays, the incidence of malignant tumors is relatively high, posing a considerable threat to human health. Glioma is the most common malignant tumor, and it mainly includes high-grade glioma (HGG) and low-grade glioma (LGG). Magnetic resonance imaging (MRI) can produce high-quality brain images without damage and skull artifacts, which provides comprehensive information for the diagnosis and treatment of brain tumors ^[2]. In addition, accurate brain tumor segmentation can not only provide valuable information such as morphology, size, and location of the tumor but also help clinicians improve the diagnosis of brain tumors. Therefore, the automatic segmentation of brain tumors on MRI is of great significance to patient treatment and has been widely studied for a long time.

Recently, deep learning methods have demonstrated superiority in the brain tumor segmentation task compared to the traditional methods. Especially, the Multimodal Brain Tumor Segmentation Challenge (BraTS), jointly held with the International Association for Medical Image Computing and Computer Assisted Intervention since 2012, has greatly promoted the progress of deep learning-related brain tumor segmentation. The earlier deep learning methods for brain tumor segmentation mainly utilize the idea of patch classification in small-scale images based on convolutional neural networks. However, these methods have limitations in maintaining image spatial continuity, which also occupies large storage space and leads to low efficiency. Motivated by the fully convolutional network, Ronneberger et al. ^[3] present a novel symmetric 2D U-Net model which contains down-sampling layers for feature encoding, up-sampling layers for feature recovery, and skip connections to integrate up-sampling and down-sampling features. Due to its concise and efficient architecture, U-Net largely improves the segmentation performance for medical images and quickly becomes the mainstream in brain tumor segmentation tasks. Then, to discover the potential of the model, researchers combine U-Net with various other advanced methods, such as residual module ^[4], attention module ^[5] and ensemble module ^[6], to produce more robust brain tumor segmentation models. Though 2D U-Net models require low memory consumption, they have limitations in containing spatial context information of 3D MRI images. Therefore, 3D U-Net models are developed to capture more spatial context information, which also achieves higher accuracy on this medical task. Similar to 2D U-Net models, a great deal of 3D modules, including residual module, dense connection block, and attention module ^[7,8,9], are integrated to gain higher segmentation performance. Among them, Myronenko ^[10] utilizes an asymmetrical residual U-Net, in which most of the trainable parameters of the model resided in the encoder and win the first place in the BraTS 2018 challenge. Zhao et al. ^[11] utilize an 3D U-Net model with dense blocks along with three tricks including data processing methods, semi-supervised learning and optimizing processes, win the second place in the BraTS 2019 challenge. Besides, Jiang et al. ^[12] achieve the first place using the two-stage cascaded asymmetric 3D residual U-Net. Moreover, Isensee et al. ^[13] place the first in the BraTS 2020 challenge based on the 3D nnU-Net ^[14] model. The successful application of the above method is enough to prove that the transformation based on the 3D U-Net network has a huge advantage in brain tumor segmentation.

In this work, following merits of the 3D U-Net architecture, we try to construct a novel 3D U-Net model called SGEResU-Net to segment brain tumors. SGEResU-Net embeds residual modules and spatial group-wise enhance (SGE) attention modules into a single 3D U-Net architecture, where SGE attention modules are adopted to enhance the feature learning of semantic regions and reduce possible noise and interference with almost no extra parameters. In addition, motivated by the effectiveness of the self-ensemble module ^[11], we also integrate this module into the network to achieve satisfied segmentation accuracy. The overall architecture of the SGEResU-Net can be shown in Figure 1. We evaluate the SGEResU-Net model using the Brain Tumor Segmentation (BraTS) Challenge 2020 and 2021 datasets, and the experiment results on the BraTS 2021 demonstrate that it can achieve DSC values of 83.31, 91.64 and 86.85%, as well as Hausdorff distances (95%) of 19.278, 5.945 and 7.567 for the enhancing tumor, whole tumor, and tumor core, respectively.

2.   Methods

In this section, we mainly give a detailed description of the spatial group-wise enhance attention residual U-Net (SGEResU-Net) model for the brain tumor segmentation. We first introduce the overall architecture of SGEResU-Net and then describe the utilized 3D SGE, 3D residual and the self-ensemble modules. Finally, the loss function employed is given.

2.1.   Network architecture

The SGEResU-Net model is a trainable 3D brain tumor segmentation network that simultaneously embeds 3D spatial group-wise enhance attention module and residual module into a single U-Net architecture, as shown in Figure 1. The size of the input data for this network is $4\times128\times128\times128$, which means that the image size is $128\times128\times128$, and the channel number is 4. The same as the U-Net model, the core of SGEResU-Net consists of an encoder path, a decoder path and skip connection units. The encoder path includes four residual blocks and a bottom convolutional layer to capture high-level features of the context semantic information, while its decoder path utilizes four normal convolutional blocks to recover features. Besides, to provide the effective information for brain tumor segmentation, the spatial group-wise enhance (SGE) attention module is added as the horizontal connection, which may enhance the feature learning of semantic regions and reduce possible noise and interference with almost no extra parameters, and it will also gain the possible accuracy improvement. Finally, an ensemble module given by Zhao et al. ^[11] is adopted in this network, aiming to combine prediction results on different scales. Using the typical encoder-decoder structure, SGEResU-Net can well conserve the high-level context information in deep layers by fusing shallow features and deep features. To better describe the SGEResU-Net model, we will introduce the details of the utilized 3D SGE, 3D residual and the self-ensemble modules in the following context of this section.

2.1.1.   SGE module

Motivated by the idea of grouping features and attention mechanism, Li et al. ^[15] propose the spatial group-wise enhance (SGE) attention module for visual tasks. SGE module aims to change the importance of semantic features in each semantic group by using an attention factor to reduce the influence of similar patterns and noisy background. More specifically, Li et al. think that CNNs generate the feature representation of complex objects by collecting hierarchical and different parts of semantic sub-features, however, the activation of these sub-features is often spatially affected by similar patterns and noisy backgrounds, which leads to uncorrected localization and classification. To resolve these problems, SGE groups the feature along the channel dimension of the deep feature map to form several sub-features, and then utilizes the attention mechanism guided by similarities between the global and local feature representation inside each semantic feature group. The SGE model not only enhances the feature learning of semantic regions but also compresses noise and interference with almost no extra parameters. Therefore, in this work, we embed the SGE module at each horizontal connection position of the segmentation network. Besides, as the MRI brain tumor image is in 3D format, we extend the original 2D SGE attention module into the 3D module to fit for the input data, and the overall architecture of 3D SGE can be demonstrated in Figure 2. In this way, not only the shallow feature information can be better extracted, but also the deep semantic information can be well combined to achieve a better segmentation result.

2.1.2.   Residual module

As we all know, the residual module can effectively avoid the problem of vanishing or exploding gradients during model training, which also accelerates the convergence of the network ^[16]. In this work, SGEResU-Net utilizes a variety of residual blocks (Resblocks) to capture semantic features of brain tumor images during the down-sampling process, and its basic architecture can be illustrated in Figure 3. Each residual block consists of three convolutional layers, three GroupNorm layers, and two Relu layers. We do not replace all the conventional convolutions with residual modules like the traditional one, but only replace the second conventional convolution of each layer. It is worth noting that the residual module is not employed during the up-sampling process, which means that we adopt the normal convolution block (Convblock) to decode the encoding feature.

2.1.3.   Self-ensemble module

We employ the self-ensemble module given by Zhao et al. ^[11] to assist the segmentation of brain tumors. Fusing prediction results of different layers to generate the final result is a common strategy to improve the network performance in various computer vision tasks. Inspired by this idea, Zhao et al. also utilize the self-ensemble module for brain tumor segmentation, which joins predictions at each scale to obtain the final segmentation result. Following this network, we also embed the self-ensemble module into the last three layers of SGEResU-Net as shown in Figure 1. Specifically, a prediction result is firstly generated in the third-to-last layer, and then up-sampled to the size of the feature map of the second-to-last layer. Then, combined with the prediction generated by the second-to-last layer, it will be up-sampled to the last layer. Finally, it fuses the prediction result of the last layer to generate the final prediction result.

2.2.   Loss function

The brain tumor segmentation task suffers from a serious class imbalance problem. To address this problem, we adopt a combination of Dice loss and Cross-Entropy loss as the loss function for our network. The combined loss function Loss can be defined as follows:

where $\alpha$ is the balance parameter varying from 0 to 1.

The Dice loss is a commonly used loss function for the medical image segmentation, which can make the model effectively learn those samples with unbalanced categories. Meanwhile, the cross-entropy loss is a good solution to the problem of multi-task imbalance, which decreases the difference between the training samples and evaluation metrics. In this way, we combine these two losses to solve the problem of category imbalance to a certain extent. The Dice loss and the cross-entropy loss can be expressed as follows:

where N indicates the total number of classes, $y_{i}$ is the one-hot encoding (0 or 1) for class i, and $\hat{y_{i}}$ denotes the correct predicted probability for the class.

3.   Experimental datasets and settings

In this section, we first introduce the BraTS 2020 and 2021 datasets utilized to evaluate the SGEResU-Net model, followed by the brief description of the pre-processing method. After that, the evaluation metrics and experiment settings are briefly described.

3.1.   Datasets and pre-processing

We adopt two recently released MRI brain tumor benchmarks of BraTS 2020 and BraTS 2021 to evaluate the proposed SGEResU-Net model. The BraTS 2020 dataset contains a training dataset of 369 glioma patients, a validation dataset consisting of 125 patient cases, and a test dataset including 166 samples, respectively. Compared with the BraTS 2020 dataset, the BraTS 2021 dataset largely increases brain tumor patient cases from 660 to 2, 000. More specifically, it includes the training dataset of 1, 251 cases, the validation dataset of 219 cases, and the test dataset of 530 cases, respectively ^{[17,18,19,20,21,22]}. For the two benchmarks, ground truths of training datasets are provided by the BraTS organizers. However, labels of validation and test datasets are unavailable to the public. In order to ensure the fairness and accuracy of the experimental results, the evaluation results of validation and test datasets can be obtained only via the BraTS online website. Besides, each patient case includes four modalities describing with Flair, T1, T1ce and T2, and the size of each image is $240\times240\times155$ pixels. The basic labels include four types, which are called healthy parts (label 0), necrotic and non-enhancing tumors (label 1), edema around the tumor (label 2) and GD enhancing tumors (label 4). The various sub-regions considered for the segmentation evaluation are the whole tumor (the combined area of labels 1, 2 and 4), the core tumor (the combined area of labels 1, 4), and the enhancing tumor (the area of label 4). Figure 4 shows some typical images of a patient case in the BraTS 2021 training dataset.

Since the BraTS 2020 and 2021 datasets are collected in multiple institutions with different scanners and protocols, which makes the intensity values not standardized. Thus, we utilize the z-score method to standardize non-standardized brain tumors, which applies average and standard deviation to process each image. The detailed calculation formula is as follows:

where $z$ is the input image, $z^{\prime}$ is the normalized image. $\mu$ and $\delta$ denote the mean value and standard deviation of the input image, respectively. Besides, as the brain tumor image contains a lot of useless background information, all images are cropped into $128\times128\times128$ pixels as input. In addition, a variety of data transformation strategies such as random rotation, random scaling, random elastic deformation, random flip and random intensity changes are utilized to preprocess the data.

3.2.   Evaluation metrics and settings

In this paper, four commonly used metrics including Dice score/ Dice similarity coefficient (DSC), Hausdorff distance 95%, Sensitivity and Specificity are adopted to evaluate the segmentation network, and they are defined as follows:

where TP, FP, TN and FN denote the number of false negative, true negative, true positive and false positive voxels, respectively. The DSC value, as a metric of ensemble similarity, is generally used to calculate the similarity of two samples. The value range is 0–1, and values close to 1 indicate more similar contours. Sensitivity and Specificity measure voxel-wise overlap between the predicted results and the ground-truth.

where $t$ and $p$ represent the points on the the ground-truth regions $T$ and the predicted regions $P$, respectively. $d \left (t, p \right)$ is the function that computes the distance between points $t$ and $p$. Hausdorff95 is a metric of using Hausdorff distance to measure the 95% quantile of the surface distance. Since Dice score and Hausdorff95 are the overall evaluation metrics for the entire BraTS challenges, we adopt them as the key metric for evaluation consistently across all the challenges.

In our experiments, we utilize the Adam optimizer with an initial learning rate of 0.001. Momentum is 0.95, the weight decay is 1e^-5, and the batch size is set to 4. Our model is implemented using the PyTorch deep learning framework on an NVIDIA GeForce GTX 3090 GPU with 24 GB of memory.

4.   Results and discussion

To verify the effectiveness of the embedded SGE and Residual modules, we firstly conduct basic experiments on the BraTS 2020 dataset, and then compare the SGEResU-Net model with the representative brain tumor segmentation networks. Finally, we focus on experiments performed on the BraTS 2021 dataset to evaluate SGEResU-Net, which contains more brain tumor samples and may provide more stable results.

4.1.   Compared experiment results on the BraTS 2020 benchmark

Comparisons with the baseline on the BraTS 2020 training dataset. Firstly, we perform compared experiment with the baseline model of 3D U-Net on the BraTS 2020 training dataset to evaluate SGEResU-Net. We divide the 2020 training dataset into two data sets with a ratio of 8:2 for training and validation, and the five-fold cross-validation strategy is adopted to achieve more stable results. The compared results are listed in Table 1.

As can be seen in Table 1, our SGEResU-Net can achieve the better segmentation performance than its baseline model, and it obtains DSC of 0.7940, 0.9048 and 0.852 on the enhancing tumor, whole tumor and core tumor segmentation, respectively. It is superior to the U-Net with 0.89, 0.38 and 2.25% accuracy improvement on the enhancing tumor, whole tumor and core tumor, respectively. This experiment demonstrates the effectiveness of SGEResU-Net in the segmentation of brain tumors to a certain extent.

Ablation experiments on the BraTS 2020 validation dataset. Then, we perform ablation experiments to further test the effects of SGE module and Residual module on the BraTS 2020 validation dataset. Here, we utilize the whole BraTs 2020 training dataset of 369 MRI images to train segmentation models, and then submit 125 predicted results of the validation dataset via BraTS online website to achieve the final DSC and Hausdorff95 values. The ablation results on the BraTS 2020 validation dataset are reported in Table 2.

As shown in Table 2, the baseline model of U-Net achieves DSC values of 0.7686, 0.8938 and 0.7992 on enhancing tumor, whole tumor, and core tumor segmentation, respectively. After embedding Residual blocks, it accordingly gains 0.75, 0.41 and 2.62% accuracy improvements on these three regions, which demonstrates the effects of Residual block on brain tumor segmentation task. When it comes to the SGE module, we incorporate it into four skip connections for capturing long-range context features of brain tumor images. As shown in the table, SGEU-Net also gains better segmentation results than its baseline, and its DSC values of 0.7748, 0.8991 and 0.8104 on enhancing tumor, whole tumor, and core tumor segmentation, respectively. Finally, we simultaneously embed Residual blocks and SGE modules into a single 3D U-Net architecture, and construct the final segmentation model of SGEResU-Net. And it achieves DSC values of 0.7742, 0.9026 and 0.8306 on enhancing tumor, whole tumor, and core tumor segmentation, respectively. Compared to its baseline, it accordingly gains 0.56, 0.88 and 3.14% accuracy improvements on these three regions, which further demonstrates the effects of integrating Residual blocks and SGE modules on brain tumor segmentation task. Therefore, ablation experiment results well prove that the effectiveness of SGEResU-Net for the brain tumor segmentation task.

Comparisons with the state-of-the-arts on the BraTS 2020 validation dataset. Additionally, we also perform compared experiments with some state-of-the-art models on the BraTS 2020 validation dataset, whose results are listed in Table 3. In this table, SGEResU-Net achieves DSC values of 77.4, 90.3 and 83.1% on the enhancing tumor, whole tumor, and tumor core, respectively. Especially, our SGEResU-Net gets the optimal result on the whole tumor segmentation, and it outperforms by 1.3% accuracy gains over the second highest result of 0.890 ^[26,27] on the whole tumor segmentation, which further demonstrates that our proposed model is effective for the segmentation of brain tumors. On the enhancing tumor segmentation, Wang et al. ^[26] obtain the optimal DSC value of 0.785 by introducing the self-attention mechanism into the segmentation model of deep learning. Compared with Sundaresan et al. ^[27], Zhang et al. ^[30] and Huang et al. ^[29], the DSC values obtained by our method are 0.4, 0.4 and 7.4% higher on enhancing tumor, respectively. Meanwhile, Zhang et al. ^[30] obtain the optimal DSC values of 0.839 on core tumor, and our SGEResU-Net is slightly lower than it on core tumor segmentation. They simultaneously integrate asymmetric convolution blocks and an expectation-maximization attention module into the DMF-Net architecture to capture more powerful deep features for the brain tumor segmentation, which may lead to higher DSC values on core tumor segmentation. Note that SGEResU-Net outperforms by 1.7% accuracy gains over the second highest result of 0.814 ^[24,26] on core tumor segmentation. When it comes to Hausdorff95 results, SGEResU-Net gains Hausdorff95 distances are 29.8, 7.0 and 11.6 on the enhancing tumor, whole tumor, and tumor core, respectively. Compared with latest counterpart models, SGEResU-Net is only inferior to Sundaresan et al. ^[27] on enhancing tumor, and Wang et al. ^[26] on whole tumor and core tumor. It consistently has a better segmentation performance than models given by Fang et al. ^[28] and Zhang et al. ^[30]. Nevertheless, SGEResU-Net still ranks the third and fourth places on the core tumor and enhancing tumor segmentation considering the Hausdorff95 metric. In general, the compared experiment results demonstrate the competitive performance of SGEResU-Net and again prove its effectiveness for segmenting brain tumors.

4.2.   Compared experiment results on the BraTS 2021 benchmark

Comparisons with the baseline on the BraTS 2021 training dataset. In this experiment, we utilize 1251 labeled cases from BraTS 2021 training dataset to randomly split the data for training and validation with a ratio of 8:2, and the five-fold cross-validation is also adopted. The experiment results are reported in Table 4. As shown in Table 4, SGEResU-Net achieves DSC values of 0.8790, 0.9365 and 0.9219 on enhancing, whole and core tumor segmentation, respectively. The average DSC on three segmentation models is 0.9125. These results well show the effectiveness of SGEResU-Net for brain tumor segmentation. In addition, we also provide some visualization results of SGEResU-Net in Figure 5. Different colors represent the different types of tumors, i.e., the red regions are necrosis and non-enhancing, the yellow regions represent enhancing tumor, the green regions indicate edema. Meanwhile, images from left to right are Flair, Ground Truth, 3D U-Net and SGEResU-Net segmentation results overlaid on flair image, respectively. As shown in this figure, SGEResU-Net can well segment the enhancing, whole and core tumor regions.

Experiment results on the BraTS 2021 validation dataset. In this experiment, we leverage the whole BraTS 2021 training dataset of 1251 MRI images to train the brain tumor segmentation model. Then, it is used to segmentation tumor images in the BraTS 2021 validation dataset, whose results will be uploaded to the official website platform to obtain the final validation dataset results. We tabulate the final evaluation results of SGEResU-Net on the BraTS 2021 validation dataset in Table 5. Quantitatively, SGEResU-Net achieves DSC values of 83.31, 91.64 and 86.85%, as well as Hausdorff distances (95%) of 19.278, 5.945 and 7.567 for the enhancing tumor, whole tumor, and tumor core, respectively. These results further demonstrate the effectiveness of SGEResU-Net for brain tumor segmentation. In addition, this table also lists the mean, standard deviation, median, and 25th and 75th percentile of each metric. Finally, Figure 6 shows the box plot of DSC values obtained on the BraTS 2021 validation dataset. These results to a certain prove the effects of our SGEResU-Net model on this medical image segmentation application.

5.   Conclusions

In this work, we propose a novel 3D U-Net model called SGEResU-Net to segment brain tumors. SGEResU-Net replaces the horizontally linked part of the baseline network with an improved 3D SGE module to adpat the complex feature distribution of brain tumor images. The 3D SGE module can learn sub-features and suppress noise in a targeted manner for each group, and is also a lightweight module that can adapt brain tumor segmentation models without stress. Experiment results on the Brain Tumor Segmentation Challenge 2020 and 2021 benchmarks demonstrate its effectiveness for the MRI brain tumor segmentation task. In the future, we will evaluate SGEResU-Net on other typical medical image segmentation applications. Additionally, we will also exploit where the SGE module is more beneficial for brain tumor segmentation performance in 3D U-Net networks.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China under Grant 61972062, the Young and Middle-aged Talents Program of the National Civil Affairs Commission, the Liaoning BaiQianWan Talents Program, the University-Industry Collaborative Education Program under Grant 201902029013, and the Henan Provincial Department of Science and Technology Research Project under Grant 202102210168.

Conflict of interest

The authors have no conflict of interest.