An underwater image inpainting method based on an improved generative adversarial network model

Abstract

The application discloses an underwater image repairing method based on an improved generative adversarial network model, and comprises the following steps: obtaining an underwater image to be processed; constructing a generative adversarial network model, training the generative adversarial network model by using image samples of an underwater image sample data set, and taking the trained generative adversarial network model as an underwater image repairing model; inputting the underwater image to be processed into the underwater image repairing model, generating an image by using a generator, and completing the repairing of the underwater image. The application takes the generative adversarial network as a basic framework, corrects the color cast of the underwater image by inputting a white balance image in a pre-processing manner, combines a multi-scale structure and a triple attention to design a multi-scale triple attention module and introduce the module into the generator, realizes the cross-dimension interaction of the channel dimension and the spatial dimension of different levels of features, makes the network better learn underwater image features and suppress noise features, and improves the feature extraction capability of the network.

Description

Technical Field

[0001] This invention relates to, and in particular to, an underwater image restoration method based on an improved generative adversarial network model. Background Technology

[0002] With the massive depletion of land resources, the development of marine resources is becoming increasingly urgent. Underwater images serve as carriers of marine information and are of great significance for acquiring marine information and resource exploration. Different wavelengths of light attenuate at different rates underwater. Among them, red light, with the longest wavelength, attenuates the fastest in water, while blue and green light, with shorter wavelengths, attenuate relatively slowly. Therefore, underwater images typically appear blue-green. Furthermore, due to light scattering, underwater images also suffer from low contrast and blurred clarity. Therefore, solving the problems of color distortion, loss of detail, and low contrast in underwater images is a major challenge currently faced.

[0003] Currently, underwater image enhancement methods are generally divided into traditional methods and data-driven methods. Traditional methods mainly fall into two categories: one is image enhancement based on pixels. While this improves the visual quality to some extent, it is computationally expensive and the enhancement effect is unstable. The other is image enhancement based on physical models. This method mathematically models the underwater image degradation process and estimates unknown parameters based on the model. However, this method relies heavily on imaging models and prior information, resulting in poor generalization ability for underwater images in different scenarios and unsatisfactory image enhancement effects. Summary of the Invention

[0004] Purpose of the invention: To address the above problems, the purpose of this invention is to provide an underwater image restoration method based on an improved generative adversarial network model.

[0005] Technical solution: The present invention provides an underwater image inpainting method based on an improved generative adversarial network model, comprising the following steps:

[0006] Step 1: Obtain the underwater image to be processed;

[0007] Step 2: Construct a generative adversarial network (GAN) model and train it using image samples from an underwater image sample dataset. Use the trained GAN model as the underwater image inpainting model. The GAN model includes a generator and a discriminator. The generator includes a first encoding block, a second encoding block, a multi-scale triple attention module, and a decoding block. The output of the first encoding block is connected to the input of the multi-scale triple attention module. The outputs of the multi-scale triple attention module and the second encoding block are respectively connected to the inputs of the decoding block. The output image of the decoder is used as the image generated by the generator.

[0008] Step 3: Input the underwater image to be processed into the underwater image restoration model, and use the generator to generate an image to complete the underwater image restoration.

[0009] Furthermore, the steps for training a generative adversarial network model using image samples from an underwater image sample dataset include:

[0010] Image samples are input into the generator, and then the real image and the image output by the generator are input into the discriminator. The generator and discriminator are alternately iterated until a Nash equilibrium is reached and the loss function value meets the set value, at which point the training is complete.

[0011] Furthermore, during the training of the generative adversarial network model, the calculation of the loss function includes calculating the adversarial loss, L1 loss, and content-aware loss. The adversarial loss is used to capture high-frequency features, the L1 loss to capture low-frequency features, and the content-aware loss to obtain the realistic image. The calculation process is as follows:

[0012] The expression for calculating adversarial loss is:

[0013] L GAN =E(y){log[D(y)]}+E(x){log{1-D[G(x,z)]}}

[0014] In the formula, y is the reference image, E represents the mathematical expectation, D is the discriminator, G is the generator, x is the underwater image, z is the underwater white balance image, G(x,z) represents the fake sample generated by the generator, and D[G(x,z)] is the discriminator's judgment of whether the sample generated by the generator is real or fake. The real image is judged as 1, and the generated fake image is judged as 0.

[0015] The expression for calculating L1 loss is:

[0016] L1 = E x,y [||yG(x,z)||1]

[0017] In the formula, E x,y Denotes the mathematical expectation, and ||||1 denotes the l1 norm;

[0018] The expression for calculating the content-aware loss is:

[0019] L con =E x,y {||Φ(y)-Φ[G(x,z)]||2}

[0020] In the formula, Φ() represents the content function; ||||2 represents the l2 norm;

[0021] By adding hyperparameters to adjust the above three loss functions, the final expression for the total loss function is as follows:

[0022] L = L GAN +λ1L1+λ2L con

[0023] In the formula, λ1 is the weight of L1 loss, λ2 is the weight of content-aware loss, and λ1+λ2=1.

[0024] Furthermore, the multi-scale triple attention module includes three branches, each of which includes convolutional layers with different dilation rates. A triple attention mechanism module is connected after the convolutional layer of each branch. The features of the three branches are then concatenated together, and after passing through convolution, batch normalization, and activation functions, they are finally added to the input of the multi-scale triple attention module as the output of the multi-scale triple attention module.

[0025] Furthermore, there are skip connections between the first coded block and the decoded block, and between the second coded block and the decoded block.

[0026] Furthermore, the first and second encoding modules have the same structure, both including 5 convolutional blocks, each including a convolutional layer, a batch normalization layer and an activation function layer arranged sequentially.

[0027] Furthermore, the input item for the first coding block is an underwater image, and the input item for the second coding block is an underwater white balance image.

[0028] Beneficial Effects: Compared with existing technologies, the significant advantages of this invention are as follows: This invention uses a generative adversarial network (GAN) as its basic framework. It corrects the color cast of underwater images by preprocessing the input white balance image. It combines a multi-scale structure and a triple attention module to design a multi-scale triple attention module and introduce it into the generator. This enables cross-dimensional interaction between the channel and spatial dimensions of different feature levels, allowing the network to better learn underwater image features and suppress noise features, thus improving the network's feature extraction capabilities. Batch normalization and activation functions are introduced in the two encoding and decoding blocks. Batch normalization improves the model training speed, thereby enhancing the network's generalization ability and preventing overfitting. Activation functions accelerate model convergence. Attached Figure Description

[0029] Figure 1 This is a schematic diagram of the generator structure in the embodiment;

[0030] Figure 2 This is a structural block diagram of the improved generative adversarial network model in the embodiment.

[0031] Figure 3 This is a block diagram of the generator in the embodiment;

[0032] Figure 4 This is a structural block diagram of the triple attention mechanism in the embodiment;

[0033] Figure 5 This is a structural block diagram of the multi-scale triple attention module in the embodiment;

[0034] Figure 6 This is a block diagram of the discriminator in the embodiment. Detailed Implementation

[0035] To make the objectives, technical solutions, and advantages of this application clearer, the following detailed description is provided in conjunction with the accompanying drawings and embodiments.

[0036] The underwater image inpainting method based on an improved generative adversarial network model described in this embodiment includes the following steps:

[0037] Step 1: Obtain the underwater image to be processed;

[0038] Step 2: Construct a Generative Adversarial Network (GAN) model and train it using image samples from an underwater image sample dataset. The trained GAN model will then serve as the underwater image inpainting model. The GAN model includes a generator and a discriminator, such as... Figure 1 As shown, the generator includes a first encoding block, a second encoding block, a multi-scale triple attention module, and a decoding block. The output of the first encoding block is connected to the input of the multi-scale triple attention module, and the output of the multi-scale triple attention module and the output of the second encoding block are respectively connected to the input of the decoding block. The image output by the decoder is used as the image generated by the generator.

[0039] Step 3: Input the underwater image to be processed into the underwater image restoration model, and use the generator to generate an image to complete the underwater image restoration.

[0040] In this embodiment, a generative adversarial network (GAN) is used as the basic framework. The underwater image is input into the first coding block. Preprocessing is used to correct the color cast of the underwater image by inputting a white balance image. A multi-scale triple attention module with a multi-scale structure and triple attention design is introduced into the generator to realize cross-dimensional interaction between the channel dimension and spatial dimension of different levels of features. This enables the improved GAN model to better learn underwater image features and suppress noise features, thereby improving the network's feature extraction capability.

[0041] Specifically, the steps for training a generative adversarial network model using image samples from an underwater image sample dataset include:

[0042] Image samples are input into the generator, and then the real image and the image output by the generator are input into the discriminator. The generator and discriminator are alternately iterated until a Nash equilibrium is reached and the loss function value meets the set value, thus completing the training of the network model.

[0043] like Figure 2 As shown, in one example, the large underwater image dataset EUVP is selected as the image sample input to the generator. 11375 image pairs are selected from this dataset as the training set, with all images having a size of 256×256×3. The underwater images are input into the first encoding block, and the preprocessed underwater white balance image is input into the second encoding block. Then, a multi-scale triple attention module is used to obtain multi-scale, cross-dimensional interactive feature information. The image output from the multi-scale triple attention module and the image output from the second encoding block are input into the decoding block. The decoding block upsamples the input image features to restore the image to its true size. The enhanced image output from the decoding block is the generator's output image. This enhanced image, along with the real image, is input into the discriminator. The discriminator determines whether the generator's output image is a real image or a generated fake image, classifying the real image as 1 and the generated fake image as 0. The generator network and the discriminator network iterate alternately, playing a game with each other, until a Nash equilibrium is reached.

[0044] Specifically, during the training of a generative adversarial network (GAN) model, the calculation of the loss function includes calculating the adversarial loss, L1 loss, and content-aware loss. The adversarial loss is used to capture high-frequency features, the L1 loss is used to capture low-frequency features, and the content-aware loss is used to obtain the realistic image. The calculation process is as follows:

[0045] The expression for calculating adversarial loss is:

[0046] L GAN =E(y){log[D(y)]}+E(x){log{1-D[G(x,z)]}}

[0047] In the formula, y is the reference image, E represents the mathematical expectation, D is the discriminator, G is the generator, x is the underwater image, z is the underwater white balance image, G(x,z) represents the fake sample generated by the generator, and D[G(x,z)] is the discriminator's judgment of whether the sample generated by the generator is real or fake. The real image is judged as 1, and the generated fake image is judged as 0.

[0048] The expression for calculating L1 loss is:

[0049] L1 = E x,y [||yG(x,z)||1]

[0050] In the formula, E x,y Let ||||1 represent the mathematical expectation, and let ||||1 represent the L1 norm. The L1 loss function can make the output image of the generator consistent with the reference image.

[0051] The expression for calculating the content-aware loss is:

[0052] L con =E x,y {||Φ(y)-Φ[G(x,z)]||2}

[0053] In the formula, Φ() represents the content function, which can be extracted from the high-level features of the pre-trained VGG-19 network conv5_2 layer; ||||2 represents the l2 norm; content-aware loss can be used to achieve difference finding, information matching, and feature extraction.

[0054] By adding hyperparameters to adjust the above three loss functions, the final expression for the total loss function is as follows:

[0055] L = L GAN +λ1L1+λ2L con

[0056] In the formula, λ1 is the weight of the L1 loss, λ2 is the weight of the content-aware loss, and λ1 + λ2 = 1. In one example, λ1 = 0.7 and λ2 = 0.3 can be selected. In this example, the above network model training process can be implemented based on the PyTorch deep learning framework under the Linux operating system. The computer configuration is a 2080Ti GPU, using the Adam optimizer with a batch size of 32 and a learning rate of 0.0001.

[0057] like Figure 2-3 The schematic diagram of the generative adversarial network (GAN) model shows that the first and second encoding modules have the same structure, each consisting of five convolutional blocks. Each convolutional block includes a convolutional layer, a batch normalization layer, and an activation function layer arranged sequentially (in the diagram, Conv represents convolution, BN represents batch normalization, and LeakyReLU represents activation function), used to downsample the original network image. The decoding block also includes five convolutional blocks. The first four convolutional blocks each include a deconvolutional layer, a BN layer, and a LeakyReLU layer arranged sequentially. The fifth convolutional block includes a convolutional layer and an activation function layer. The decoding block is used to upsample the input image features, restoring the image to its true size. Batch normalization and activation functions are introduced into the first and second encoding blocks and the decoding block. Batch normalization aims to improve the model training speed, enhance the network's generalization ability, and prevent overfitting; activation functions accelerate model convergence. Finally, after passing through the Tanh function, an enhanced image of size 256×256×3 is obtained.

[0058] There are skip connections between the first coded block and the decoded block, and between the second coded block and the decoded block. Specifically, the fourth convolutional block in both the first and second coded blocks is skipped to the first convolutional block in the decoded block, the third convolutional block in both the first and second coded blocks is skipped to the second convolutional block in the decoded block, the second convolutional block in both the first and second coded blocks is skipped to the third convolutional block in the decoded block, and the first convolutional block in both the first and second coded blocks is skipped to the fourth convolutional block in the decoded block.

[0059] To capture features at different levels, this improved generative adversarial network model incorporates a multi-scale triple attention module (MSTA). Through three convolutions with different dilation rates, it can capture features at different scales. Different receptive fields result in different features: a smaller receptive field captures more local and detailed features; conversely, a larger receptive field captures more comprehensive features with a higher semantic level. For example... Figure 5 As shown, the Multi-Scale Triple Attention Module (MSTA) consists of three branches, each containing convolutional layers with different dilation rates d = 1, 2, and 4. After the convolutional layer of each branch, a Triple Attention (TA) module is connected. The features from the three branches are then concatenated, processed through convolution, batch normalization, and activation functions, and finally added to the input of the Multi-Scale Triple Attention Module to form the output of the module.

[0060] like Figure 4 As shown, to enhance useful feature information and suppress the influence of noise, a lightweight Triple Attention (TA) mechanism is introduced. Triple Attention uses three parallel branches to capture cross-dimensional interactions to calculate attention weights. Given an input of C×H×W, the three parallel branches are implemented as follows:

[0061] The first branch establishes spatial attention through the interaction between height H and width W: the input passes through ChannelPool, Conv, Batch Normalization (BN), and finally the Sigmoid activation function to generate spatial attention weights. The second branch is the interaction between channel C and width W: the input features are permuteed to become H×C×W, then MaxPool and AvgPool are applied to height H, resulting in 2×C×W features. After Conv, Batch Normalization, and the Sigmoid activation function, the feature size is restored to C×H×W after Permute. The third branch is the interaction between channel C and height H: the input features are permuteed to become W×H×C, then MaxPool and AvgPool are applied to width W, resulting in 2×H×C features. After Conv, Batch Normalization, and the Sigmoid activation function, the feature size is restored to C×H×W after Permute. Finally, the features C×H×W output from the three parallel branches are summed and averaged.

[0062] Specifically, such as Figure 6 As shown, the discriminator sequentially includes convolutional blocks with kernels of 3×3, 3×3, 3×3, 3×3, and 3×3. The first layer of the five 3×3 convolutional blocks is a convolutional layer, followed by a batch normalization (BN) layer and a LeakyReLu activation function layer, ultimately producing a 16×16×1 patch block output. If the generated image and the reference image have a small difference in a certain region, the discriminator outputs a pixel value approximately 1 for the corresponding patch block; otherwise, it outputs approximately 0. This output represents the average difference between the generated image and the reference image. In one example, the slope of the LeakyReLu activation function layer can be set to 0.2.

[0063] In one example, 90 images were selected from publicly available EUVP images as a test set. These images included underwater images with bluish tint, underwater images with greenish tint, underwater images with low illumination, underwater images with shallow water, and underwater images with yellowish tint. The underwater image restoration model trained using the generative adversarial network model training method described in the above embodiments resulted in enhanced images with better restoration of details, brightness, saturation, and color compared to the original images.

Claims

1. An underwater image inpainting method based on an improved generative adversarial network model, characterized in that, Includes the following steps: Step 1: Obtain the underwater image to be processed; Step 2: Construct a generative adversarial network (GAN) model and train it using image samples from an underwater image sample dataset. Use the trained GAN model as the underwater image inpainting model. The GAN model includes a generator and a discriminator. The generator includes a first encoding block, a second encoding block, a multi-scale triple attention module, and a decoding block. The output of the first encoding block is connected to the input of the multi-scale triple attention module. The outputs of the multi-scale triple attention module and the second encoding block are respectively connected to the inputs of the decoding block. The output image of the decoder is used as the image generated by the generator. The multi-scale triple attention module includes three branches, each of which includes convolutional layers with different dilation rates. A triple attention mechanism module is connected after the convolutional layer of each branch. The features of the three branches are then concatenated together, and after convolution, batch normalization, and activation functions, they are finally added to the input of the multi-scale triple attention module as the output of the multi-scale triple attention module. There is a skip connection between the first encoded block and the decoded block, and a skip connection between the second encoded block and the decoded block; The input item for the first coding block is an underwater image, and the input item for the second coding block is an underwater white balance image; Step 3: Input the underwater image to be processed into the underwater image restoration model, and use the generator to generate an image to complete the underwater image restoration.

2. The underwater image inpainting method based on an improved generative adversarial network model according to claim 1, characterized in that, The steps for training a generative adversarial network model using image samples from an underwater image sample dataset include: Image samples are input into the generator, and then the real image and the image output by the generator are input into the discriminator. The generator and discriminator are alternately iterated until a Nash equilibrium is reached and the loss function value meets the set value, at which point the training is complete.

3. The underwater image inpainting method based on an improved generative adversarial network model according to claim 2, characterized in that, During the training of a generative adversarial network (GAN) model, the calculation of the loss function includes calculating the adversarial loss, L1 loss, and content-aware loss. The adversarial loss is used to capture high-frequency features, the L1 loss is used to capture low-frequency features, and the content-aware loss is used to obtain the realistic image. The calculation process is as follows: The expression for calculating adversarial loss is: L GAN =E(y){log[D(y)]}+E(x){log{1-D[G(x,z)]}} In the formula, y is the reference image, E represents the mathematical expectation, D is the discriminator, G is the generator, x is the underwater image, z is the underwater white balance image, G(x,z) represents the fake sample generated by the generator, and D[G(x,z)] is the discriminator's judgment of whether the sample generated by the generator is real or fake. The real image is judged as 1, and the generated fake image is judged as 0. The expression for calculating L1 loss is: L1=E x,y [||yG(x,z)||1] In the formula, E x,y Let |||||1 represent the mathematical expectation, and |||||1 represent the l1 norm. The expression for calculating the content-aware loss is: L con =E x,y {||Φ(y)-Φ[G(x,z)]||2} In the formula, Φ() represents the content function; |||2 represents the l2 norm; By adding hyperparameters to adjust the above three loss functions, the final expression for the total loss function is as follows: L=L GAN +λ1L1+λ2L con In the formula, λ1 is the weight of L1 loss, λ2 is the weight of content-aware loss, and λ1+λ2=1.

4. The underwater image inpainting method based on an improved generative adversarial network model according to claim 1, characterized in that, The first and second encoding modules have the same structure, both including 5 convolutional blocks. Each convolutional block includes a convolutional layer, a batch normalization layer, and an activation function layer arranged sequentially.

Images