Underwater optical image enhancement method

Abstract

According to the underwater optical image enhancement method provided by the invention, distortion of an underwater image can be reduced, a clearer and vivid visual effect can be recovered, a matching pair data set containing an original image and a reference image is constructed, and gamma transformation is used for preprocessing; the enhancement network adopts a coding-decoding structure, a coder comprises multi-scale convolution and a residual structure to extract global features, and a decoder recovers a clear image through up-sampling; meanwhile, a multi-loss function containing color, structure and texture constraints is designed, and a VGG19 network is adopted to extract features and calculate structure loss so as to ensure a detail recovery effect; by continuously optimizing the loss, the model obtained by training can effectively enhance underwater images in a test set, and the indexes such as the peak signal-to-noise ratio and the structural similarity are improved.

Description

Technical Field

[0001] The present invention relates to the field of image processing, and particularly to an underwater optical image enhancement method. Background Art

[0002] Underwater optical images are often used for underwater target guidance and confirmation at close range. Due to the different attenuation of light of each wavelength in water, serious distortion problems exist in underwater optical images, which affect subsequent image processing and target detection. This patent proposes an underwater optical image enhancement method to reduce the distortion degree of underwater optical images. Summary of the Invention

[0003] In view of this, the present invention provides an underwater optical image enhancement method, and the method includes the following steps:

[0004] Step S1: Obtain an underwater optical image data set;

[0005] Step S2: Construct an underwater optical image training set and a test set, and construct an underwater optical image preprocessing module;

[0006] Step S3: Construct an underwater optical image enhancement network, and the underwater optical image enhancement network includes a feature extraction module, a feature fusion module and an image restoration module;

[0007] Step S4: Construct a training loss function, set hyperparameters, and train the underwater optical image enhancement network;

[0008] Step S5: Test the underwater optical image enhancement network on the underwater optical image test set, count the results, if the results meet the index requirements, stop training, otherwise, modify the hyperparameters and retrain until the index requirements are met.

[0009] Particularly, the step S1 includes: the underwater optical image data set consists of an original image and its reference image.

[0010] Particularly, the step S2 includes: dividing the data set into a training set and a test set according to a preset ratio, and preprocessing the underwater optical image using the gamma transformation method to obtain a preprocessed underwater optical image.

[0011] Particularly, the step S3 includes: the feature extraction module includes a plurality of encoding modules, and the plurality of encoding modules include:

[0012] The preprocessed image passes through a first convolutional module with a stride of 1 and a size of (3×3), then through a second convolutional module with a stride of 1 and a size of (5×5), then through a third convolutional module with a size of (2×2), through a fourth convolutional module with a stride of 1 and a size of (3×3), then through a fifth convolutional module with a stride of 1 and a size of (5×5), and then through a multi-scale feature extraction module, which consists of a fifth convolutional module with a size of (3×3), a sixth convolutional module with a size of (5×5), and a residual module.

[0013] Specifically, in step S3, a residual module and an attention mechanism are introduced into the feature fusion module. The residual module can avoid gradient disappearance, transfer features to deeper network layers, and improve the learning ability of the network. The pixel attention mechanism weights each pixel point, and the channel attention can weight the channels, focusing on local important features, thereby improving the pertinence of feature extraction. The encoder structure of the feature fusion module is described as follows:

[0014] The upper layer output passes through two seventh convolutional modules with a stride of 1 and a size of (3×3), and then passes through a pixel point attention module and a channel attention module in parallel;

[0015] In the pixel point attention module, it passes through two eighth convolutional modules with a stride of 1 and a size of (1×1), and then multiplies with the input; [[ID=X]]

[0016] The channel attention module includes that the input is transformed to a size of (1×1×C) through a fully connected layer, and then multiplies with the output of the pixel point attention module.

[0017] Specifically, in step S3, the image restoration module includes: passing through a ninth convolutional module with a stride of 1 and a size of (3×3), then through an upsampling module, and finally through a tenth convolutional module with a stride of 1 and a size of (3×3), and the output image is restored to the size of the input image.

[0018] Specifically, step S4 includes: constructing a training loss function, which is the weighted sum of the RGB channel loss, the structure loss, and the smoothness loss;

[0019] The RGB channel loss is as follows:

[0020]

[0021] img pre represents the image output by the enhancement network, img gt represents the original unenhanced image, c represents the three RGB channels, c = 0, 1, 2 correspond to the R, G, B channels respectively, w c represents the weight of the c-th channel, and M and N respectively represent the height and width of the image;

[0022] The structural loss includes using the VGG19 network to extract the feature maps of the enhanced image and the original image. The VGG19 network includes multiple convolutional and pooling layers, and the feature map extracted by the 4th convolutional layer is used to obtain the feature map F4.

[0023] The feature map F4 and the original image are input into the VGG19 network with the fully connected layer removed, and the following structural loss is obtained:

[0024]

[0025] Loss total = η1Loss rgb + Loss struct

[0026] Among them, Loss struct represents the structural loss, and VGG19() represents the function of using the VGG19 network to extract image features; C represents the number of channels of the feature map, H represents the height of the feature map, and W represents the width of the feature map.

[0027] Loss total represents the total loss function, Loss rgb represents the RGB channel loss, which is used to constrain the color information; η1 represents the weight factor set for the RGB channel loss, and Loss struct represents the structural loss, which is used to constrain the structural or detail information.

[0028] Initialize the learning rate and batch size, and set the number of iterations to implement the training process of the underwater optical image enhancement network.

[0029] Specifically, step S5 includes: testing the trained underwater optical image enhancement network on the test set, and statistically calculating two indicators of peak signal-to-noise ratio and structural similarity. If the results meet the indicator requirements, stop training; otherwise, modify the hyperparameters and retrain until the indicator requirements are met.

[0030] Beneficial effects:

[0031] Through the solution of the present invention, the distortion degree of the underwater optical image can be reduced, and the clarity of the image can be enhanced.

[0032] Through the solution of the present invention, the attenuation difference of different channels of the underwater image during propagation can be compensated.

[0033] Through the solution of the present invention, the attention mechanism can be used to focus on local important features, and the pertinence of feature extraction can be improved.

[0034] Through the solution of the present invention, a matching pair dataset containing the original image and the reference image can be constructed, providing strong supervision for network training.

[0035] Through the solution of the present invention, a multi-constraint loss including color, structure, and details can be adopted to make the enhancement result comprehensively approximate the real effect.

[0036] Through the solution of the present invention, an encoder-decoder network structure can be cascaded to focus on global and local features in sequence.

[0037] Through the solution of the present invention, a test set evaluation process can be established to guide the network training direction and ensure the enhancement quality. BRIEF DESCRIPTION OF THE DRAWINGS

[0038] Figure 1 is a flowchart of the underwater optical image enhancement method proposed in the present invention;

[0039] Figure 2 is a schematic diagram of the feature extraction module of the optical image enhancement network proposed in the present invention;

[0040] Figure 3 is a schematic diagram of the feature fusion module of the optical image enhancement network proposed in the present invention;

[0041] Figure 4 is a schematic diagram of the image restoration module of the optical image enhancement network proposed in the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0042] The present invention will be described in detail below with reference to the accompanying drawings and by way of examples.

[0043] The present invention provides an underwater optical image enhancement method, the flowchart of which is as Figure 1 shown, and the method includes the following steps:

[0044] Step S1: Obtain an underwater optical image dataset; use the publicly available UFO-120 underwater dataset as the experimental dataset for this time. This dataset includes 1620 pairs of underwater optical images in different scenarios;

[0045] Step S2: Construct an underwater optical image training set and a test set, and construct an underwater optical image preprocessing module; pair the processed original images and reference images, and divide the training set and the test set according to a certain ratio. Repeat the above steps until enough image pairs with rich scenarios are obtained.

[0046] Divide the dataset into a training set and a test set according to a preset ratio. If the preset ratio is set to 3:1, use the gamma transformation method to preprocess the underwater optical image. The steps of using gamma transformation to preprocess the underwater optical image are as follows: Read out the RGB three channels of the original image separately. Perform gamma transformation on each channel individually: I_new = I_old^gamma; where I_old is the original pixel intensity, I_new is the adjusted intensity, and gamma is the gamma coefficient.

[0047] Repeat the gamma transformation for each channel, using different gamma values. Usually, the gamma value for the R channel is taken as a value less than 1, and the G and B channels are taken as values greater than 1. Recombine the processed three channels into one image. In this way, the gamma transformation can compensate for the attenuation differences of different channels underwater, serving as a simple and effective color correction preprocessing. Filtering and other operations can also be added after the gamma transformation to further reduce noise.

[0048] Step S3: Construct an underwater optical image enhancement network, which includes a feature extraction module, a feature fusion module, and an image restoration module; in step S3, the feature extraction module is as Figure 2 shown, which includes multiple encoding modules, and the multiple encoding modules include:

[0049] The preprocessed image passes through a first convolutional module of (3×3) with a stride of 1, then passes through a second convolutional module of (5×5) with a stride of 1, then a third convolutional module of (2×2), passes through a fourth convolutional module of (3×3) with a stride of 1, then passes through a fifth convolutional module of (5×5) with a stride of 1, and then passes through a multi-scale feature extraction module, which is composed of a fifth convolutional module of (3×3), a sixth convolutional module of (5×5), and a residual module.

[0050] In step S3, the feature fusion module is as Figure 3 shown, which introduces a residual module and an attention mechanism. The residual module can avoid gradient disappearance, transfer features to deeper network layers, and improve the learning ability of the network. The pixel attention mechanism weights each pixel point, and the channel attention can weight the channels, focusing on local important features, thereby improving the pertinence of feature extraction; the encoder structure of the feature fusion module is described as follows:

[0051] The upper-layer output passes through two seventh convolutional modules of (3×3) with a stride of 1, and then passes through a pixel point attention module and a channel attention module in parallel;

[0052] In the pixel point attention module, it passes through two eighth convolutional modules of (1×1) with a stride of 1, and then multiplies with the input;

[0053] The channel attention module includes transforming the input to a size of (1×1×C) through a fully connected layer, and then multiplying it by the output of the pixel attention module.

[0054] In step S3, the image restoration module is as Figure 4 shown, which includes: passing through a ninth convolutional module with a stride of 1 and a size of (3×3), then passing through an upsampling module, and finally passing through a tenth convolutional module with a stride of 1 and a size of (3×3), and the output image is restored to the size of the input image.

[0055] Step S4: Construct a training loss function, set hyperparameters, and train the underwater optical image enhancement network; step S4 includes: constructing a training loss function, and the training loss function is weighted and implemented by RGB channel loss, structural loss, and smoothness loss;

[0056] The RGB channel loss is as follows:

[0057]

[0058] img pre represents the image output by the enhancement network, and img gt represents the original unenhanced image, c represents the three RGB channels, c = 0, 1, 2 correspond to the R, G, and B channels respectively, and w c represents the weight of the c-th channel, and M and N represent the height and width of the image respectively;

[0059] In the structural loss, the VGG19 network is used to extract the feature maps of the enhanced image and the original image. The VGG19 network includes multiple convolutional and pooling layers, and the feature map F4 is obtained using the feature map extracted by the 4th convolutional layer;

[0060] Input the feature map F4 and the original image into the VGG19 network without the fully connected layer, and the structural loss is obtained as follows:

[0061]

[0062] Loss total = η1Loss rgb + Loss struct

[0063] where Loss struct represents the structural loss, and VGG19() represents the function of using the VGG19 network to extract image features; C represents the number of channels of the feature map, H represents the height of the feature map, and W represents the width of the feature map;

[0064] Loss total represents the total loss function, and Lossrgb represents the RGB channel loss, which is used to constrain color information; η1 represents the weight factor set for the RGB channel loss, Loss struct represents the structure loss, which is used to constrain structural or detailed information;

[0065] Initialize the learning rate and batch, and set the number of iterations to implement the training process of the underwater optical image enhancement network.

[0066] Step S5: Test the underwater optical image enhancement network on the underwater optical image test set, and count the results. If the results meet the index requirements, stop training; otherwise, modify the hyperparameters and retrain until the index requirements are met. Specifically, it includes: testing the trained underwater optical image enhancement network on the test set, and counting two indicators, the peak signal-to-noise ratio and the structural similarity. If the results meet the index requirements, stop training; otherwise, modify the hyperparameters and retrain until the index requirements are met.

[0067] In summary, the above is only a preferred embodiment of the present invention and is not intended to limit the protection scope of the present invention. Any modifications, equivalent replacements, improvements, etc. made within the spirit and principle of the present invention shall be included within the protection scope of the present invention.

[0068] For those skilled in the art, it is obvious that the embodiments of the present invention are not limited to the details of the above exemplary embodiments, and can be implemented in other specific forms without departing from the spirit or basic characteristics of the embodiments of the present invention. Therefore, from any point of view, the embodiments should be regarded as exemplary and non-restrictive. The scope of the embodiments of the present invention is defined by the appended claims rather than the above description. Therefore, it is intended to cover all changes falling within the meaning and scope of the equivalent elements of the claims in the embodiments of the present invention. Any reference signs in the claims should not be regarded as limiting the claimed rights. In addition, it is obvious that the word "comprising" does not exclude other units or steps, and the singular does not exclude the plural. The multiple units, modules or devices stated in the system, apparatus or terminal claims can also be implemented by the same unit, module or device through software or hardware. First, second, etc. are used to represent names and do not represent any specific order.

[0069] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the embodiments of the present invention and not to limit them. Although the embodiments of the present invention have been described in detail with reference to the above preferred embodiments, those of ordinary skill in the art should understand that the technical solutions of the embodiments of the present invention can be modified or equivalently replaced without departing from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A method for underwater optical image enhancement, characterized in that: The method comprises the following steps: Step S1: Acquire underwater optical image dataset; Step S2: constructing underwater optical image training set and test set, and constructing underwater optical image preprocessing module; Step S3: constructing an underwater optical image enhancement network, wherein the underwater optical image enhancement network includes a feature extraction module, a feature fusion module and an image restoration module; Step S4: constructing a training loss function, setting hyperparameters, and training the underwater optical image enhancement network; Step S5: Testing the underwater optical image enhancement network on the underwater optical image test set, and collecting statistics on the results. If the results meet the index requirements, stop training; otherwise, modify the hyperparameters and retrain until the index requirements are met.

2. The underwater optical image enhancement method according to claim 1, characterized in that: Said step S1 includes: the underwater optical image data set is composed of an original image and a reference image thereof.

3. The underwater optical image enhancement method according to claim 1, characterized in that: The step S2 includes: dividing the data set into a training set and a test set according to a preset ratio, and preprocessing the underwater optical image using a gamma transform method to obtain a preprocessed underwater optical image.

4. The underwater optical image enhancement method according to claim 1, characterized in that: The step S3 includes: the feature extraction module includes multiple encoding modules, and the multiple encoding modules include: the preprocessed image passes through a first convolution module (3×3) with a step size of 1, and then passes through a second convolution module (5×5) with a step size of 1, and then a third convolution module (2×2), and then passes through a fourth convolution module (3×3) with a step size of 1, and then passes through a fifth convolution module (5×5) with a step size of 1, and then passes through a multi-scale feature extraction module, which consists of a fifth convolution module (3×3), a sixth convolution module (5×5) and a residual module.

5. The underwater optical image enhancement method according to claim 4, characterized in that: In step S3, the feature fusion module introduces a residual module and an attention mechanism. The residual module can avoid gradient vanishing, pass features to deeper network layers, and improve the learning ability of the network. The pixel attention mechanism weights each pixel point, and the channel attention mechanism weights the channel, which can focus on important local features, thereby improving the pertinence of feature extraction. The encoder structure of the feature fusion module is described as follows: The upper layer output passes through two (3×3) seventh convolution modules with a stride of 1, and then passes through a pixel attention module and a channel attention module in parallel; The pixel attention module passes through two eighth convolution modules with a stride of 1 (1×1) and then multiplies the input; The channel attention module includes an input that is fully connected to a size of (1×1×C) and then multiplied by the output of the pixel attention module.

6. The underwater optical image enhancement method according to claim 5, characterized in that: In step S3, the image restoration module includes: passing through a ninth convolution module (3×3) with a step size of 1, then passing through an upsampling module, and finally passing through a tenth convolution module (3×3) with a step size of 1, and the output image is restored to the size of the input image.

7. The underwater optical image enhancement method according to claim 1, characterized in that: The step S4 comprises: constructing a training loss function, wherein the training loss function is a weighted sum of RGB channel loss and structural loss; The RGB channel losses are as follows: img pre Represents the image output by the enhanced network, img gt represents the original unenhanced image, c is the channel index, c = 0, 1, 2 correspond to R, G, B channels respectively, w c represents the weight of the cth channel, M and N represent the height and width of the image respectively; The structural loss includes using a VGG19 network to extract feature maps of the enhanced image and the original image. The VGG19 network includes multiple convolutional and pooling layers, and the feature map F4 is obtained using the feature map extracted by the fourth convolutional layer. The feature map F4 and the original image are input into the VGG19 network with the fully connected layer removed, and the structural loss is as follows: Loss total =η1Loss rgb +Loss struct Among them, Loss struct Represents structural loss, VGG19() represents a function for extracting image features using the VGG19 network; C represents the number of channels of the feature map, H represents the height of the feature map, and W represents the width of the feature map; Loss total Represents the total loss function, Loss rgb Represents the RGB channel loss, which is used to constrain color information; η1 represents the weight factor set for the RGB channel loss, Loss struct Represents structural loss, used to constrain structure or detail information; Initializing the learning rate, batch, and number of iterations to implement the training process of the underwater optical image enhancement network.

8. The underwater optical image enhancement method according to claim 1, characterized in that: Step S5 includes testing the trained underwater optical image enhancement network on a test set, calculating two indicators, peak signal-to-noise ratio and structural similarity. If the result meets the indicator requirements, training is stopped; otherwise, hyperparameters are modified and training is repeated until the indicator requirements are met.

Images