A haze image restoration method based on generative adversarial network with memory perception module

Abstract

The present invention relates to a haze image restoration method based on a generative adversarial network with a memory perception module. First, under the framework of a Pix2Pix conditional generative adversarial network, a generator G capable of generating a restored image based on a haze image and a discriminator D for judging whether the restored image belongs to a real image are introduced, and network training is performed. The adversarial loss, pixel reconstruction loss and feature perception loss are used as the overall target loss function of the network. During the network training process, paired haze-clear image data sets are input into the network, the generator G and the discriminator D are trained, and the optimizer is used to update the network parameters of G and D. Whether the target loss function value converges is used as the judgment condition for the end of training, and finally the local optimal solution of the network model is obtained. Compared with the traditional image defogging method based on a physical model, the present invention realizes the one-step restoration of haze images, and the network has a high operating efficiency and can obtain high-quality restored images.

Description

Technical Field

[0001] The present invention belongs to the technical field of computer digital image processing, and in particular relates to a haze image restoration method based on a generative adversarial network with a memory perception module. Background Art

[0002] In application fields such as daily photography or medical image detection, due to the existence of the atmosphere, the image quality is inevitably affected by physical factors such as radiation, absorption and scattering of particles in the atmosphere, especially in haze or rainy and snowy weather, the image quality degradation will be more serious. Image degradation will seriously damage the recognition of image details and reduce its use value. Therefore, the study of image dehazing methods has an urgent application demand and is of great significance to fully explore the application value of images.

[0003] Traditional single image dehazing is generally based on the physical model method, which has two shortcomings. First, physical model dehazing is a typical ill-posed problem that requires strong prior conditions or assumptions and is not suitable for image dehazing in any scene. Second, to estimate the restored image, it is necessary to first estimate the intermediate variables related to haze imaging, which causes the accumulation of intermediate errors and the instability of the final result. Summary of the invention

[0004] In view of the above-mentioned deficiencies in the prior art, the object of the present invention is to provide a haze image restoration method based on a generative adversarial network with a memory perception module, which can realize one-step restoration of haze images under haze weather.

[0005] To achieve the above object, the technical solution adopted by the present invention is as follows:

[0006] The present invention provides a method for restoring haze images using a generative adversarial network based on a memory perception module, and the steps are as follows:

[0007] Step 1, construct a haze-clear dataset S: the dataset S includes haze images and clear images, and the haze images and the clear images correspond one to one;

[0008] Step 2: Based on the Pix2Pix conditional generative adversarial network framework and combined with the image dehazing task, a generator G and a discriminator D are constructed and introduced; wherein the generator G is used to learn the mapping relationship between the haze image domain and the clear image domain, and convert the input haze image to the clear image domain to obtain a restored image; the discriminator D is used to determine whether the restored image is a real image, and the discriminator D cooperates with the generator G for adversarial training;

[0009] Step 3, initialize the network parameters and construct the target loss function to constrain the optimization direction of the generator G and the discriminator D in the adversarial training; wherein the target loss function includes pixel reconstruction loss, feature perception loss and adversarial loss generated by adversarial training;

[0010] Step 4, using the paired haze images and clear images in the haze-clear data set S to perform forward calculations on the generator G and the discriminator D, and obtaining the target loss functions of the generator G and the discriminator D respectively;

[0011] Step 5: Fix the parameters of the generator G network and use the Adam optimizer to update the parameters of the discriminator D network.

[0012] Step 6: Fix the parameters of the discriminator D network and use the Adam optimizer to update the parameters of the generator G network.

[0013] Step 7, loop through steps 4, 5, and 6 until the loss function values ​​of the generator G and the discriminator D calculated in step 4 converge at the same time. At this time, the local optimal solution of the generator network parameters can be obtained. Using the fully trained generator, the estimated value of the restored image can be solved.

[0014] Furthermore, in step 2, the defogging network uses the Pix2Pix generative adversarial network as the basic framework of image defogging, and the U-Net network as the backbone network of the generator network G. In order to enable G to better remove image haze and restore image texture, structure and other features, the present invention introduces an attention mechanism and a recurrent neural network, and designs G into a network with stronger generation capabilities, that is, the generator G network is constructed based on the haze information perception memory subnetwork and the improved U-Net subnetwork, and the improved U-Net subnetwork realizes image restoration, including defogging and restoring the texture structure of the image; the haze information perception memory subnetwork realizes different degrees of attention to different fog concentrations.

[0015] Furthermore, the structural design of the generator G includes a haze information perception memory sub-network and an improved U-Net sub-network. Specifically, the steps for constructing the G network are as follows:

[0016] Step 2-1, introduce the haze information perception memory sub-network to guide the network to focus on the areas with thick fog and output the haze information attention map;

[0017] Step 2-2, the improved U-Net subnetwork is derived from the U-Net network. Based on the U-Net, the present invention abandons the pooling layer and the batch normalization layer, introduces a residual dense module to replace the simple convolution layer, and introduces a LeakyReLU function to replace the ReLU function, thereby improving the performance of the network from the aspects of feature extraction capability and nonlinear expression capability;

[0018] Among them, the residual dense block (RDB) is constructed by a dense connection layer and a feature fusion layer. The dense connection layer contains three ResNet residual modules. The input of each residual module is the sum of the outputs of all previous residual modules on the corresponding channel, which can directly pass shallow features to deeper networks. The feature fusion layer connects the features of all previous modules, and finally performs feature fusion through a convolution layer with a convolution kernel of 1×1 to adaptively retain dense features.

[0019] Step 2-3, the input image and the output of the haze information perception memory sub-network are superimposed in the channel direction, and input into the improved U-Net sub-network to obtain the final restored image. That is, the haze image is input into the haze information perception memory sub-network to obtain the haze information attention map, the haze information attention map is superimposed with the haze image in the channel direction, and then input into the improved U-Net sub-network, and the U-Net sub-network outputs the final restored image.

[0020] Furthermore, the haze information perception memory subnetwork is constructed based on the Inception perception module, the ResNet residual module and the LSTM long short-term memory module, wherein the Inception perception module and the ResNet residual module are respectively used to expand the width and depth of the network, and the LSTM long short-term memory module transmits the low-level feature information to the high-level network, so that the network pays more attention to the feature information related to the haze step by step. The haze information perception memory subnetwork includes several layers, and the above modules are sequentially connected into a layer structure. In each layer, the Inception perception module and the ResNet residual module are respectively used to expand the width and depth of the network, and the LSTM long short-term memory module is used to transmit information. The LSTM long short-term memory module is connected to the LSTM module of the next layer to transmit the low-level feature information to the high-level network, so that the network pays more attention to the feature information related to the haze step by step.

[0021] Furthermore, the residual dense module in the improved U-Net subnetwork is constructed based on the dense communication layer and the feature fusion layer. The dense communication layer contains 3 ResNet residual modules. The input of each residual module is the sum of the outputs of all the previous residual modules on the corresponding channel, which can directly pass the shallow features to the deeper network; the feature fusion layer connects the features of all the previous residual modules, and finally performs feature fusion through the convolution layer of the convolution kernel to adaptively retain the dense features.

[0022] Furthermore, in step 2, the discriminator D is used to determine whether the image generated by G belongs to the clear image domain, and cooperates with the generator for adversarial training. Its network mainly includes two parts, as described below:

[0023] The first part is a multi-scale convolution module, which performs operations of halving the size and doubling the channels of the input image. After the input image passes through the convolution modules with a convolution kernel of 5×5, a step size of 1, a step size of 2, and a step size of 4, the original image size feature map, 1 / 2 the original image size feature map, and 1 / 4 the original image size feature map are obtained respectively. Among them, each scale convolution module contains 1 convolution layer and 1 LeakyReLU activation function layer.

[0024] The second part is the feature extraction module. This module is composed of multiple convolution modules connected in series, and performs multiple convolution feature extraction on the input feature map. Each convolution module consists of a convolution layer with a convolution kernel of 5×5 and an activation function layer. Among them, the features of different scales extracted by the multi-scale convolution module in the first part are sent to the convolution layers of different depths of this module to improve the loss of shallow feature information in the deep network.

[0025] The third part is the output module, which consists of a convolution with a stride of 1 and obtains an output feature map with 1 channel.

[0026] Furthermore, the overall objective loss function in step 3 is mainly composed of WGAN-GP adversarial loss, Smooth-L 1 The pixel reconstruction loss and feature perception loss are composed, and the calculation steps are as follows:

[0027] Step 3-1, adversarial loss refers to the loss during adversarial training between the generator G and the discriminator D, which is composed of the generator loss component and the discriminator loss component. Drawing on the idea of ​​WGAN-GP, the wasserstein distance is used as the adversarial loss. This loss function calculates the degree to which the real image is more realistic than the fake image.

[0028] Step 3-2, the reconstruction loss is used to measure the loss value between the generated image and the original input image. 1 and L 2 The loss cannot make the reconstruction loss converge to a meaningful state, so Smooth-L 1 The loss function is used as a pixel reconstruction loss to compare the difference between the generated image and the real image and perform constrained optimization.

[0029] Step 3-3, feature perception loss is used for higher-level feature requirements. The perception loss of this network consists of haze information perception loss and multi-scale perception loss. The haze information perception loss is given by the attention map, the MSE loss of the difference image between the clear image and the haze image.

[0030] Furthermore, the multi-scale perceptual loss is given by the VGG19 classification network trained on the ImageNet dataset. The restored image generated by the generator G and the clear image are simultaneously input into the VGG19 network, and the 2nd, 7th, 12th, 21st and 30th layers are extracted as feature perception layers respectively. 2 The function calculates the loss value between the corresponding feature layers of the two, which is recorded as multi-scale perceptual loss.

[0031] Furthermore, the forward calculation in step 4 refers to inputting the haze image into G, obtaining the restored image and the loss function value of G, inputting the restored image and the clear image into D, calculating the loss function value of D, and the loss function of G and D is the corresponding overall loss function in step 4. BRIEF DESCRIPTION OF THE DRAWINGS

[0032] Figure 1 is a flow chart of the method of the present invention;

[0033] Figure 2 A network principle diagram in an embodiment of the present invention;

[0034] FIG3 is a structural diagram of the generator haze information attention subnetwork in an embodiment of the present invention;

[0035] FIG4 is a diagram showing the structure of a generator automatic encoding and decoding subnetwork in an embodiment of the present invention;

[0036] Figure 5 This is a diagram showing the structure of a discriminator network in an embodiment of the present invention;

[0037] Figure 6 (a)(b)(c) are comparison diagrams of the first set of clear images, haze images, and restored images in an embodiment of the present invention;

[0038] Figure 7 (a)(b)(c) are comparison diagrams of the second set of clear images, haze images, and restored images in an embodiment of the present invention;

[0039] Figure 8 (a)(b)(c) are comparison diagrams of the second set of clear images, haze images, and restored images in an embodiment of the present invention. DETAILED DESCRIPTION

[0040] In order to facilitate the understanding of those skilled in the art, the present invention is further described below in conjunction with embodiments and drawings. The contents mentioned in the implementation modes are not intended to limit the present invention.

[0041] The present invention discloses a method for restoring haze images using a generative adversarial network based on a memory perception module. The method is based on a conditional generative adversarial network framework, and introduces an improved generator G and a discriminator D, respectively, wherein the generator G learns the mapping relationship from haze images to clear image domains, and the discriminator D cooperates with G for adversarial training. During the adversarial training process, a loss function is introduced to constrain the optimization direction of G and D. When the loss function converges, the generator reaches a local optimum, and the restored image estimation value can be solved using the well-trained generator.

[0042] The method proposed in the present invention can realize one-step restoration of haze images under haze weather, has high operating efficiency, and can obtain high-quality restored images.

[0043] like Figure 1 As shown, a method for restoring haze images using a generative adversarial network based on a memory perception module of the present invention comprises the following steps:

[0044] Step 1: construct a haze-clear dataset S, which includes haze images and clear images, and the haze images and the clear images correspond one to one.

[0045] Step 2, such as Figure 2 As shown in the figure, the dehazing network is based on the Pix2Pix generative adversarial network framework, combined with the image dehazing task, and introduces the improved generator G and discriminator D. Among them, the generator G learns the mapping relationship between the haze image domain and the clear image domain, converts the input haze image to the clear image domain, and obtains the restored image. The discriminator D determines whether the restored image generated by G is a real image. The discriminator D cooperates with the generator G for adversarial training.

[0046] In order to make the generator G better remove haze information and restore the texture and structural features of the image, the present invention designs G into a network with stronger generation ability based on the idea of ​​attention mechanism and residual network. The structural design of G includes haze information perception memory subnetwork and improved U-Net subnetwork. The specific steps of G network construction are as follows:

[0047] Step 2-1, the haze information perception and memory subnetwork is shown in Figure 3(a). Its function is to guide the network to focus on areas with denser fog and output a haze information attention map. The network is constructed by the Inception perception module, the ResNet residual module and the LSTM long short-term memory module. The Inception perception module is shown in Figure 3(b). The Inception perception module and the ResNet residual module are used to expand the width and depth of the network respectively. The LSTM long short-term memory module transfers the low-level feature information to the high-level network, making the network pay more attention to the feature information related to haze stage by stage.

[0048] Step 2-2, the improved U-Net subnetwork is shown in Figure 4(c). Based on the U-Net network, the present invention abandons the pooling layer and the batch normalization layer, introduces a residual dense module to replace the simple convolution layer, and introduces a LeakyReLU function to replace the ReLU function, thereby improving the performance of the network in terms of feature extraction capability and nonlinear expression capability.

[0049] Among them, Figure 4(a) is a network diagram of the residual dense block. The residual dense block (RDB) is constructed by a dense connection layer and a feature fusion layer. The dense connection layer contains 3 ResNet residual modules. Figure 4(b) is a network diagram of the ResNet residual module. The input of each residual module is the sum of the outputs of all previous residual modules on the corresponding channel, which can directly pass shallow features to deeper networks. The feature fusion layer connects the features of all previous layers, and finally performs feature fusion through a convolution layer with a convolution kernel of 1×1 to adaptively retain dense features.

[0050] In step 2-3, the input image and the output of the haze information perception memory subnetwork are superimposed in the channel direction and input into the improved U-Net subnetwork to obtain the final restored image.

[0051] In step 2, the discriminator D is Figure 5 As shown in the figure, D is used to determine whether the image generated by G is a clear image, and cooperates with the generator for adversarial training. Its network mainly consists of three parts, as described below:

[0052] The first part is a multi-scale convolution module, which halves the size of the input image and multiplies the channels. After the input image passes through the convolution modules with a convolution kernel of 5×5, a step size of 1, a step size of 2, and a step size of 4, the original image size feature map, 1 / 2 the original image size feature map, and 1 / 4 the original image size feature map are obtained. Each scale of the convolution module contains 1 convolution layer and 1 LeakyReLU activation function layer.

[0053] The second part is the feature extraction module. This feature extraction module is composed of multiple convolution modules connected in series. It performs multiple convolution feature extraction on the input feature map. Each convolution module consists of a convolution layer with a convolution kernel of 5×5 and an activation function layer. Among them, the shallow features of different scales extracted by the first part of the multi-scale convolution module are sent to the convolution layers of different depths of the module to improve the loss of shallow feature information in the deep network.

[0054] The third part is the output module, which consists of a convolution with a stride of 1 and obtains an output feature map with 1 channel.

[0055] Step 3: Initialize the network parameters and construct a target loss function to constrain the optimization direction of the generator G and the discriminator D in the adversarial training; the target loss function includes pixel reconstruction loss, feature perception loss, and adversarial loss generated by adversarial training.

[0056] Adversarial loss refers to the loss during adversarial training between the generator G and the discriminator D, which is composed of the generator loss component and the discriminator loss component. Drawing on the idea of ​​WGAN-GP, the wasserstein distance is used as the adversarial loss. This loss function calculates the degree to which the real image is more realistic than the fake image.

[0057] The reconstruction loss is used to measure the loss between the reconstructed image and the original input image. 1 and L 2 The loss cannot make the reconstruction loss converge to a meaningful state, so Smooth-L 1 The loss function is used as a pixel reconstruction loss to compare the difference between the generated image and the real image and perform constrained optimization.

[0058] Feature perception loss is used for higher-level feature requirements. The perception loss of this network consists of haze information perception loss and multi-scale perception loss. Among them, haze information perception loss is given by the MSE loss of the attention map and the difference image between the clear image and the haze image. The multi-scale perception loss is given by the VGG19 classification network trained on the ImageNet dataset. The restored image generated by the generator G and the clear image are simultaneously input into the VGG19 network, and the 2nd, 7th, 12th, 21st and 30th layers are extracted as feature perception layers respectively. L 2 The function calculates the loss value between the corresponding feature layers of the two, which is recorded as multi-scale perceptual loss.

[0059] Step 4: Use the paired data sets to perform forward calculations on the generator G and the discriminator D to obtain the loss function values ​​of the two. Forward calculation refers to inputting the haze image data set into G to obtain the loss value and restored image of G, and inputting the restored image and the clear image into D to obtain the loss function value of D.

[0060] Step 5, fix the parameters of the generator G network, and use the Adam optimizer to update the parameters of the discriminator D network. The momentum term of the optimizer is 0.5, the decay rate is 0.999, the initial learning rate is 0.0001, and the learning rate uses a linear decay strategy. When the number of training times reaches the 50th time, the learning rate begins to decay linearly from 0.0001 until the learning rate decays to 0 at the 100th training time, and the network stops training.

[0061] Step 6: Fix the parameters of the discriminator D network and use the Adam optimizer to update the parameters of the generator G network. The learning rate parameter setting of the optimizer is the same as step 5.

[0062] Step 7, loop through steps 4, 5, and 6 until the loss function value calculated in step 4 converges. At this point, the local optimal solution for the generator network parameters can be obtained. Using the fully trained generator, the estimated value of the restored image can be calculated.

[0063] like Figure 6 (a)(b)(c), Figure 7 (a)(b)(c), Figure 8 As shown in (a)(b)(c), there are comparison diagrams of clear images, haze images, and restored images of three groups of images, respectively. It can be seen intuitively from the figures that the haze image restoration method based on the generative adversarial network of the memory perception module proposed in the present invention can restore the haze images under haze weather in one step, and can restore the haze images with high quality, significantly improve the image quality, and enhance the ability to resolve image details.

[0064] The present invention has many specific application paths. The above is only a preferred embodiment of the present invention. It should be pointed out that for ordinary technicians in this technical field, several improvements can be made without departing from the principle of the present invention. These improvements should also be regarded as the protection scope of the present invention.

Claims

1. A haze image restoration method based on a generative adversarial network with a memory perception module, characterized in that: Here are the steps: Step 1, construct a haze-clear dataset S: the dataset S includes haze images and clear images, and the haze images and the clear images correspond one to one; Step 2, based on the Pix2Pix conditional generative adversarial network framework, combined with the image dehazing task, construct and introduce the generator G and the discriminator D; wherein the generator G is used to learn the mapping relationship between the haze image domain and the clear image domain, convert the input haze image to the clear image domain, and obtain the restored image; the discriminator D is used to determine whether the restored image is a real image, and the discriminator D cooperates with the generator G for adversarial training; wherein, The generator G network uses the improved U-Net sub-network as the backbone network, and introduces the haze information perception memory sub-network based on the attention mechanism and recurrent neural network, so that the generator G can achieve the defogging function while also restoring the texture and structural features of the image. The specific steps of constructing the generator G are as follows: Step 2-1, introduce the haze information perception memory sub-network and output the haze information attention map; Step 2-2, the improved U-Net sub-network is based on the classic U-Net network. On the basis of U-Net, the pooling layer and batch normalization layer are discarded, the residual dense module is introduced to replace the simple convolution layer, and the LeakyReLU function is introduced to replace the ReLU function, thereby improving the performance of the network in terms of feature extraction ability and nonlinear expression ability; Step 2-3: superimpose the input image and the output of the haze information perception memory sub-network in the channel direction and input them into the improved U-Net sub-network to obtain the final restored image; Step 3, initialize the G and D network parameters, and construct a target loss function to constrain the optimization direction of the generator G and the discriminator D in the generative adversarial training; wherein the target loss function includes pixel reconstruction loss, feature perception loss, and adversarial loss generated by adversarial training; Step 4, using the paired haze images and clear images in the haze-clear data set S to perform forward calculations on the generator G and the discriminator D, and obtaining the target loss functions of the generator G and the discriminator D respectively; Step 5: Fix the network parameters of the generator G and use the Adam optimizer to update the network parameters of the discriminator D. Step 6: Fix the network parameters of the discriminator D and use the Adam optimizer to update the network parameters of the generator G; Step 7, determine whether the objective loss functions of the generator G and the discriminator D converge at the same time. If not, return to step 4, otherwise end the training and obtain the local optimal solution of the network parameters of the generator G. Using the trained generator, the estimated value of the restored image can be calculated to obtain the dehazed image.

2. According to claim 1, a method for restoring haze images using a generative adversarial network based on a memory perception module is characterized in that: The haze information perception memory subnetwork is constructed based on the Inception perception module, the ResNet residual module and the LSTM long short-term memory module. The haze information perception memory subnetwork includes several layer structures. The above modules are sequentially connected into a layer structure. In each layer, the Inception perception module and the ResNet residual module are respectively used to expand the width and depth of the network, and the LSTM long short-term memory module is used to transmit information. The LSTM long short-term memory module is connected to the LSTM module of the next layer to transmit the low-level feature information to the high-level network, so that the network pays more attention to the feature information related to the haze step by step.

3. According to claim 2, the method for restoring haze images using a generative adversarial network based on a memory perception module is characterized in that: The residual dense module in the improved U-Net subnetwork is constructed based on a dense communication layer and a feature fusion layer. The dense communication layer includes three ResNet residual modules. The input of each residual module is the sum of the outputs of all previous residual modules on the corresponding channel, which can directly pass shallow features to deeper networks; the feature fusion layer connects the features of all previous residual modules, and finally performs feature fusion through the convolution layer of the convolution kernel to adaptively retain dense features.

4. A method for restoring haze images using a generative adversarial network based on a memory perception module according to any one of claims 1 to 3, characterized in that: In the step 2, the discriminator D is used to determine whether the image generated by G is a clear image, and cooperates with the generator G to perform adversarial training. The discriminator D network includes a multi-scale convolution module, a feature extraction module and an output module; first, the multi-scale convolution module is used to reduce the size of the input image by half and multiply the channels, and then, the feature extraction module is used to perform multiple convolution feature extraction on the input feature map, and finally, the feature map is output through the output module.

5. According to claim 4, a method for restoring haze images using a generative adversarial network based on a memory perception module is characterized in that: The pixel reconstruction loss in step 3 is used to measure the loss value between the generated image and the original input image, to compare the difference between the generated image and the real image, and to perform constrained optimization; The feature perception loss includes haze information perception loss and multi-scale perception loss, wherein the haze information perception loss is obtained by the MSE loss of the difference image between the attention map, the clear image and the haze image; the multi-scale perception loss is obtained by the classification network, the restored image generated by the generator G and the clear image are simultaneously input into the classification network, the feature perception layer is extracted, and the loss value between the feature layers is calculated as the multi-scale perception loss; The adversarial loss includes a generator loss component and a discriminator loss component, and adopts the wasserstein distance as the adversarial loss function to calculate the degree of authenticity of the real image relative to the fake image.

6. The method for restoring haze images using a generative adversarial network based on a memory perception module according to claim 5 is characterized in that: The forward calculation in step 4 is as follows: input the haze image into the generator G, obtain the restored image and the loss function value of the generator G, and then input the restored image and the clear image into the discriminator D to calculate the loss function value of D.

Images