A method for generating pneumonia CT images based on UNet-GAN network

Abstract

The present invention discloses a method for generating pneumonia CT images based on a UNet-GAN network. It belongs to the field of computer vision technology. The generative adversarial network in the present invention includes two main bodies, a generator and a discriminator; first, an existing novel coronavirus infection public data set image is used, and the data set is preprocessed; secondly, a convolutional neural network combined with U-Net and DCGAN is established, and the improved U-Net is embedded in the generator part of DCGAN to realize the process of real image analysis and regeneration, and the discriminant structure in DCGAN realizes the process of network discrimination between real images and generated images; then the training data is input into the constructed generative adversarial network, so that the model can generate high-quality images similar to real pictures; finally, the real image is input into the trained network to obtain the required pneumonia CT image. The newly generated image of the present invention can more clearly retain the feature information of small targets in the sample image, improve the image generation quality and network convergence speed, and provide data for the training of a large number of network models.

Description

Technical Field

[0001] The present invention belongs to the technical field of computer vision and relates to a pneumonia CT image generation method based on a UNet-GAN network. Background Art

[0002] In the prior art, neural network training requires a large data set, but due to the privacy of medical images, there are very few public detection data sets for novel coronavirus infection. Using deep convolutional generative adversarial networks to generate images is a practical way to solve this problem. There are also other medical imaging data expansion methods, such as the algorithm proposed in patent number CN201910028616.2, which has complex training steps and a long network training time.

[0003] Generative Adversarial Networks, also known as GAN networks, are mainly composed of two parts: the generator and the discriminator. The generator is an upsampling structure, which is used to generate an image similar to the real image through deconvolution of the input random noise z, and then input the generated image and the real image into the discriminator together. The discriminator is a downsampling structure, which judges the two input images and outputs the corresponding probability based on the judgment result. If it is judged to be a real image, the output probability result is 1; if it is judged to be a directly generated image, the probability output is 0. When the generator and the discriminator compete with each other, and the output probability of the discriminator reaches 0.5, the discriminator cannot judge whether it is true or false, and the model reaches the optimal level.

[0004] Generally, network models are more accurate in identifying large-scale objects. Large-scale objects generally have high resolution and clear local information. Even in the deep convolution process, they still retain rich features. However, small target objects have low resolution. In the deep convolution process, as high-level features are extracted, local information is gradually marginalized, resulting in weak expression of small target features. Therefore, it is particularly important for the model to be able to sensitively extract small target features. Currently, because the detection algorithm has a high sampling rate, a large receptive field, and a contradiction between space and semantics, it is very unfavorable for the detection of small targets and seriously affects the accuracy of discrimination. Summary of the invention

[0005] Purpose of the invention: In order to overcome the deficiencies in the prior art, the purpose of the present invention is to provide a method for generating pneumonia CT images based on a UNet-GAN network.

[0006] Technical solution: The method for generating pneumonia CT images based on the UNet-GAN network described in the present invention has the following specific operation steps:

[0007] (1) Use the existing public dataset images of novel coronavirus infection and preprocess the dataset;

[0008] (2) Build a training model based on the improved U-Net and GAN network structure;

[0009] (3) Train the generative adversarial network model, feed the pneumonia CT images into the built model in batches, and set parameters to generate data with different quality diversity;

[0010] (4) Input the real pneumonia image into the trained model and finally obtain the generated pneumonia image.

[0011] Furthermore, in step (1), the preprocessing operation on the data set is specifically:

[0012] First, the image is translated, cropped, rotated, and flipped to expand the data samples;

[0013] Then, the sample is gray-scale processed and the image is sharpened to enhance the light and dark contrast of the lesion in the image.

[0014] Furthermore, in step (2), the improved U-Net and GAN network structures are specifically:

[0015] The generator model in the GAN network structure adopts an improved U-Net network, adding a residual structure and multi-scale feature fusion to the U-Net.

[0016] The discriminator model in the network structure consists of five convolutional layers and one global pooling layer. The scales are three 5*5 and two 3*3 convolutional layers, and the output result of the last layer is a probability value with a scale of 1*1.

[0017] Furthermore, the residual structure added to U-Net is specifically: the improved U-Net network downsampling is a structure of connecting a standard block and three residual blocks, the standard block is composed of two 3*3 convolution kernels superimposed, and the residual block is to add a 1*1 convolution kernel on both sides of a 3*3 convolution kernel, and connect a 1*1 convolution kernel at the input and output ends; upsampling is composed of two residual blocks and one standard block.

[0018] Furthermore, the added multi-scale feature fusion is specifically: the generator model jumps the feature maps of corresponding scales in the upsampling and downsampling processes, and upsamples the output feature maps of the first three convolutional layers, and performs fractional convolution with step sizes of 1, 2, and 4, respectively, and fuses them with the feature maps obtained by the last layer of deconvolution.

[0019] Furthermore, in step (3), the pneumonia CT images are sent in batches to the built model, and the parameters are set as follows: the network training batch processing number is set to 20, the initial learning rate is set to 0.0004, the momentum optimization value is set to 0.8, and the number of iterations is set to 10,000.

[0020] The U-Net has superior performance in medical image segmentation and other aspects, and is more sophisticated in medical image processing. The idea of ​​combining U-Net with GAN network is used in image analysis and reconstruction, image translation algorithm, high-resolution image generation and other aspects. The U-Net network is used as a model generator and is more suitable for the reconstruction of pneumonia CT images.

[0021] Beneficial effects: Compared with the prior art, the present invention has the following features: 1. Using an improved U-Net network, a residual block is added to the original network to avoid the disadvantage of small target information loss caused by deep convolution and solve the model degradation problem; multi-scale feature fusion fuses local information with global information to facilitate more complete image recovery; the generator in the GAN network is replaced with an improved U-Net, so that the pre-training data input to the generator has more image feature factors, which is conducive to maintaining model stability and accelerating network convergence. The present invention adopts the idea of ​​combining improved U-Net with GAN, which can enable the network to better learn the characteristics of different feature areas of the target image, improve the problems of small target feature loss in deep convolution, lack of information connection between features, and non-convergence of the GAN network model, and achieve relatively good results. Compared with the original model, the improved model converges faster and the accuracy of generated images is greatly improved; 2. The present invention connects a batch normalization layer after each convolution layer. The BN layer standardizes the data distribution of the network, accelerates the model learning efficiency, reduces the impact of different parameter sizes on the network, enhances the network generalization ability, and alleviates the gradient disappearance problem; 3. The generator model of the present invention improves the image generation effect of pneumonia CT and lays the foundation for good network training. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] Figure 1 It is the improved U-Net network model in the present invention;

[0023] Figure 2 It is the improved generative adversarial network in the present invention;

[0024] Figure 3 Schematic diagram of the improved generative adversarial network structure in the present invention;

[0025] Figure 4 It is the original CT image of covid-19;

[0026] Figure 5 It is a model of covid-19 that generates CT images;

[0027] Figure 6 It is an operation flow chart of the present invention. DETAILED DESCRIPTION

[0028] The present invention is further described below in conjunction with the accompanying drawings and specific embodiments.

[0029] As shown in the figure, the pneumonia CT image generation method based on the UNet-GAN network described in the present invention has the following specific operation steps:

[0030] (1) Use the existing public dataset images of novel coronavirus infection and preprocess the dataset;

[0031] (2) Build a training model based on the improved U-Net and GAN network structure;

[0032] (3) Train the generative adversarial network model, feed the pneumonia CT images into the built model in batches, and set parameters to generate data with different quality diversity;

[0033] (4) Input the real pneumonia image into the trained model and finally obtain the generated pneumonia image.

[0034] Furthermore, in step (1), the preprocessing operation on the data set is specifically:

[0035] First, the image is translated, cropped, rotated, and flipped to expand the data samples;

[0036] Then, the sample is gray-scale processed and the image is sharpened to enhance the light and dark contrast of the lesion in the image.

[0037] Furthermore, in step (2), the improved U-Net and GAN network structures are specifically:

[0038] The generator model in the GAN network structure adopts an improved U-Net network, adding a residual structure and multi-scale feature fusion to the U-Net.

[0039] The discriminator model in the network structure consists of five convolutional layers and one global pooling layer. The scales are three 5*5 and two 3*3 convolutional layers, and the output result of the last layer is a probability value with a scale of 1*1.

[0040] Furthermore, the residual structure added to U-Net is specifically: the improved U-Net network downsampling is a structure of connecting a standard block and three residual blocks, the standard block is composed of two 3*3 convolution kernels superimposed, and the residual block is to add a 1*1 convolution kernel on both sides of a 3*3 convolution kernel, and connect a 1*1 convolution kernel at the input and output ends; upsampling is composed of two residual blocks and one standard block.

[0041] Furthermore, the added multi-scale feature fusion is specifically: the generator model jumps the feature maps of corresponding scales in the upsampling and downsampling processes, and upsamples the output feature maps of the first three convolutional layers, and performs fractional convolution with step sizes of 1, 2, and 4, respectively, and fuses them with the feature maps obtained by the last layer of deconvolution.

[0042] Furthermore, in step (3), the pneumonia CT images are sent in batches to the built model, and the parameters are set as follows: the network training batch processing number is set to 20, the initial learning rate is set to 0.0004, the momentum optimization value is set to 0.8, and the number of iterations is set to 10,000.

[0043] Specifically, based on the superior performance of the U-Net network in medical image processing, the U-Net network is selected as the basic network. Aiming at the problem that small target information is easily lost, residual blocks and multi-scale feature fusion are added to the U-Net on the basis of the original network, avoiding the disadvantage of small target information loss caused by deep convolution and solving the degradation problem of the deep network of the model. The improved U-Net network structure is as follows: Figure 1 .

[0044] At the same time, the improved U-Net network is embedded into the GAN network as the generator part of the GAN network. This makes the pre-training data input to the generator part contain the characteristic factors of pneumonia lesions, which is conducive to maintaining the stability of the model, accelerating the convergence of the network, and improving the real effect of the generated image. The schematic diagram of the improved GAN model is shown in the figure. Figure 2 .

[0045] A pneumonia CT image generation model based on a UNet-GAN network comprises a generator model and a discriminator model; the generator model adopts an improved U-Net network, the generator model is improved by adding a residual structure and multi-scale fusion to the U-Net, firstly, a standard block is connected to three residual blocks for downsampling, then two residual blocks and a standard block are used for upsampling, and features of the same size are jump-connected during the upsampling and downsampling processes; during the downsampling process, the standard block is formed by superimposing two 3*3 convolution kernels, and the residual block is formed by adding a 1*1 convolution kernel on both sides of a 3*3 convolution kernel, and connecting a 1*1 convolution kernel at the input and output ends; the upsampling is composed of two residual blocks and a standard block, and the output feature maps of the first three convolution operation blocks are upsampled by fractional convolutions with step sizes of 1, 2, and 4, respectively, and fused with the final feature map.

[0046] The discriminator model includes five convolutional layers and one global pooling layer. The first three convolutional layers all use convolutional kernels with a scale of 5*5, and the fourth and fifth layers use convolutional kernels with a scale of 3*3.

[0047] The generator model is connected to a BN layer and a ReLU activation function after each 3*3 convolutional layer in the improved U-Net network.

[0048] Furthermore, the discriminator model adds a BN layer and a LeakyRelu activation function after each convolutional layer.

[0049] The specific steps are:

[0050] First, we obtain the covid-19-CT related data set, perform image enhancement and uniform scaling on the training data, and convert the training data into a form suitable for network reading; secondly, we establish a convolutional neural network that combines U-Net and DCGAN, and embed the improved U-Net into the generator part of DCGAN to realize the process of real image analysis and regeneration. The discriminant structure in DCGAN realizes the process of network discrimination between real images and generated images; then we input the training data into the constructed generative adversarial network, and through continuous iteration and optimization of the network structure parameters, enable the model to generate high-quality images similar to real images; finally, we input the real image into the trained network to obtain the required pneumonia CT image;

[0051] First, obtain the covid-19-CT related data set, perform image enhancement and uniform scaling on the training data, and convert the training data into a form suitable for network reading;

[0052] Second, we added residual structure and multi-scale feature fusion to the U-Net network to obtain an improved U-Net, and used the improved U-Net as the generator part of the DCGAN network. We added BN layers and activation functions to the corresponding positions in the previous description, and built a generative adversarial network model under the pytorch framework.

[0053] 3. Take the processed covid-19 CT image as input, first train D to maximize the value of the objective function logD(x)+log(1-D(G(Z)), then keep D unchanged and train G to maximize the value of logD(G(z)), repeat the training of D and G, feed the result of the loss function back to the network during the iteration, and use the Adam optimizer to optimize the network parameters until G and D reach Nash equilibrium;

[0054] Fourth, input the covid-19 test sample into the trained GAN network to obtain the final covid-19 CT generated image.

[0055] In the present invention, in step 1, the CT image is translated, cropped, rotated, and flipped to expand the sample image; the image is grayed using the average method to enhance the high-frequency increment of the image, reduce the blur of the image, sharpen the image, and enhance the image to increase the light and dark contrast of the lesion, so that the lesion is clearer; the original image is as follows Figure 4 .

[0056] In step 2, a GAN network is built, and a residual structure is used. The downsampling of the U-Net network is composed of a standard block and three residual blocks. The standard block is formed by superposition of two 3*3 convolution kernels. The residual block is based on adding a 1*1 convolution kernel on both sides of a 3*3 convolution kernel, and connecting a 1*1 convolution kernel at the input and output ends; upsampling is composed of two residual blocks and a standard block; the residual block is transformed from the original two 3*3 convolution kernels to add a 1*1 convolution kernel on both sides of a 3*3 convolution kernel, and connect a 1*1 convolution kernel at the input and output ends. The lesion images input to the network during the experiment of the present invention are unified into 256*256 size. When the residual structure implements the convolution operation, the first 1*1 convolution layer in the new structure reduces the number of channels of the feature map, reduces the amount of calculation of the convolution process, and the second 1*1 convolution layer restores the dimension of the feature map, retains the accuracy of the feature map, and prevents the model degradation problem caused by deep convolution.

[0057] In step 2, multi-scale feature fusion is used to perform fractional convolution upsampling on the output feature maps of the network standard block and the first two residual blocks, and deconvolution with step sizes of 1, 2, and 4 are used respectively. Then, they are fused with the feature maps output by the last layer of standard blocks. In this way, the shallow feature information of the image can be retained, and the local information can be fused with the global information learned in the deep layer, which can improve the problem of information disconnection between the feature maps.

[0058] In step 3, the network training batch processing number is set to 20, the learning rate is 0.0004, the momentum optimization value is 0.8, and the number of iterations is 10000; the Adam optimizer is used as the optimizer.

[0059] In addition, a BN layer is added after each down-sampling convolution layer. Adding a BN layer can speed up the network training speed and avoid the gradient disappearance or explosion of the network. The present invention connects a batch normalization layer after each convolution layer. The BN layer makes the data distribution of the network more uniform, accelerates the model learning efficiency, reduces the influence of different parameter sizes on the network, enhances the generalization ability of the network, and alleviates the gradient disappearance problem.

[0060] The discriminator of pneumonia CT images consists of five convolutional layers and one global pooling layer. The scale sizes are three 5*5 and two 3*3 convolutional layers, and the output result of the last layer is a probability value with a scale of 1*1.

[0061] The hardware system used in this embodiment is Windows 10, the processor is Intel (R) Core (TM) i5-1035G1 CPU @ 1.00GHz 1.19GHz, the machine has 16G RAM, the software operating environment is Pycharm, the programming language is python 3.7, and the Pytorch framework is adopted.

[0062] The present invention adopts two mechanisms, residual structure and multi-scale feature fusion, which can accelerate the gradient flow during training and convergence, so that the network will not be as convolutional as the convolution block, resulting in too serious information loss, and solve the problems of gradient diffusion and network degradation; at the same time, it fuses local information with global information learned in the deep layer, solves the problem of small target feature information loss caused by deep convolution, and is more convenient to recognize the characteristics of the lesion. Through experiments, the time cost of the present invention is reduced, the generation effect is improved, and the model generates images such as Figure 5 .

[0063] The above are only preferred embodiments of the present invention. The protection scope of the present invention is not limited to the above embodiments. All technical solutions under the concept of the present invention belong to the protection scope of the present invention. It should be pointed out that for ordinary technicians in this technical field, some improvements and modifications without departing from the principle of the present invention should be regarded as the protection scope of the present invention.

Claims

1. A pneumonia CT image generation method based on UNet-GAN network, characterized in that: The specific steps are as follows: (1) Use the existing public dataset images of novel coronavirus infection and preprocess the dataset; Specifically: First, the image is translated, cropped, rotated, and flipped to expand the data samples; Then, the samples are gray-scaled and image sharpened to enhance the light-dark contrast of the lesions in the image; (2) Build a training model based on the improved U-Net and GAN network structure; in, The generator model in the GAN network structure adopts an improved U-Net network, adding a residual structure and multi-scale feature fusion to the U-Net. The discriminator model in the network structure includes five convolutional layers and one global pooling layer. The scales are three 5*5 and two 3*3 convolutional layers, and the output of the last layer is a probability value with a scale of 1*1. The residual structure added to U-Net is: the improved U-Net network downsampling is a structure of connecting a standard block and three residual blocks, the standard block is composed of two 3*3 convolution kernels superimposed, the residual block is a 1*1 convolution kernel added to both sides of a 3*3 convolution kernel, and a 1*1 convolution kernel is connected to the input and output ends; the upsampling is composed of two residual blocks and a standard block; Specifically, the GAN network is built using a residual structure. The downsampling of the U-Net network consists of a standard block and three residual blocks. The standard block is composed of two 3*3 convolution kernels superimposed. The residual block is based on adding a 1*1 convolution kernel on both sides of a 3*3 convolution kernel, and connecting a 1*1 convolution kernel at the input and output ends; the upsampling is composed of two residual blocks and a standard block; the residual block is transformed from the original two 3*3 convolution kernels superimposed to adding a 1*1 convolution kernel on both sides of a 3*3 convolution kernel, and connecting a 1*1 convolution kernel at the input and output ends; The added multi-scale feature fusion is: the generator model jumps the feature maps of the corresponding scales in the upsampling and downsampling processes, and upsamples the output feature maps of the first three convolutional layers, performs fractional convolutions with step sizes of 1, 2, and 4, and fuses them with the feature maps obtained by the last deconvolution layer; Specifically, the output feature maps of the network standard block and the first two residual blocks are fractionally upsampled using multi-scale feature fusion, and deconvolution with step sizes of 1, 2, and 4 is used respectively. Then, they are fused with the feature maps output by the last layer of standard blocks, so that the shallow feature information of the image can be retained, and the local information can be fused with the global information learned in the deep layer. (3) Train the generative adversarial network model, feed the pneumonia CT images into the built model in batches, and set parameters to generate data with different quality diversity; Specifically, the network training batch size is set to 20, the initial learning rate is 0.0004, the momentum optimization value is 0.8, and the number of iterations is 10000; (4) Input the real pneumonia image into the trained model and finally obtain the generated pneumonia image.

Images