Residual based attention-Unet combing DAC and RMP modules for automatic liver tumor segmentation in CT

1.   Introduction

Among all organ tumors, the mortality of liver tumors is much higher than that of others ^[1]. Many treatments are available for liver tumors, in which the directional resection of the focus area is one of the fundamental approaches for liver tumors. In addition, the rapid development of various imaging techniques (e.g., Ultrasound, Computed Tomography (CT), Optical coherence tomography (OCT), Magnetic Resonance (MRI), etc.) also provides the possibility of clinical computer-aided diagnosis ^[2,3,4]. However, the tumor's location, size, and contrast affect the difficulty of resection and postoperative recovery (e.g., two challenges are shown in Figure 1). Therefore, automatic and accurate segmentation of liver tumors is extremely valuable in the clinical environment.

With the rapid development of computer hardware for the past few decades, scholars have proposed many sophisticated deep learning-based segmentation methods. For example, Mu et al. ^[5] proposed a progressive global perception and local polishing (PCPLP) network to automatically segment COVID-19 infections in CT images, integrating multi-scale features with multi-level features. Experimental evaluations show that the proposed PCPLP effectively improves the learning ability to identify well-defined lung infection regions accurately. Zhou et al. ^[6] suggested a 3D semantic V-net (SV-net) for extraocular muscle (EOM) and optic nerve (ON) extractions in orbital CT images. Qualitative and Quantitative evaluation showed that the method has good performance in automatically extracting total EOM and ON in orbital CT images and measuring their volumes. Finally, Zhao et al. ^[7] proposed a combination of a 3D network and a shape deformation module implemented by a spatial transformation network (STN) to achieve automatic segmentation of lung parenchyma, and the experimental results showed that the method is helpful for clinical diagnosis of COVID-19 infection in CT images. Ronneberger et al. ^[8] proposed the U-Net network, which adopts the skip connection to minimize the loss of detail caused by upsampling. Thanks to its excellent segmentation ability, the U-Net network has become the benchmark network of many medical segmentation models.

Many scholars have made extensive efforts on U-Net. For example, Oktay et al. ^[9] applied the attention mechanism to the U-Net and proposed the attention-Unet. The network can adjust the weight of some feature map elements to make the network focus on easily misclassified targets. Chen et al. ^[10] proposed DeepLab series networks. First, they added void convolution to the convolution module of FCN ^[11]. Then, the authors ^[12] also proposed an atrous spatial pyramid pooling (ASPP) to extract context information of different scales and applied this module to a DeepLab network. Alom et al. ^[13] applied cyclic residual convolution to U-Net and proposed the R2UNet network with the highest sensitivity on the ISIC dataset. Then, they added the attention mechanism threshold and proposed the attention-R2UNet. In addition, Wang et al. ^[14] introduced Squeeze-and-Excitation (SE) block, ASPP, and residual block into U-Net for robust liver segmentation. Zongwei et al. ^[15] improved the hop connection of U-Net and proposed a mesh UNet++ network, which connects different levels of UNet networks through the cascade method.

Moreover, there are also variants of U-type networks, such as V-type and M-type networks. Arnab et al. ^[16] proposed the V-Net network, which covers the scope of application of U-Net to 3D medical images for the first time. Li et al. ^[17] proposed a 2.5D semantic segmentation model H-DenseUNet based on V-Net. They present a transformation processing function to transform 3D CT volumes into 2D adjacent slices. Finally, Jun et al. ^[18] proposed the M-Net network, which uses multi-scale input to extract the features of four-scale input images. To improve the network's ability on multi-scale targets, the author proposed a loss function based on multi-label and used polar coordinates to convert to different scale layers to generate local prediction maps. As a result, the network achieves the state-of-the-art effect on the dataset of optic cups and discs.

Some scholars also employed a non-coding-decoding structure, such as bilateral networks proposed by Chang et al. ^[19]. They abandoned the traditional encoding-decoding design but adopted a novel bilateral network instead. The bilateral network comprises spatial path, semantic path, and feature fusion module connecting the two, in which the spatial path mainly extracts the detailed features of the image. Meanwhile, the semantic path extracts the abstract semantic information by increasing the range of receptive fields through continuous downsampling.

Nevertheless, the end-to-end semantic segmentation model still has some limitations, which can be outlined as follows:

(i) The conventional U-Net is suitable for large targets but not low contrast tumors and small ones.

(ii) Although the segmentation effect of attention-Unet has been improved, the model's speed is limited since there is no downsampling process in the spatial branch.

This paper added the CE structure with the DAC and RMP modules to attention-UNet and infused it with the Residual module to propose the ResCEAttUNet network. In summary, the main contributions of this paper are as follows:

● Use a 2D end-to-end model, which converges faster and has fewer parameters than 2.5D and 3D networks.

● Employ the CE structure via extracting spatial details and multi-scale semantic information to improve tumor edges' segmentation effect and reduce the misclassification rate of small tumors.

● Integrate the residual and batch normalization (BN) modules to the network to improve convergence speed.

We organize the rest of the paper as follows: Section 2 describes the proposed method in detail. Section 3 provided the experiments and results. Finally, Section 4 summarizes the work.

2.   Methods

2.1.   Proposed ResCEAttUnet network

To improve the segmentation ability of the network for tumors with different scales, we employed the improved CE structure proposed by Gu et al. ^[20]. It comprises DAC and RMP blocks. We directly apply this CE module to the attention-Unet. To ensure fair comparisons, we use the official source code provided by Attention-Unet and CE-Net. All the experimental models in this paper are based on Attention-Unet, using the same parameters.

Figure 2(a) shows the structure of the DAC module. The dimension of the output feature map of the structure is consistent with the input one, which minimizes the loss of spatial information while extracting semantic information. We connect the DAC module to the end of the shrinkage path of attention U-Net for multi-scale features.

Figure 2(b) shows the structure of RMP. We adjusted some designs in PSPNET ^[21] to adapt to liver tumor segmentation. Specifically, the global pooling is removed, and the pool sizes of the four branches are set to 2, 3, 5, and 6, respectively. In addition, to reduce the model parameters, the number of channels of the feature map is reduced to 1 by using 1 ⊆ 1 convolution kernel. Finally, we link the DAC and RMP modules to form the CE net. Then, we connect this structure with the shrinkage path of attention-Unet and get the CE-Att-Unet.

Furthermore, to improve the model's generalization ability, we add the BN layer to all the convolution modules in the CE-Att-UNet to reduce the model's sensitivity to parameter initialization. Moreover, to avoid the problem of network degradation, we employ the residual module based on bottleneck convolution (as shown in Figure 2(c)). Finally, the residual and BN form a new convolution module, ResNet-BN. Then, we applied the ResNet-BN to CE-Att-UNet and obtained the proposed network ResCEAttUnet (shown in Figure 3).

Compared with CE-NET, the proposed ResCEAttUnet is improved in the following aspects: First, we use attention-Unet as the basic framework. Then, we use DAC and RMP in CE-Net as the connection point between encoding and decoding regions in attention-Unet. Finally, each convolutional sequence in attention-Unet was replaced with residual connection and BN layer.

2.2.   Loss function

For semantics segmentation, the loss function based on cross-entropy treats each pixel independently, which results in the loss function lacking global information attention. Meanwhile, the Tversky function based on cross-ratio can evaluate the segmentation result as a whole. Its definition is shown in Eq (1).

where i is the index of the channel in the real image, j is the index of the pixel in the predicted image, l is the total number of classes in the segmentation task, N is the total number of pixels in the image, p is the probability of the output of the classification model, and g is true. We use p_ij to denote the probability that pixel j belongs to class i at the time of prediction, and g_ij to denote the probability that pixel j of the real image belongs to class i. If pixel j in the input image belongs to class 0, then g_0j = 1, and g_1j, g_2j, …, g_lj = 0. In addition, two parameters α and β are used to adjust the ratio between false positives and false negatives, and the sum of α and β is always equal to 1.

However, the training process and the segmentation effect are not stable enough. Therefore, we developed a hybrid loss function that combines the two parts with its equation shown in Eq (2).

where L_CE represents a loss function based on cross-entropy, L_TL denotes the Tversky loss function, n denotes the number of iterations of the current training, and N denotes the total number of training iterations. ρ represents the equalization factor used to control two-loss values on the same order of magnitude. α and β represent the weighted proportions of Tversky false-negative and false-positive, respectively.

For semantic segmentation of images, the loss function based on cross-entropy treats each pixel independently without paying attention to the global information. As a result, the optimization process is simple, and the training is stable. Therefore, the cross-entropy based on the pixel points can be used to segment the target quickly and stably in the early stage of training. In the later stage of training, the rough results obtained from the previous segmentation are refined based on the result-oriented Tversky loss function. This coarse-to-fine strategy can ensure that the hybrid loss function can take advantage of different strengths at different stages.

2.3.   Evaluation metrics

We used five metrics to evaluate the proposed method in the experiment, including Sensitivity, Specific, Intersection over Union (IoU), Dice coefficient, and Hausdorff. Their definitions are provided in Eqs (3)–(7) below:

where TP, TN, FP, and FN represent true positive, true negative, false positive, and false negative. d_XY/d_YX represents the maximum value of the shortest distance from any point on edge X/Y to edge Y/X. The Hausdorff coefficient of edge X and Y is the larger value of d_XY and d_YX, which is used to measure the correlation measure of the surface distance between the labeled tumor and the predicted tumor.

3.   Experiments and results

3.1.   Data and implementation

The experiment is evaluated in the LiTS17 ^[22] and 3Dircadb (https://www.ircad.fr/research/3dircadb/) databases. The LiTS17 dataset contains 201 CT datasets, and each set is annotated by three different doctors for the location of the liver and the tumor. The dataset is divided into the training set and test set, of which 131 for training and 70 for the test. The slice spacing was between [0.45 mm, 6.0 mm], the cross-sectional resolution distribution was between [0.55 mm, 1.0 mm], and the number of Z-axis slices varied from 42 to 1026.

The 3Dircadb dataset contains venous phase abdominal-enhanced CT data of ten men and ten women, 15 of which include liver tumors with ground truth.

In our experiments, 15 sets of LiTS dataset were randomly selected for testing, 5 sets of 3Dircadb were randomly selected for testing, and the remaining data were used in a 3:1 ratio for training and validation. We did not perform data augmentation in this experiment.

All the experiments are run on a workstation with Ubuntu 18.04 operating system, graphics card RTX2080Ti, memory 64G, single CPU Intel Xeon Silver 4110, and using the Pytorch1.4 deep learning framework for implementation. The initial learning rate of Adam of the network is set to 0.01, and the initial learning rate of Adam of other networks is 1e-4.

3.2.   Comparison of the different loss function

To illustrate the impact of the hybrid loss function hyperparameter setting more directly on the segmentation results. We compared the segmentation results of the CE-Att-UNet network using different hyperparameters in Figure 4.

It can be seen from Figure 4 that, when the parameter α of the hybrid loss function is relatively large, the proportion of non-tumor pixels in CT slices that are incorrectly classified as tumor pixels is relatively low. It indicates that a larger α can effectively suppress false positives. Meanwhile, When the loss function β parameter is relatively large, the segmentation effect of the tumor edge and small target tumor in the CT slice is better, which indicates that a larger β can effectively suppress false negatives.

Figure 5 shows the change curve of the loss function value of the CE-Att-UNet network during training. The red/blue/green dashed line represents the Dice change of the cross-entropy/Tversky/mixed loss training model on the training set.

It can be seen from Figure 5 that, the cross-entropy loss function based on single-pixel point optimization can converge quickly at the beginning of model training. In the first training round, the Dice of the model reaches 0.60. Because the cross-entropy function has a higher weight at the beginning of training, the convergence speed of the hybrid loss function is also faster.

In addition, Figure 5 also shows that the training networks using cross-entropy loss and hybrid loss are close to full convergence when the model training reaches the seventh round. At the beginning of training, the segmentation effect of the Tversky is poor, so the function converges slowly. When the training reaches the 12th round, the training effect of the Tversky function exceeds the cross-entropy loss function, and it will not approach full convergence until the 16th round.

Figure 5 demonstrates that the hybrid loss function combines the advantages of the cross-entropy and the Tversky loss function. At the beginning of training, the hybrid loss function enables a quick converge, meanwhile refines the segmentation result at the end of the training, thus proving in a superior performance than the other two-loss functions.

3.3.   Ablation

This section evaluates the effectiveness of our proposed network framework by employing different structures and loss functions.

Table 1 shows that CE-Att-Unet(H) and ResCEAttUnet (H) obtained the best performance with the hybrid loss function of α = 0.3, β = 0.7. Furthermore, for the soft-dice loss functions, the Dice score of ResCEAttUnet(s), which added the ResNet-BN module, is higher than that of CE-Att-Unet(s). In contrast, the ResCEAttUnet(H) improves the Dice score compared with CE-Att-Unet(H) for the hybrid loss functions. Thus, we can conclude that the enhanced ResCEAttUnet based on residual and BN convolution modules can improve the effect of liver tumor segmentation.

The ResCEAttUnet network trained by loading the pre-training model weight of Resnet-34 improves the Dice coefficient more than the network trained from zero. Besides, the attention-Unet showed the best performance on the specificity, while the ResCEAttUnet network is slightly inferior to the benchmark network. Therefore, the RMP module in the CE structure weakens attention's ability on suppressing false positives. Nevertheless, the ResCEAttUnet network is optimal on Dice and other evaluation metrics. Therefore, the most significant effect of migration learning by loading the pre-training model is speed up the convergence.

In Figure 6, we visualize the role of adding residual structure and the effect of loading the pre-training model for migration learning through the change process of the loss function. Finally, we intercept the changes of the hybrid loss function in CE-Att-Unet, ResCEAttUnet, and the migration learning network loaded with pre-training network Resnet-34 in the training and verification set in the first 14 epochs.

The loss functions of the above networks are dynamic weighted hybrid loss functions, with α = 0.3, β = 0.7. In the training process, we can see that the convergence speed of ResCEAttUnet is faster. The reasons are as follows: (i) The convolution module of ResCEAttUnet has a BN layer, effectively reducing the internal skew variable shift problem. Therefore, the initial learning rate of the network is set relatively large. (ii) The convolution module of ResCEAttUnet employs the residual structure, which converts the original mapping into identity mapping, speeding up the model's convergence.

By comparing the loss function curves of CE-Att-Unet and ResCEAttUnet, we can see that ResCEAttUnet shows a superior segmentation effect on the verification set. Besides, the convergence speed of the model has been significantly improved by directly loading the weight of the pre-training model. In the primary stage of model training, although the effect of the verification set fluctuates, it eventually tends to be stable.

We can conclude from the ablation analysis that with the continuous improvement of the network structure, the final ResCEAttUnet based on the residuals and BN convolution module effectively improves the accuracy of liver tumor segmentation, and finally, with the ResCEAttUnet network loaded with the pre-trained model weights of Resnet-34. Furthermore, experimental results show that migration learning by loading pre-trained models can further improve segmentation accuracy and accelerate model convergence.

3.4.   Comparison with state-of-the-art methods

To verify the effectiveness and robustness of the ResCEAttUnet network proposed in this paper for liver tumor segmentation, we selected some classical semantic segmentation models for comparison, including UNet, R2UNet, D-LinkNet, and UNet++.

Table 2 compares the segmentation effects of the four commonly used benchmark networks and our proposed ResCEAttUnet. On Dice, IOU, Hausdorff, and sensitivity, the proposed ResCEAttUnet in this paper are optimal, but only the specificity is slightly inferior to D-LinkNet.

Since the CE structure of the D-LinkNet network deploys only a DAC module without the multi-scale pooling module, the reason could be that the RMP module weakens the ability to suppress false positives. Nevertheless, our proposed in this paper are significantly superior to other networks on the other metrics.

Figure 7 shows some typical segmentation results of ResCEAttUnet and its baseline network Attention-Unet on some of the test sets. The red, blue, and green lines indicate the tumor contours declined by the ground truth, the attention-UNet, and our proposed ResCEAttUnet, respectively.

These test datasets are from the validation and test sets that do not participate in the model training. The main challenges include multiple liver tumors in some slices, uneven edges, and low contrast between tumors and normal tissues.

From the experimental results, some normal tissues are wrongly segmented as tumors by the attention-Unet (Figure 7(d), (f), (h), (i)). Still, our proposed ResCEAttUnet successfully avoids this kind of wrong segmentation.

Besides, for a single tumor with distinct contour and high contrast, the segmentation results of the two networks are both satisfying (Figure 7(e), (j)). However, when two tumors are close to each other, Attention-Unet tends to segment them into one (Figure 7(c)).

Furthermore, in Figure 7(a), (b), the attention-UNet mistakenly segmented many tumor edge pixels into the background. Therefore, the Dice is much lower than that of the ResCEAttUnet network. Also, for the large tumors in Figure 7(f), (g), (h), the segmentation results of ResCEAttUnet are very close to the ground truth.

Finally, for the complex tumor edges in Figure 7(e), (f), the segmentation effect of the ResCEAttUnet network is also significantly improved.

3.5.   Comparison of running time of different models

Tables 3 and 4 list the training and testing time (per case) using different models on datasets of LiTS17 and 3DIRCADb, respectively. It can be seen from the tables that, the training and testing time of the proposed Pre-ResCEAttUnet are the highest. Specifically, due to the introductions of the residual module, CE, and Attention-Unet in the ablation experiment, the training and testing time of ResCEAttUnet is gradually increased compared with the other methods. However, more image features are obtained through a deeper network, and the accuracy of the proposed ResCEAttUnet is greatly improved, which has been verified in Section 3.3. On the whole, the training time of ResCEAttUnet is higher than other models, while the test time is only slightly higher than other models. Nevertheless, this strategy of improving accuracy at the time cost is still relatively meaningful in computer-aided diagnosis.

4.   Conclusions

This paper proposed a hybrid network for liver tumor segmentation, which leverages residual, BN, CE technologies to improve the attention-Unet. Compared with the attention-Unet, our proposed network structure achieved the following advantages:

(i) All convolution modules on the encoder and decoder paths adopt the ResNet-BN structure, which can accelerate the convergence speed of the model, and the shrinking path is consistent with ResNet-34. Furthermore, we can load pre-trained ResNet-34 weights for migration learning to make the model converge faster and segmentation better. (ii) The DAC and RMP modules in the context information extraction structure proposed to enhance the network's ability to extract multi-scale features and improve the network's preservation of spatial details. In addition, the CE structure can effectively improve segmentation performance for easily misclassified targets, such as liver tumor edges and small target tumors. (iii) The proposed hybrid loss function based on cross-entropy and Tversky can dynamically allocate the weights of the two primary functions through the current number of iterations to fully play the advantages of different loss functions in various stages. Although our proposed method obtains a higher dice score than other methods, it still has limitations. For one, the large number of network parameters requires significant computing resources. For another, the proposed method tends to result in poor segmentation accuracy when dealing with tumors affected by adjacent organs, as shown in Figure 8.

In conclusion, compared with the Attention-Unet, our proposed network accelerates the convergence speed. Meanwhile, it improved the segmentation performance on small target tumors, the tumor edge pixels, and liver tumors with similar contrast to normal tissues. Thus, it could be a promising tool for liver tumor segmentation in clinical assistance. In addition, the growth and quantification of tumors over time need to be premised on accurate contour segmentation and then calculate its max diameter or volume. However, the segmentation accuracy of liver tumors by state-of-the-art methods is still not high. Therefore, we will focus on the growth of the tumor in future work based on further accurate segmentation.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China (No. 61741106), and the Heilongjiang Provincial Natural Science Foundation of China (No. LH2019F023).

Conflict of interest

The authors declare there is no conflict of interest in this study.