Underwater image color restoration method based on lightweight attention mechanism

Abstract

The application discloses an underwater image color restoration method based on a lightweight attention mechanism, and comprises the following steps: obtaining an underwater image; constructing a generative adversarial network model, including designing a generator network structure and simultaneously adding a lightweight attention mechanism with better effect in a shallow layer of an encoder-decoder to better extract feature information of the underwater image, and including designing a discriminator network structure to distinguish true and false underwater images generated; determining a loss function based on a mean square logarithmic error calculation method for updating network parameters; training the designed network structure using a training set, continuously optimizing the network through cyclic iteration until an output network model is obtained, and inputting a test set into the trained network model to obtain a clear image. The underwater image enhanced by the application is more in line with the observation habits of human eyes, and is improved compared with existing deep learning methods in terms of underwater image quality evaluation indicators such as peak signal-to-noise ratio and structural similarity.

Description

TECHNICAL FIELD

[0001] The present application relates to the field of underwater image enhancement, in particular to an underwater image color restoration method based on a lightweight attention mechanism. BACKGROUND

[0002] Underwater image enhancement technology is one of the important research directions of current underwater image processing technology, which can improve the efficiency of underwater computer vision tasks at a lower shooting cost, and is widely used in fields such as marine resource development, utilization, protection and control. At present, underwater image enhancement algorithms based on deep learning provide more and more help for marine fisheries, marine ecological research, marine terrain mapping and other fields. The marine environment is complex and changeable, and physical phenomena such as the absorption of red light in natural light by water molecules, the refraction of light by different water concentrations, and the scattering of light by suspended particles in water will all affect the imaging effect of underwater light photography. Marine organisms are diverse, and most of them are small in size, so the accuracy requirements for target capture and recognition of underwater robots are high. Although the optical images captured by the camera are rich in details, the weakening of image information in water still cannot meet the requirements, and the mapping of marine terrain also requires high accuracy to ensure its accuracy. Image enhancement algorithms based on deep learning can more easily achieve the purpose of extracting deep information of underwater images, and thus make the algorithm more beneficial in improving the clarity of underwater images, while also adjusting the color tone of underwater images and improving the success rate of underwater robots in marine organism recognition, positioning and tracking tasks, providing technical support for marine fisheries development and marine organism research.

[0003] In terms of processing methods to solve the problems of image detail blur and color distortion in underwater photography, they can be generally divided into two categories: underwater image restoration methods relying on underwater imaging models and underwater image enhancement methods independent of models. Underwater image enhancement algorithms use mathematical methods to process images to make them clear, while underwater image restoration algorithms restore images to "on-land" images according to physical models.

[0004] The algorithm based on deep learning is mainly based on the principle of end-to-end image enhancement. A pair of degraded underwater and clear underwater images are used to form a data set, and then different deep learning frameworks are used to build a model. Li C et al. designed an algorithm named UWCNN based on the principle of convolution operation (Li C, Anwar S, Porikli F. Underwater scene prior inspired deep underwater image and video enhancement[J]. Pattern Recognition, 2020, 98: 107038.). Although this algorithm can realize the function of color correction of degraded images, the output image is also fuzzy and smooth, but the details in the image are not very clear, and the correction of local color is also not good. Liu X et al. designed an algorithm named MLFcGAN based on the principle of conditional generative adversarial network (Liu X, Gao Z, Chen BM. MLFcGAN: Multilevel feature fusion-based conditional GAN for underwater image color correction[J]. IEEE Geoscience and Remote Sensing Letters, 2019, 17(9): 1488-1492.). This algorithm has good ability of local color correction and global tone processing, but the network structure is complex, the model occupies large memory, the running speed is slow, and the details are fuzzy. Uplavikar P M et al. designed an algorithm named UIE-DAL based on the principle of domain adversarial learning (Uplavikar P M, Wu Z, Wang Z. All-in-one underwater image enhancement using domain-adversarial learning[C] / / CVPR workshops. 2019: 1-8.). The details of the underwater image generated by this algorithm are better, but the color correction of the local, especially the boundary, is prone to distortion. SUMMARY

[0005] In order to overcome the shortcomings of the prior art, the present application provides a kind of underwater image color restoration method based on light attention mechanism, to solve the color imbalance problem existing in the processing of traditional enhancement algorithm.

[0006] The technical scheme for achieving the object of the present application is as follows: an underwater image color restoration method based on light attention mechanism, comprising the following steps:

[0007] Step 1, obtaining an underwater image dataset composed of underwater original images and corresponding normal light images;

[0008] Step 2, constructing a generative adversarial network model, which includes a generator and a discriminator, the network structure of the generator is an improved U-shaped semantic segmentation model, and the network structure of the discriminator adopts a full convolutional adversarial network structure;

[0009] Step 3, determining a loss function based on a mean square logarithmic error calculation method for updating network parameters;

[0010] Step 4, training the designed network structure using the underwater image dataset, and continuously optimizing the network through loop iteration until the output network model is obtained;

[0011] Step 5, inputting the collected underwater original images into the trained generative adversarial network model to obtain clear images.

[0012] Preferably, the generator includes an input layer, a hidden layer and an output layer, the input layer is used for inputting underwater original images, the hidden layer is used for convolution calculation and deconvolution calculation on the input images, and the output layer is used for outputting results;

[0013] The hidden layer adopts an encoder-decoder structure composed of 4 pairs of encoding layers and decoding layers;

[0014] The output of each encoding layer is input to the mirror decoding layer according to a skip connection processing method.

[0015] Preferably, the specific structure of the hidden layer of the generator is as follows:

[0016] Encoding layer 1: input a 256*256 feature map with 3 channels, perform one convolution and one pooling to output a 128*128 feature map with 32 channels;

[0017] Encoding layer 2: input the feature map output by the encoding layer 1, perform two convolutions and one pooling to output a 64*64 feature map with 64 channels;

[0018] Encoding layer 3: input the feature map output by the encoding layer 2, perform two convolutions and one pooling to output a 32*32 feature map with 128 channels;

[0019] Encoding layer 4: input the feature map output by the encoding layer 3, perform one convolution to output a 32*32 feature map with 256 channels;

[0020] Decoding layer 1: input the feature map output by the encoding layer 4, perform one upsampling to output a 64*64 feature map with 256 channels;

[0021] Decoding layer 2: input the feature map output by the decoding layer 1, splice the feature map with the feature map output by the encoding layer 3, perform convolution and up-sampling once after splicing, and output a 128*128 feature map with a channel number of 256;

[0022] Decoding layer 3: input the feature map output by the decoding layer 2, splice the feature map with the feature map output by the encoding layer 2, perform convolution and up-sampling once after splicing, and output a 256*256 feature map with a channel number of 128;

[0023] Decoding layer 4: input the feature map output by the decoding layer 3, add a lightweight attention mechanism to the feature map output by the encoding layer 1, and splice the feature map with the feature map output by the previous layer, perform convolution three times after splicing, and output a 256*256 feature map with a channel number of 3.

[0024] Preferably, the discriminator comprises an input layer, a hidden layer and an output layer, the input layer is used for inputting two pictures with a total size of 256*256*6, the hidden layer is used for convolution calculation on the input image, and the output layer is used for outputting a result.

[0025] Preferably, the discriminator hidden layer is composed of 5 convolution filters, specifically:

[0026] Convolution layer 1: input two pictures with a total size of 256*256*6 as a feature map, perform convolution once to output a 128*128 feature map with a channel number of 32;

[0027] Convolution layer 2: input the feature map output by the convolution layer 1, perform convolution once to output a 64*64 feature map with a channel number of 64;

[0028] Convolution layer 3: input the feature map output by the convolution layer 2, perform convolution once to output a 32*32 feature map with a channel number of 128;

[0029] Convolution layer 4: input the feature map output by the convolution layer 3, perform convolution once to output a 16*16 feature map with a channel number of 256;

[0030] Convolution layer 5: input the feature map output by the convolution layer 4, perform convolution once to output a 16*16*1 information distribution matrix.

[0031] Preferably, the loss function in step 3 is:

[0032]

[0033] In the formula, L cGAN is the discriminator loss calculated by the discriminator network Patch GAN, λ1 and λc is an increase formula of the global similarity of the enhanced image, L con is an increase formula of the constitutive content similarity between the generated image and the target image, G is a mapping equivalent to the generation network, and D is a mapping equivalent to the discrimination network.

[0034] Preferably, the discriminator loss calculated by the discrimination network Patch GAN is specifically:

[0035] L cGAN (G,D)=E X,Y [logD(Y)]+E X,Y [log(1-D(X,G(X,Z)))]

[0036] wherein X and Y respectively represent the underwater image information to be enhanced and the enhanced underwater information, Z represents the input random noise, E X,Y (*) represents the expected value of the distribution function when X, Y are independent variables. Wherein E X,Y [log(1-D(X,G(X,Z)))] is the generator loss function, denoted as L G .

[0037] Preferably, the increase formula of the global similarity of the enhanced image is specifically:

[0038] L1(G)=E X,Y,Z [||Y-G(X,Z)||1]

[0039] wherein E X,Y,Z represents the expected value of the distribution function when X, Y, Z are independent variables.

[0040] Preferably, the increase formula of the constitutive content similarity between the generated image and the target image is specifically:

[0041] L con (G)=E X,Y,Z [||Θ(Y)-Θ(G(X,Z))||2]

[0042] wherein Θ(·) represents a feature extraction mapping function fitted in the block5_conv2 layer of the disclosed pre-trained VGG-19 network model, which is used to extract the image content feature information of the underwater image.

[0043] Preferably, the specific steps of training the network structure in step 4 are:

[0044] input the training data set in step 1 into the generative adversarial network model;

[0045] define the hyperparameters for guiding the training process;

[0046] The generator outputs a generated image according to the underwater original image and random noise;

[0047] The discriminator outputs an information distribution matrix according to the generated image and the underwater original image, and calculates the difference between the matrix and a zero matrix by using a mean square logarithmic error calculation method;

[0048] Meanwhile, the discriminator outputs an information distribution matrix according to the corresponding normal light image and the underwater original image, and calculates the difference between the matrix and a matrix with all elements being 1 by using the mean square logarithmic error calculation method;

[0049] The average of the two differences is the discriminator loss function value, and the parameters of the discriminator are updated by minimizing the discriminator loss function value by using an ADAM optimizer;

[0050] The parameters of the discriminator are frozen, the generated image and the corresponding underwater original image are input into the discriminator to obtain the discriminator output information distribution matrix, the difference between the matrix and a zero matrix is calculated by using the mean square logarithmic error calculation method, and the discriminator loss L cGAN is obtained; meanwhile, the difference between the generated image and the corresponding underwater original image is calculated by using the mean square logarithmic error calculation method;

[0051] The difference is brought into the formula to obtain the global similarity increasing formula L1 and the content similarity increasing formula L con ; the L cGAN , L1, L con are weighted and added together to obtain the generator loss function value, and the parameters of the generator are updated by minimizing the generator loss function value by using the ADAM optimizer;

[0052] The parameters of the generator and the discriminator are continuously updated by repeatedly executing the above process to update the model until the loop ends;

[0053] The model corresponding to the minimum loss function value is selected from the recorded generator loss function value as the final model.

[0054] Compared with the prior art, the present application has the following advantages: (1) the network structure is designed, a fully convolutional CGAN model is constructed, and a lightweight attention mechanism is added in the shallow layer of the encoder-decoder to better restore color and increase contrast; (2) the loss function is designed based on the MSLE calculation method, the performance of extracting underwater image deep learning features is improved, and the loss of underwater image details is reduced.

[0055] The present application will be further described below in conjunction with the drawings of the specification. BRIEF DESCRIPTION OF DRAWINGS

[0056] Figure 1 It is a flowchart of the underwater image color restoration method based on the lightweight attention mechanism.

[0057] Figure 2 is a generator structure diagram after design and improvement.

[0058] Figure 3 is a discriminator structure diagram.

[0059] Figure 4 is a model training flowchart of the present application.

[0060] Figure 5 is a specific flowchart of model testing of the present application.

[0061] Figure 6 is a test result schematic diagram of each method on six groups of images, wherein the images from left to right are (a) original underwater images, (b) FUnIE-GAN test result images, (c) UWCNN test result images, (d) MLFcGAN test result images, (e) UIE-DAL test result images, and (f) the present application test result images. Specific embodiments

[0062] Due to the loss of underwater image information in the transmission process, the designed network structure cannot fully extract the deep learning features of the underwater image, and other reasons, the clear underwater image synthesized by the network model based on deep learning has problems such as local color imbalance and fuzzy details. Therefore, the present application provides a kind of underwater image color restoration method based on lightweight attention mechanism. By designing network structure, a kind of fully convolutional CGAN model is constructed, attention mechanism is added in the generation network structure, and Loss function is designed, the extraction performance of deep learning features of underwater image is improved, so that the model can reduce the loss of underwater image details while restoring color, increase contrast, provide protection for exploring and developing the accurate completion of computer vision tasks in marine processes. Through qualitative and quantitative analysis of the generated images and underwater enhanced images obtained by other deep learning algorithms, the effectiveness and robustness of the method are verified.

[0063] A kind of underwater image color restoration method based on lightweight attention mechanism, comprising the following steps:

[0064] Step 1, obtaining underwater images. The data set used for training of the present application is the EUVP (Enhancing Underwater Visual Perception) data set, which is a set of underwater images of different brightness, different clarity, different color tone, collected by Islam et al. using 7 different cameras in different sea areas, which can represent the similarity of underwater images to a certain extent. The network training of the present application only needs to use the paired data set. The present application uses 3700 pairs of data sets for training, and randomly selects 23 pairs from the test data set for detailed qualitative and quantitative analysis, multi-aspect evaluation of the model and comparison of the performance of the model for comparison experiments.

[0065] Step 2, constructing the generator network structure and discriminator network structure in the generative adversarial network model. The present application will construct a fully convolutional CGAN model to improve the extraction performance of the deep learning features of the underwater image, so that the model can restore the color and enhance the contrast while reducing the loss of details, and provide protection for the accurate completion of the computer vision task in the exploration and development of marine processes.

[0066] The generator of the present application is an improved U-shaped semantic segmentation model (U-Net model), which includes an input layer, a hidden layer and an output layer. The input layer is used to input the original underwater image, the hidden layer is used to perform convolution calculation and deconvolution calculation on the input image, and the output layer is used to output the result. The generator structure is as follows Figure 2 The hidden layer is an encoder-decoder structure composed of 4 pairs of encoding layers and decoding layers. The 4 encoding layers together form a shrinkage path, i.e. step-by-step down-sampling of deep information; the 4 decoding layers together form an expansion path, i.e. up-sampling of the enhanced image by restoring the extracted features; because the output of each encoding layer is input to the mirror decoding layer according to the jump connection processing method, the two paths are almost symmetrical and are used for feature extraction and resolution improvement for precise positioning, respectively. The discriminator network structure also includes an input layer, a hidden layer and an output layer, which is a patch GAN with 5 layers and a patch size of 16*16.

[0067] Step 2.1, the image input to the generator is 256*256*3 in size, which is an RGB image with color deviation and degraded by simulated underwater environment. The image is then processed by 4 pairs of encoding layers and decoding layers in the hidden layer, and finally a 256*256*3 RGB image is obtained as the output after processing by the generation network. The specific structure of the hidden layer is as follows:

[0068] Encoding layer 1: input a 256*256 feature map with 3 channels; perform a convolution (convolution kernel size 5*5, convolution kernel number 32, convolution step length 1, padding same) and a pooling (select maximum pooling, size 2*2, padding same); output a 128*128 feature map with 32 channels;

[0069] Encoding layer 2: input the feature map output by the last layer; perform two convolutions (convolution kernel size 4*4, convolution kernel number 64, convolution step length 1, padding same, and use the BN algorithm) and a pooling (select maximum pooling, size 2*2, padding same); output a 64*64 feature map with 64 channels;

[0070] Encoding layer 3: input the feature map output by the last layer; perform two convolutions (convolution kernel size 3*3, convolution kernel number 128, convolution step length 1, padding same, and use the BN algorithm) and a pooling (select maximum pooling, size 2*2, padding same); output a 32*32 feature map with 128 channels;

[0071] Encoding layer 4: input the feature map output by the last layer; perform a convolution (convolution kernel size 3*3, convolution kernel number 256, convolution step length 1, padding same, and use the BN algorithm); output a 32*32 feature map with 256 channels;

[0072] Decoding layer 1: input the feature map output by the last layer; perform an up-sampling (up-sampling factor 2); output a 64*64 feature map with 256 channels;

[0073] Decoding layer 2: input the feature map output by the last layer; splice the feature map with the feature map output by the encoding layer 3, and then perform a convolution (convolution kernel size 3*3, convolution kernel number 256, convolution step length 1, padding same) and an up-sampling (up-sampling factor 2); output a 128*128 feature map with 256 channels;

[0074] Decoding layer 3: input the feature map output by the last layer; splice the feature map with the feature map output by the encoding layer 2, and then perform a convolution (convolution kernel size 3*3, convolution kernel number 128, convolution step length 1, padding same) and an up-sampling (up-sampling factor 2); output a 256*256 feature map with 128 channels;

[0075] Decoding layer 4: input the feature map output by the last layer; add a light attention mechanism (LCBAM) to the feature map output by the encoding layer 1, and splice the feature map with the feature map output by the last layer, then perform convolution once (the convolution kernel size is 3*3, the number of convolution kernels is 128, the convolution step is 1, and the padding is the same), then perform convolution once (the convolution kernel size is 3*3, the number of convolution kernels is 256, the convolution step is 1, and the padding is the same), and then perform convolution once (the convolution kernel size is 4*4, the number of channels is 3, the convolution step is 1, and the padding is the same); output a feature map with 3 channels and 256*256.

[0076] The light attention mechanism (LCBAM) added therein combines the channel attention mechanism and the spatial attention mechanism, and better results can be achieved. The position of the added attention mechanism can ensure effective extraction of features.

[0077] The generator structure is as shown in Figure 2 . Figure 2 The numbers such as 32, 128, 256, etc. marked in the middle represent the number of layers of deep features extracted from the input of each encoding and decoding layer, and 128*128, 64*64 represent the information distribution matrix extracted at each layer after the input information is processed by convolution operation of each encoding and decoding layer.

[0078] Step 2.2, constructing the discriminator network structure. The discriminator is a post-invalidation process Patch GAN with 5 layers and a patch size of 16*16. Compared with the ordinary discriminator, the output of the Patch GAN discriminator is no longer a scalar value in the interval of 0 and 1, but an average of N*N scalar matrix values. The advantage of this processing is that it can distinguish true and false for each patch, that is, it can distinguish local features of the image, and can also consider the influence of different regions on the image, thereby improving the accuracy of local discrimination and meeting the requirements of higher resolution and details. The discriminator structure is as shown in Figure 3 , the discriminator also includes an input layer, a hidden layer and an output layer, the input layer is used for inputting two pictures with a total size of 256*256*6, the hidden layer is used for convolution calculation on the input image, and the output layer is used for outputting results. The discriminator structure is as shown in Figure 3 . The hidden layer is composed of 5 convolution filters. The specific structure of the hidden layer is as follows:

[0079] Convolution layer 1: input two pictures with a total size of 256*256*6 as a feature map; perform convolution once (the convolution kernel size is 3*3, the number of convolution kernels is 32, the convolution step is 2, and the padding is the same); output a feature map with 32 channels and 128*128;

[0080] Convolutional layer 2: input the feature map output by the last layer; perform a convolution (convolution kernel size 3*3, convolution kernel number 64, convolution step 2, padding same, and use BN algorithm); output a 64*64 feature map with 64 channels;

[0081] Convolutional layer 3: input the feature map output by the last layer; perform a convolution (convolution kernel size 3*3, convolution kernel number 64, convolution step 2, padding same, and use BN algorithm); output a 32*32 feature map with 128 channels;

[0082] Convolutional layer 4: input the feature map output by the last layer; perform a convolution (convolution kernel size 3*3, convolution kernel number 64, convolution step 2, padding same, and use BN algorithm); output a 16*16 feature map with 256 channels;

[0083] Convolutional layer 5: input the feature map output by the last layer; perform a convolution (convolution kernel size 4*4, channel number 1, convolution step 1, and padding same); output a 16*16*1 information distribution matrix.

[0084] Among them, the Patch GAN discriminator is compared with the ordinary discriminator, the output is no longer a scalar value in the interval of 0 and 1, but the average of the scalar matrix value of N*N. The advantage of such processing is that each patch can be discriminated true or false, that is, the local features of the image can be discriminated, and the influence of different regions on the image can also be considered, the precision of local discrimination is improved, and the task with higher requirements for resolution and details can be met.

[0085] The model generator structure diagram is the same as that shown in the description, Figure 3 The numbers 32, 64, 128 and the like in the description represent the number of layers of the output matrix of each processing layer, and 128*128, 64*64 and the like represent the size of the output matrix.

[0086] Step 3, define a suitable loss (Loss) function for updating the network parameters, which can improve the model enhancement effect. The Loss function designed in the application can connect local structure and style, global similarity and image composition, and optimize feature extraction.

[0087] The application integrates three linear integrations together, and the Loss function expression of the model network is as follows:

[0088] LOSS=L cGAN (G,D)+λ1L1(G)+λ c L con (G)

[0089] Wherein, λ1=0.7 and λ c= 0.3 is an empirical hyper-parameter scaling factor during the experiment.

[0090] Therefore, the loss function of the whole model is

[0091]

[0092] That is, assuming the G (generator) mapping is invariable first, record the corresponding maximum loss function when the D (discriminator) mapping changes, then change the G mapping, and execute the above process cyclically. Finally, select the minimum loss value corresponding to the recorded generator loss function value as the final loss function value.

[0093] wherein L cGAN The discriminator loss calculated by the discriminator network PatchGAN is:

[0094] L cGAN (G, D) = E X,Y [log D(Y)] + E X,Y [log (1-D (X, G (X, Z)))]

[0095] Wherein, here X and Y represent the underwater image information to be enhanced and the enhanced underwater information respectively, Z represents the input random noise, G is the mapping equivalent to the generation network, and D is the mapping equivalent to the discriminator network. Wherein E X,Y [log (1-D (X, G (X, Z)))] is the generator loss function, denoted as L G .

[0096] Adding a loss term in the loss function can enhance the corresponding ability, for example, adding the formula. Wherein L1 is an increase formula, which can improve the global similarity of the image:

[0097] L1(G) = E X,Y,Z [||Y-G(X,Z)||1]

[0098] Wherein L con is also an increase formula, which can enhance the similarity of the composition content between the generated image and the target image:

[0099] L con (G) = E X,Y,Z [||Θ(Y)-Θ(G(X,Z))||2]

[0100] Wherein, represents the feature extraction mapping function fitted in the block5_conv2 layer of the disclosed pre-trained VGG-19 network model, which is used to extract the image content feature information of the underwater image here. The VGG-19 network model is composed of 19 hidden layers. Generally speaking, the shallower layers in the network are used to extract globally generalized features, while the abstract information that can be extracted by the deep layers is related to the training data set, so the shallower block5_conv2 layer is adopted in the present application.

[0101] For the calculation method of L in the formula of step 3.1 cGAN , L1, L con The calculation method of L, in order to supervise the adversarial training, according to the overall content, color, local texture and style information of the perception image, the present application formulates an objective function (MSLE calculation method) for evaluating the quality of the perception image, which is used to better train the model parameters and improve the model enhancement effect. MSLE measures the difference between actual value and expected value, which is used to improve the calculation accuracy of the Loss function. MSLE considers small differences between small actual values and predicted values as large differences between large actual values and predicted values, and the addition of logarithm reduces the attention of MSLE to the percentage difference between actual values and predicted values and the relative difference between them, which can make more feature points appear and increase the accuracy. The calculation formula of MSLE is:

[0102]

[0103] Wherein, N represents the number of samples, y i represents the true value, represents the predicted value.

[0104] Step 4, using the training set in the underwater image data set to train the designed network structure, and continuously optimizing the network through loop iteration until the output network model is obtained.

[0105] In order to realize the model training, first input the training data set in step 1, which is composed of underwater original image and corresponding normal light image.

[0106] Then define the hyperparameters that guide the training process, and the selection of numerical values is based on past experimental experience and error correction in the experimental process. Among them, the hyperparameter epoch is 200, which represents that the EUVP data set has been trained for 200 epochs in the experimental process, and the batch processing amount of data is 4. As shown in step 1, the total number of training data is 3700 pairs; the value of val_interval is 2000, and the generated samples are verified and saved every 2000 times of training.

[0107] Then, according to the principle of the generative adversarial network, the generator outputs a generated image according to the underwater original image and random noise;

[0108] The discriminator outputs an information distribution matrix according to the generated image and the underwater original image, and calculates the difference between the matrix and a zero matrix by using a mean square logarithmic error (MSLE) calculation method. Meanwhile, the discriminator also outputs an information distribution matrix according to the corresponding normal light image and the underwater original image, and calculates the difference between the matrix and a matrix with all elements being 1 by using the mean square logarithmic error (MSLE) calculation method. The average of the two differences is the discriminator loss function value, and the parameters of the discriminator are updated by minimizing the discriminator loss function value using an ADAM optimizer.

[0109] Subsequently, the parameters of the discriminator are frozen, the generated image generated by the generator and the corresponding underwater original image are input into the discriminator to obtain the discriminator output information distribution matrix, and the difference between the matrix and a zero matrix is calculated by using a mean square logarithmic error (MSLE) calculation method to obtain the discriminator loss L cGAN . Meanwhile, the difference between the generated image generated by the generator and the corresponding underwater original image is calculated by using the mean square logarithmic error (MSLE) calculation method. The difference is brought into the formula to calculate the global similarity increasing formula L1 and the content similarity increasing formula L con . Finally, L cGAN , L1, L con are weighted and added together to obtain the generator loss function value, and the parameters of the generator are updated by minimizing the generator loss function value using an ADAM optimizer.

[0110] The parameters of the generator and the discriminator are continuously updated by repeatedly executing the above process to update the model until the loop ends. Finally, the model corresponding to the minimum loss function value is selected from the recorded generator loss function value as the final model.

[0111] According to the principle of the generative adversarial network, the discriminator will try to maximize the above Loss function during the adversarial training, while the generator will minimize the Loss function during the adversarial training. The model structure and parameters need to be saved by the last sentence of code output. During training, the generator Loss function value and the discriminator Loss function value are output once every training to verify the training effect in real time and understand the experimental progress.

[0112] Step 5, input the underwater original image test data set in step 1 into the trained generative adversarial network model to perform image enhancement. Then save the test image generated by the model enhancement in a specified folder to obtain the final enhanced image.

[0113] Embodiment

[0114] This embodiment compares our method with four deep learning algorithms—FUnIE-GAN, MLFcGAN, UWCNN, and UIE-DAL—on the underwater image dataset obtained in step 1. These four models are also based on deep learning, utilizing neural network knowledge for improvement, and are trained on large amounts of data; they can also be used to complete augmentation tasks. The FUnIE-GAN and MLFcGAN models, like the model proposed in this invention, are improvements based on conditional generative adversarial networks (GANs). The main differences lie in the choice of model structure and the definition of the loss function. The UWCNN model's structure is based on a deep convolutional neural network, while the UIE-DAL model is based on domain adversarial networks. The experiment compares these three types of models to summarize their characteristics. During the testing process, 23 pairs of underwater images were randomly selected, and each image was input into the algorithm model. The output images were processed by each algorithm, and then these images were used to perform qualitative and quantitative tests on the models, including those of this invention. Some experimental results are shown below. Figure 6 As shown, the images from left to right are (a) the original underwater image, (b) the FUnIE-GAN test result, (c) the UWCNN test result, (d) the MLFcGAN test result, (e) the UIE-DAL test result, and (f) the test result of the present invention. Table 1 shows the quantitative evaluation index results obtained from 23 pairs of underwater images used in the test. In summary, the three indicators used in the experiment are: Peak Signal-to-Noise Ratio (PSNR), which evaluates the objective overall similarity between an image and a reference image by measuring the differences between corresponding pixels; Structural Similarity (SSIM), which evaluates the structural composition similarity between an image and a reference image by judging the correlation between image pixels; and Underwater Image Quality Evaluation Index (UIQM), which evaluates whether the image composition matches human visual perception by judging the linear sum of three measurement indicators: color, sharpness, and contrast. The higher the value of each of the three, the better the performance in the evaluated aspect. Clearly, the model of the present invention shows better performance and a more significant enhancement effect.

[0115] Table 1: Evaluation Index Results

[0116]

[0117] Based on the above analysis, the underwater image color restoration method based on a lightweight attention mechanism proposed in this invention produces underwater images that outperform those generated by other models in all aspects. It successfully obtains an enhanced model that can generate richer detail information and achieve local color correction and global tone processing capabilities.

Claims

1. A method for underwater image color restoration based on a lightweight attention mechanism, characterized in that, Includes the following steps: Step 1: Obtain an underwater image dataset, which consists of raw underwater images and corresponding normal light images; Step 2: Construct a generative adversarial network (GAN) model, which includes a generator and a discriminator. The generator's network structure is an improved U-shaped semantic segmentation model, and the discriminator's network structure adopts a fully convolutional adversarial network (FAN) structure. The generator includes an input layer, a hidden layer, and an output layer. The input layer is used to input the original underwater image, the hidden layer is used to perform convolution and deconvolution calculations on the input image, and the output layer is used to output the result. The hidden layer adopts an encoder-decoder structure consisting of four pairs of coding layers and decoding layers; The output of each coding layer is input into the mirror decoding layer using a skip-connection processing method; The specific structure of the generator hidden layer is as follows: Encoding layer 1: Input a 256*256 feature map with 3 channels, perform one convolution and one pooling to output a 128*128 feature map with 32 channels; Encoding layer 2: Input the feature map output from encoding layer 1, perform two convolutions and one pooling, and output a 64*64 feature map with 64 channels; Encoding layer 3: Input the feature map output from encoding layer 2, perform two convolutions and one pooling, and output a 32*32 feature map with 128 channels; Encoding layer 4: Input the feature map output from encoding layer 3, perform one convolution, and output a 32*32 feature map with 256 channels; Decoding layer 1: Input the feature map output from encoding layer 4, perform one upsampling, and output a 64*64 feature map with 256 channels; Decoding layer 2: Input the feature map output by decoding layer 1, concatenate this feature map with the feature map output by encoding layer 3, perform one convolution and upsampling after concatenation, and output a 128*128 feature map with 256 channels; Decoding layer 3: Input the feature map output by decoding layer 2, concatenate this feature map with the feature map output by encoding layer 2, perform one convolution and upsampling after concatenation, and output a 256*256 feature map with 128 channels; Decoding layer 4: Input the feature map output by decoding layer 3, add a lightweight attention mechanism to the feature map output by encoding layer 1, and concatenate this feature map with the feature map output by the previous layer. After concatenation, perform three convolutions to output a 256*256 feature map with 3 channels. The discriminator includes an input layer, a hidden layer, and an output layer. The input layer is used to input two images with a total size of 256*256*6. The hidden layer is used to perform convolution calculations on the input images. The output layer is used to output the results. The discriminator hidden layer consists of 5 convolutional filters, specifically: Convolutional layer 1: Take two images with a total size of 256*256*6 as input feature maps, perform one convolution to output a 128*128 feature map with 32 channels; Convolutional layer 2: Input the feature map output from convolutional layer 1, perform one convolution, and output a 64*64 feature map with 64 channels; Convolutional layer 3: Input the feature map output from convolutional layer 2, perform one convolution, and output a 32*32 feature map with 128 channels; Convolutional layer 4: Input the feature map output from convolutional layer 3, perform one convolution, and output a 16*16 feature map with 256 channels; Convolutional layer 5: Input the feature map output from convolutional layer 4, perform one convolution, and output a 16*16*1 information distribution matrix; Step 3: Determine a loss function based on the mean square logarithmic error calculation method to update the network parameters, specifically: In the formula, L cGAN To discriminate the discriminator loss computed by the Patch GAN network, λ1 and λ c L is the hyperparameter scaling factor, L1 is the increment formula to improve the global similarity of the image, and L... con To enhance the similarity of the constituent content between the generated image and the target image, an additive formula is used, where G is a mapping equivalent to the generation network and D is a mapping equivalent to the discrimination network. Step 4: Train the designed network structure using the underwater image dataset, and continuously optimize the network through iterative iteration until the network model is output. Step 5: Input the acquired raw underwater images into the trained generative adversarial network model to obtain clear images.

2. The underwater image color restoration method based on a lightweight attention mechanism according to claim 1, characterized in that, The discriminator loss calculated by the Patch GAN discriminator network is as follows: L cGAN (G,D)=E X,Y [logD(Y)]+E X,Y [log(1-D(X,G(X,Z)))] Where X and Y represent the underwater image information to be enhanced and the enhanced underwater information, respectively, Z represents the input random noise, and E X,Y (*) represents the expected value of the distribution function when X and Y are independent variables, where E X,Y [log(1-D(X,G(X,Z)))] is the generator loss function, denoted as L. G .

3. The underwater image color restoration method based on a lightweight attention mechanism according to claim 1, characterized in that, The specific formula for increasing the global similarity of images is as follows: L1(G)=E X,Y,Z [||YG(X,Z)||1] In the formula, E X,Y,Z (*) represents the expected value of the distribution function when X, Y, and Z are independent variables.

4. The underwater image color restoration method based on a lightweight attention mechanism according to claim 1, characterized in that, The specific formula for enhancing the similarity of the constituent content between the generated image and the target image is as follows: L con (G)=E X,Y,Z [||Θ(Y)-Θ(G(X,Z))||2] Where Θ(·) represents the feature extraction mapping function fitted to the block5_conv2 layer in the publicly available pre-trained VGG-19 network model, which is used to extract image content feature information of underwater images.

5. The underwater image color restoration method based on a lightweight attention mechanism according to claim 1, characterized in that, The specific steps for training the network structure in step 4 are as follows: Input the training dataset from step 1 into the generative adversarial network model; Define the hyperparameters that guide the training process; The generator produces an image based on the original underwater image and random noise. The discriminator outputs an information distribution matrix based on the generated image and the original underwater image, and calculates the difference between the matrix and the zero matrix using the mean square logarithmic error calculation method. Meanwhile, the discriminator outputs an information distribution matrix based on the corresponding normal light image and the original underwater image, and uses the mean square logarithmic error calculation method to calculate the difference between the matrix and the matrix with all elements being 1; The average of the two differences is the discriminator loss function value, and the discriminator parameters are updated by minimizing the discriminator loss function value using the ADAM optimizer. Freeze the discriminator parameters, input the generator-generated image and the corresponding original underwater image into the discriminator to obtain the discriminator output information distribution matrix, and calculate the difference between this matrix and the zero matrix using the mean square logarithmic error calculation method to obtain the discriminator loss L. cGAN Simultaneously, the difference between the image generated by the generator and the corresponding original underwater image is calculated using the mean square logarithmic error method. Substituting the difference into the formula yields the global similarity increment formula L1 and the content similarity increment formula L2. con ; will L cGAN L1, L con The generator loss function value is obtained by weighted summation of the three values, and the generator parameters are updated by minimizing the generator loss function value using the ADAM optimizer. The above process is repeated to continuously update the parameters of the generator and discriminator to update the model until the loop ends; Select the model with the minimum loss function value from the recorded generator loss function values ​​as the final model.

Images