SRV-GAN: A generative adversarial network for segmenting retinal vessels

1.   Introduction

Retinal blood vessels are continuous and have dendritic structures. The branches start from the optic disc, and the width of the blood vessel decreases as it moves away from the optic disc. At the same time, the optic disc is the confluence of the main blood vessels. Vascular caliber is important for assessing cardiovascular disease risk ^[1]. The diameter, size and morphology of retinal blood vessels are closely related to diabetes, macular disease, glaucoma and other diseases, and there will be hard exudate and other pathological features in the diseased retina. Therefore, the challenges faced by retinal vessel segmentation technology include low capillary and background contrast, mis-segmentation of optic disc boundaries and interference from pathological spots. In the past, doctors determined the morphology of retinal blood vessels by manual segmentation, but this method is time-consuming and laborious, and the efficiency is very low. Therefore, it is of great significance to segment the morphology of retinal blood vessels by computer vision. Over the years, many scholars worldwide have studied automatic retinal segmentation algorithms. At present, there are supervised and unsupervised methods to segment retinal vessels according to whether labeled or unlabeled data are required. Common unsupervised retinal vessel segmentation methods include conventional matched filtering, image morphology processing, vessel tracking, threshold segmentation, region growth, active contour-based methods and graph-based methods.

Based on the centerline of blood vessels, Mendoca et al. ^[2] adopted an iterative region growing method in the segmentation, which combined images generated by morphological filters and achieved good results. Based on probability tracking, Yin et al. ^[3] used a Bayesian method to detect blood vessel edge points and achieved good segmentation accuracy on three publicly available retinal datasets. Ye et al. ^[4] proposed a three-dimensional multi-scale enhancement filter. This method uses three-dimensional Hessian matrix eigenvalues, which can improve the saliency of tiny blood vessels and increase the speed of calculation. Lazar et al. ^[5] proposed a new region growth method that defines the pixel response as a vector, and the nearest neighbor classifier is used to filter the seed points. In order to overcome the false response of the optic disc boundary, a symmetrically constrained multi-scale filtering technique was also proposed. Neto et al. ^[6] proposed a coarse-to-fine retinal vessel segmentation method, which uses spatial correlation, probability and statistics data, curvature analysis, morphological reconstruction and adaptive local thresholds to improve segmentation accuracy on multiple datasets. A method proposed by Nguyen et al. ^[7] is to obtain line detectors of different scales by changing the length of a basic line detector and linearly combining the line responses of different scales. This method is efficient and scalable.

The supervised method is to first extract the features of retinal blood vessels, then train the classifier with manually labeled images and finally use the trained classifier to segment retinal blood vessels. Feature extraction methods include discrete wavelet change, Gaussian filtering, vascular filtering, etc., and the classifier usually adopts a support vector machine, artificial neural networks, or the k-nearest neighbor algorithm. Staal et al. ^[8] proposed a supervised model based on ridge lines for automatic segmentation of retinal vessels. The method uses the sequential forward selection algorithm to get the best eigenvalue of the pixel on the ridge, and it then uses the k-nearest neighbor algorithm to classify each pixel. Ricci et al. ^[9] proposed a retinal vascular segmentation method combining line operations with a support vector machine algorithm. In this method, two orthogonal detectors are combined with the gray values of pixels to extract feature images, and then the support vector machine algorithm is used to complete the classification of pixels. Wilfred et al. ^[10] used the structure of an artificial neural network with multiple hidden layers to segment the retinal vascular structure and experiments show that the accuracy of this method on the DRIVE data set is good.

2.   Related works

In recent years, the deep learning method has been widely applied in various fields. Due to its good segmentation effects and high computational efficiency, more and more researchers have used the deep learning method to segment retinal vessels. Scholars have proposed AlexNet ^[11], VGG ^[12], GoogLeNet ^[13], Residual Net ^[14], DenseNet ^[15] and other models. In terms of semantic segmentation, the fully convolutional neural network proposed by Long et al. has achieved better performance than other convolutional neural networks. Since 2015, the U-Net model proposed by Ronneberger et al. ^[16] is a kind of fully convolutional neural network. The U-Net model is composed of an encoder and decoder as well as skip connections, and it has the abilities of pixel location and feature extraction. Based on the U-Net model, many derivative methods have been produced. Zhou et al. ^[17] proposed UNet++ for lung nodule segmentation, colon polyp segmentation, cell nucleus segmentation and liver segmentation. The core idea is to use dense connection modules in the skip connection part to fuse the semantic gap. Jin et al. ^[18] introduced deformable convolution in retinal image segmentation. The model combines the advantages of deformable units and U-shaped networks. Oktay et al. ^[19] proposed attention U-Net. Its core idea is to introduce an additive attention gate (AG) at the skip connection, which can suppress the feature response of irrelevant background areas. Ding et al. ^[20] proposed a multi-channel neural network that can accurately segment the end of blood vessels and achieved good results on retinal datasets. Sun et al. ^[21] proposed two new data enhancement modules, channel random gamma correction and channel random blood vessel enhancement, so that the model can recognize more features globally and locally. The dense connection network proposed by Li et al. ^[22] extracted retinal vascular information through dense connection blocks, which could alleviate the gradient disappearance in the feature extraction process. Alom et al. ^[23] proposed a cyclic residual network based on the U-Net model, in which residual connections can alleviate gradient disappearance and train deep network information. Combined with a recurrent neural network (RCNN), it can accumulate features and achieve high performance in segmentation tasks, and it performed well on retinal data sets.

To produce a composite image that more closely matches the real data, Goodfellow et al. ^[27] proposed GAN, which generally consists of a generator and discriminator. These two models compete against each other in the process of training and learning and finally reach a game equilibrium. The purpose of a discriminator is to determine whether the input data comes from real data or from a generator. The purpose of a generator is to learn the characteristics of a sample to produce data that confuses the judgment of the discriminator. In the traditional GAN, the input of the generator is a random noise signal. Through the training, finally the generator can output a high quality sample; however, because the random noise of input is uncontrollable, the sample type obtained by the generator is difficult to control. This is not suitable for an accurate-to-pixel task. Therefore, the conditional GAN proposed by Mirza and Osindro ^[28] adds conditional information to both the generator and the discriminator to guide the training of the model, realizing the controllable generated content. FCN and U-Net are mature models, which are widely used in various computer vision tasks with remarkable effects. However, it can be found that the detail features of images are often ignored in previous studies, while GAN can improve the image synthesis performance of the convolution model ^[29]. Radford et al. ^[30] proposed a deep convolutional generative adversarial network. By combining a deep convolutional network with GAN, GAN can be more stable in the training process and accelerate its training, which performs well in various fields of medical image processing. Since then, deep convolutional GANs have been widely used. For example, Pix2Pix proposed by Isola et al. ^[31] is a general framework for image translation based on conditional GAN, which realizes the generalization of model structure and loss function, and it has achieved remarkable results on many image translation data sets. Isola et al. designed a generator structure similar to U-Net and a convolutional discriminator structure, PatchGAN, which inputs local image blocks into the discriminator and achieves superior performance on various data sets. The convolutional GAN proposed by Yang et al. ^[32], which combines short and dense connections, can detect more tiny vessels and is superior to many methods previously proposed in terms of sensitivity and specificity. The GAN proposed by Son et al. ^[33] was used in retinal vascular segmentation and optic disc segmentation. The results showed that the indexes in the vascular segmentation task were significantly improved, while optic disc segmentation was not. In the GAN proposed by Dong et al. ^[34], U-Net was used as the generator, and a fully convolutional neural network was used as the discriminator to perform segmentation experiments on multiple thoracic organs, which proved the reliability and feasibility of GAN for medical images. Zhang et al. ^[35] proposed an improved dense GAN combined with U-Net and proposed a multi-layer attention mechanism for lung CT image segmentation, which improved the segmentation accuracy compared with other methods.

Although there are extensive studies on GAN in medical image processing in the field of radiology, there are few applications of GAN in retinal image processing in the field of ophthalmology. In previous studies on retinal vascular segmentation, it can be found that conditional GAN performs better than U-Net and other convolutional models ^[33]. GAN can deliver important performance in the absence of large tagged datasets and data shortages ^[36,37]. Therefore, this paper proposes a conditional GAN model based on deep convolution for retinal vascular image segmentation, in which a controllable variable is used as an additional input of the generator and discriminator, so as to control the output types of the generator. In our network, this controllable variable is the original fundus image. This setting ensures that the samples generated by the generator in our network are controllable, and the input image pair of the discriminator ensures the mapping between the original fundus image and the vascular segmentation image. Figure 1 shows the overall network structure, where x represents the original fundus image, G represents the generator, G(x) is the segmentation image generated by the generator, D represents the discriminator, and y is the manually labeled image. In the network we designed, the input of the generator is the original fundus image, and the output is the probability map of the same size as the fundus image we input. Obviously, the value range on the probability graph is 0–1, where the value corresponding to each pixel point represents the probability value of a blood vessel. The input of the discriminator is an image pair, namely, the original image and vascular diagram. The task of the discriminator is to distinguish whether the vascular diagram in the image pair is artificially annotated or generated by the generator.

The main work of this paper includes the following: This generator combines a residual unit and cyclic unit, and it uses the R2U-Net model, which can accumulate characteristic information and alleviate gradient disappearance. In the last layer of convolution, spatial attention and channel attention are used to extract global information features and reduce the interference of redundant information. In previous studies, densely connected networks have realized high accuracy in the classification task, so we use dense connection modules in the discriminator. They can alleviate the gradient dissipation problem in the process of training, and due to the large number of features reused, a large number of features can be generated using a small amount of convolution kernels. According to previous experience, we usually mix the loss function of the GAN with the traditional loss function to get good results. However, we cannot guarantee that the generated blood vessel map and the original fundus image correspond pixel by pixel, that is, the generated blood vessel map is not very close to our labeled results, so we also need a loss function that makes the results of the generator correspond to the labeled map. Kamran et al. ^[38] proposed using the mean square error loss function to generate a probability segmentation graph that is closer to ground truth through RV-GAN. Therefore, we adopted this loss in this experiment.

3.   Methods

3.1.   Generator

RCNN and its variants have shown superior performance in target recognition tasks using different benchmarks ^[39,40]. Alom et al. proposed R2U-Net and applied it to medical image segmentation, which has excellent segmentation performance in a variety of data sets. Therefore, we take R2U-Net as the basic model of our generator. As shown in Figure 2, R2U-Net uses recurrent residual blocks instead of the traditional conv+relu layer in the encoding and decoding process, which can effectively increase the network depth; using feature summation at different time steps to obtain more expressive features helps to extract lower levels. In skip connections, instead of cutting in the original U-Net, a cascading operation is used. The loop structure deepens the network level, and the residual structure avoids the problem of gradient disappearance as the depth increases. The advantages of U-Net, residual network and RCNN are combined.

In previous studies, we can see that attention mechanism has been widely used in image segmentation tasks. Liu et al. ^[24] proposed a residual network model fused with an attention mechanism, which highlights shallow details in channel and spatial dimensions through a reverse attention mechanism, thereby effectively fusing deep local features and shallow global information, and high segmentation indexes were obtained on three retinal data sets. The multi-scale fusion network proposed by Yang et al. ^[25] introduces both channel attention and spatial attention, which adapts the weights through the channel attention module to improve the segmentation performance, and it captures long-range feature dependencies through the position attention module. It is superior to other methods on publicly available retinal data sets. Based on R2U-Net, this generator introduces a channel attention module and position attention module to reduce the redundancy of information in the physical signs and extract more effective information. As shown in Figures 3 and 4, we adopt the dual-attention module proposed by Fu et al. ^[26] to establish the interdependence relationship in the channel and spatial dimensions, respectively.

Feature A ∈ R^{C × H × W}, and after feature A is subjected to the convolution operation, features B, C, D are obtained, {B, C} ∈ R^{C × H × W}. B is subjected to reshape and transpose operations, and C is subjected to reshape operation. N = H × W, where N refers to pixels, so R^{C × N} is obtained. The new features obtained by B and C are matrix multiplied and then go through softmax to obtain S, S ∈ R^{N × N}, as shown in formula (1).

Feature D ∈ R^{C × H × W}. Reshape feature D to get R^{C × N}, and then matrix multiply with S. Then, reshape them into R^C×H×W, and finally element-wise sum it with feature A to get the final feature E, which is multiplied by the scale parameter α. α learns weight from 0, E ∈ R^{C × H × W}, as shown in formula (2).

First, the reshaped feature A and the reshaped and transpose A are matrix multiplied and then go through softmax to obtain X, X ∈ R^{C × C}. Then, X and the reshaped A are matrix multiplied, and get R^{C × H × W} after reshaping. Then, element-wise sum it with feature A to get the final feature E, which is multiplied by scale parameter β. β learns weight from 0, E ∈ R^{C × H × W}, as shown in formulas (3) and (4).

3.2.   Discriminator

If the discriminator is slow to respond, the resulting images will converge, and the patterns will start to collapse. Conversely, when the discriminator performs well, the gradient of the generator's loss function vanishes, and learning is slow. Earlier, we mentioned that R2U-Net has a high segmentation accuracy in the retina datasets. Therefore, we should explore an excellent discriminator, compared to the deep residual network (ResNet), to achieve the same accuracy. The number of parameters to be learned by DenseNet is much lower than that of the ResNet, so the learning efficiency is higher. Compared with the ordinary convolutional neural network, the special structure of DenseNet can not only reduce the gradient disappearance problem faced by the deep network but also strengthen the "understanding ability" of the network due to the repeated use of feature graphs. Therefore, this paper adopts the dense connection module in the discriminator. As shown in Figure 5, the discriminator is composed of two convolutional layers and four densely connected blocks. The structure of the convolutional layer is Conv3 × 3-BN-ReLU. The small convolution kernel can ensure the visual perception domain and reduce the parameters, and densely connected modules are composed of BN-ReLU-Conv3 × 3. They improve the back propagation of the gradient, make the network easier to train and can achieve feature reuse, improve efficiency, reduce the amount of parameters and calculation costs and achieve superior performance. Through a series of operations, the sample features can be extracted, and finally the generated samples and the output real samples can be judged by sigmoid, so as to distinguish the ground truth and the generated retinal blood vessel segmentation images.

On the basis of the residual network, Huang et al. proposed a densely connected network. The densely connected network is similar to the residual network. The input of the latter layer is related to the previous layer. Unlike the residual network, the input of each layer is the output of all previous layers, and the output of each layer is also the input of all subsequent layers. Because of this special structure, the densely connected network has the advantage of improving the effect when the number of parameters is reduced, and the densely connected network effectively alleviates the phenomenon of gradient disappearance and reduces the loss of feature information. Its specific formula is shown in formula (5).

In the formula, X₀, X₁, X₂...X_i-1 refers to the feature maps fused by layer 0, layer 1, layer 2, layer i-1, and H_i is a composite function composed of BN-ReLU-Conv3 × 3.

3.3.   Objective function

The loss function that removes the noise vector in the original GAN is Eq (6).

x represents the input image, and y represents ground truth. However, the above formula is not our objective function. Based on the above formula, we introduce the reconstruction loss (mean squared error) in the generator to confirm the difference between the output images of the generated network and the real images, as shown in Eq (7). By using this loss, we ensure that the synthesized images contain more realistic blood vessels.

Combining the loss function of formula (6) with this function, the loss functions of the generator and discriminator are obtained as formulas (8) and (9), respectively.

In the formula, L_gen(G, D) is the generator loss, L_dis(G, D) is the discriminator loss, λ is a hyperparameter and is set to 10, and G and D are two opposing training processes. First, fix G to train D, and then fix D to train G, and so on. In the end, the capabilities of both sub-networks can be improved. When the sample image generated by G is judged by D to be an artificially annotated image, the training ends. At this time, after the newly input image passes through G, the image generated by G can be used as the correct segmentation image.

4.   Results and discussion

To complete the implementation, we used the Pytorch and Tensorflow frameworks on a single GPU machine with 16GB of RAM and an NIVIDIA GEFORCE GTX-1650 SUPER. We tested them on three retinal image datasets: DRIVE ^[8], CHASE_DB1 ^[41] and STARE ^[42]. The optimizer uses the Adam optimizer with learning rate α = 0.0002, β₁ = 0.5, β₂ = 0.999, batch_size = 24, and the number of iterations of the three data sets is 100. It took 24–48 hours to train.

4.1.   Database and preprocessing methods

Table 1 shows the pixel sizes of the three data sets and the numbers of training and test images.

Preprocessing: In the fundus image, because the green channel has high contrast and clear blood vessels, the green channel of the fundus retinal image is selected for processing in the experiment. Then, the image is preprocessed by histogram equalization, normalization and gamma transformation. Considering that the number of training sets is too small, the model is prone to overfitting, resulting in poor classification performance. Therefore, data expansion is required. The expansion methods used in this paper are rotation, mirroring, translation and so on.

4.2.   Experimental indicators and chart analysis

In order to evaluate this experiment objectively, we used 5 evaluation indicators for analysis: sensitivity (SE), specificity (SP), accuracy (AC), F1-score and AUC. The formulas are as follows.

In order to explore whether the dual attention mechanism improves the segmentation performance in this experiment, we conducted comparative experiments. Figures 6–8 are the segmentation comparison images of the SRV-GAN model with and without the dual attention mechanism on the DRIVE, CHASE-DB1 and STARE data sets. The first column is the original image, the second column is ground truth, the third column is the segmentation graph with the dual attention mechanism, and the fourth column is the segmentation graph without attention. From the figures, we can see that the segmentation graph with the dual attention mechanism is closer to ground truth and has higher classification accuracy, which indicates that after the addition of the attention mechanism, the segmentation performance of the generator is indeed improved. The attention mechanism can pay attention to the important parts of feature information and suppress the interference of invalid information. Tables 2–4 are the index data in the experiment. It can be seen from the data in the tables that various segmentation metrics of the model have been improved when the attention mechanism is added, especially the F1-score and sensitivity on the STARE dataset, which further proves that adding an attention mechanism is very necessary.

In this article, we applied our proposed model, U-Net, LadderNe and IterNet models to three publicly available retinal data sets, as shown in Figure 9. The first column shows fundus images, the second column shows the segmentation results of U-Net model, the third column shows the segmentation results of LadderNet, the fourth column shows the segmentation results of IterNet, the fifth column shows the segmentation results of SRV-GAN, and the sixth column shows ground truth. According to the experimental results, it can be seen from Figures 10–12 that SRV-GAN has the highest AUC value, indicating that compared with these models, SRV-GAN achieves the best segmentation accuracy on retinal images. In addition, we list the results of various experimental indicators of other models on the same data sets in Table 5, including sensitivity (SE), specificity (SP), accuracy (AC), F1-score (F1) and AUC. It can be seen that our model is superior to U-Net derived architecture and recent models in AUC-ROC of DRIVE, SA-UNet ^[43] performs better in CHASE-DB1, and R2U-Net performs best in the STARE data set. However, the whole, most of the indicators of the method proposed in this paper are not inferior to or even better than the original methods. Sensitivity and AUC-ROC are representative of segmentation performance, so we have to work harder to improve these two indicators.

5.   Conclusions

In this article, we propose an improved GAN for retinal image segmentation, and we achieved good segmentation results in three publicly available datasets. The experimental results show that, compared with U-Net, LadderNet and IterNet models, the SRV-GAN model proposed in this paper shows better performance in segmentation tasks. We found that GAN has not yet been used in clinical trials, so the performance of external data sets independent of training sets cannot be guaranteed. Whether GAN technology can improve the performance of machine learning in clinical diagnosis needs further research. In the future, we will explore more accurate and stable methods of adversarial training, so that they can be put into clinical trials more quickly.

Acknowledgments

This work was funded and supported in part by the National Natural Science Foundation of China, under Grant 61672386; the Anhui Provincial Natural Science Foundation of China, under Grant 1708085MF142; and the Key Research and Development Plan of Anhui Province, China, under Grant 2022a05020011.

Conflict of interest

The authors declare there is no conflict of interest.