Underwater image enhancement method based on conditional convolution and channel attention mechanism

Abstract

The invention discloses an underwater image enhancement method based on conditional convolution and a channel attention mechanism, and relates to the technical field of underwater image enhancement, and an underwater image enhancement model used in the method comprises an encoder, a decoder, a bottleneck layer and jump connection. The encoder, the decoder and the bottleneck layer all adopt a residual SE condition convolution module to extract feature information, the model utilizes a convolution down-sampling layer and a deconvolution up-sampling layer to realize down-sampling and up-sampling of a feature map respectively, and the jump connection comprises a splicing operation and a condition convolution layer. The condition convolution layer of the jump connection is used for feature dimension reduction and feature fusion of the spliced feature pattern between the codec and the decoder. According to the method, the underwater input image can be enhanced, and the underwater enhanced result image with high image quality can be obtained.

Description

Technical Field

[0001] The present invention relates to the technical field of underwater image enhancement, and in particular to an underwater image enhancement method based on conditional convolution and channel attention mechanism. Background Art

[0002] At present, the global land resources are tight and the population is growing rapidly, while the underwater resources are rich and diverse. Therefore, the research and development of underwater resources is rising worldwide, and the acquisition and processing of image information has become the first step in the development of marine resources. In this context, underwater image processing technology has received more and more attention. Underwater image processing technology plays an important role in improving the quality of underwater images and provides important help for people to perform subsequent underwater visual tasks.

[0003] The processing effect of existing underwater image processing technology on underwater images is still unsatisfactory, and the obtained underwater images still have some visual quality problems, such as image blur, color distortion, artifacts, etc.

[0004] At the same time, the traditional static convolution is a set of unchanging convolution kernels for underwater input images, which cannot effectively and dynamically cope with the underwater degradation with diverse and complex characteristics; neural networks usually cannot select feature map channels with important information. Therefore, the neural network method needs to be further improved in these aspects. Summary of the invention

[0005] In view of this, the purpose of the present invention is to propose an underwater image enhancement method based on conditional convolution and channel attention mechanism, so as to obtain underwater images with high image quality. The method aims to solve the problem that the traditional static convolution cannot dynamically process underwater input images due to a fixed set of convolution kernels, the problem that the neural network cannot screen important channel information, and the technical problems of low image quality of the processing results existing in the existing underwater image processing technology.

[0006] To this end, the present invention provides the following technical solutions:

[0007] The present invention provides an underwater image enhancement method based on conditional convolution and channel attention mechanism, comprising the following steps:

[0008] Step 1: Obtain a distorted underwater image and its corresponding underwater reference image to form a training set and a test set;

[0009] Step 2, construct an underwater image enhancement model based on conditional convolution and channel attention mechanism; the underwater image enhancement model takes the distorted underwater image as input and outputs the enhanced high-quality underwater image, and the underwater image enhancement model includes an encoder, a decoder, a bottleneck layer and a jump connection for fusing two feature maps of the same scale between the encoder and the decoder; the encoder includes a residual SE conditional convolution module and a convolution downsampling layer, and the convolution downsampling layer is used to achieve downsampling of the feature map; the decoder includes a residual SE conditional convolution module and a deconvolution upsampling layer, and the deconvolution upsampling layer is used to achieve upsampling of the feature map; the jump connection includes a splicing operation and a conditional convolution layer, and the conditional convolution layer is used to achieve feature dimension reduction and feature fusion of the spliced ​​feature map between the encoder and the decoder; the bottleneck layer includes a residual SE conditional convolution module; the residual SE conditional convolution module in the encoder, the decoder and the bottleneck layer is used as a basic processing unit to process feature information; the residual SE conditional convolution module includes a conditional squeeze-excitation module based on the channel attention mechanism, a conditional convolution layer, a ReLU function and a residual connection;

[0010] Step 3: input the training set images into the constructed underwater image enhancement model for training;

[0011] Step 4: input the test set image into the trained underwater image enhancement model, and output the underwater enhanced result image after model processing.

[0012] Advantages and positive effects of the present invention:

[0013] Firstly, the conditional convolution layer used generates a different set of convolution kernels for underwater images with different degradation conditions, thereby realizing dynamic processing of underwater images and being able to adapt to underwater degradation conditions with diverse and complex characteristics; secondly, the conditional squeeze and excitation (SE) module based on the channel attention mechanism used can emphasize the information of important channels, suppress the information of redundant channels, and has the channel-level screening capability of feature maps; finally, the conditional convolution module, the conditional squeeze and excitation module and the residual connection are integrated to form a residual SE conditional convolution module; the residual SE conditional convolution module is used as the basic processing unit of the method, so that the method can effectively process feature information and finally obtain underwater images with good image quality; the underwater image enhancement technology in the present invention will also promote the development of subsequent underwater visual tasks, such as underwater application fields such as underwater archaeology, underwater rescue, and underwater resource exploration. BRIEF DESCRIPTION OF THE DRAWINGS

[0014] Figure 1 is a flow chart of an underwater image enhancement method based on conditional convolution and channel attention mechanism in an embodiment of the present invention;

[0015] Figure 2 is a model structure diagram of an underwater image enhancement method based on conditional convolution and channel attention mechanism in an embodiment of the present invention;

[0016] Figure 3 is a schematic diagram of a jump connection structure in an embodiment of the present invention;

[0017] Figure 4 is a schematic diagram of the structure of a conditional convolutional layer in an embodiment of the present invention;

[0018] Figure 5 is a schematic diagram of the structure of a conditional squeeze-excitation module in an embodiment of the present invention;

[0019] Figure 6 Schematic diagram of the structure of the residual SE conditional convolution module in an embodiment of the present invention. DETAILED DESCRIPTION

[0020] Aiming at the color cast and blur phenomena in underwater images, this paper proposes an underwater image enhancement method based on conditional convolution and channel attention mechanism, which can be applied to underwater archaeology, underwater target detection and recognition, underwater resource exploration and other tasks.

[0021] like Figure 1 As shown, an underwater image enhancement method based on conditional convolution and channel attention mechanism in an embodiment of the present invention includes the following specific steps:

[0022] Step 1: Obtain a distorted underwater image and its corresponding underwater reference image to form a training set and a test set;

[0023] Step 2: Construct an underwater image enhancement model based on conditional convolution and channel attention mechanism; Figure 2 As shown in , the network model includes an encoder, a decoder and a bottleneck layer, as well as a jump connection used to connect the encoder and decoder feature maps; the encoder includes a residual SE conditional convolution module and a kernel size of , a convolutional downsampling layer with a stride of 2; the decoder includes a residual SE conditional convolution module and a kernel size of , a deconvolution upsampling layer with a step size of 2; the bottleneck layer includes a residual SE conditional convolution module; the jump connection includes a splicing operation and a kernel size of , a conditional convolution layer with a step size of 1; the residual SE conditional convolution module includes a conditional squeeze and excitation (SE) module based on the channel attention mechanism and a kernel size of , conditional convolutional layer with a step size of 1, ReLU function and residual connection;

[0024] Specifically, the network structure of the underwater image enhancement model includes:

[0025] The overall structure of the underwater image enhancement model is divided into 5 layers, including encoder, bottleneck layer, decoder and skip connection;

[0026] The encoder includes four downsampling stages, each of which includes a kernel size of , a convolutional downsampling layer with a stride of 2 and 2 residual SE conditional convolution modules;

[0027] The decoder includes four upsampling stages, each of which includes a kernel size of , a deconvolution upsampling layer with a stride of 2 and a residual SE conditional convolution module;

[0028] The bottleneck layer includes two residual SE conditional convolution modules;

[0029] There are 4 jump connections in total, each of which includes a concatenation operation and a kernel size of , a conditional convolutional layer with a stride of 1, and the skip connection is used to fuse the decoder upsampled feature map with the feature map from the encoder at the same scale. Figure 3 As shown, the skip connection includes a concatenation operation and a kernel size of , a conditional convolution layer with a step size of 1; in the skip connection, the feature map from the encoder is first concatenated with the upsampled feature map of the same scale from the decoder to obtain a concatenated feature map, and then a kernel size of , a conditional convolutional layer with a step size of 1 performs feature fusion and feature dimension reduction on the concatenated feature map to obtain a fused feature map with the same feature dimension as the upsampled feature map;

[0030] The conditional convolution layer includes a routing function and multiple experts (sub-convolution kernels), such as Figure 4 As shown in the figure, the conditional convolution layer can generate a corresponding convolution kernel according to each specific underwater input image, avoiding the problem that the traditional static convolution kernel treats all samples equally, and effectively improving the dynamic processing capability of the network;

[0031] The routing function is composed of global average pooling, Sigmoid activation function and learning routing weight matrix; the specific calculation formula of the routing function is as follows:

[0032] ;

[0033] in, is the input data, represents global average pooling, is the activation function, is a learnable weight matrix;

[0034] In the conditional convolution layer, multiple weights generated by the routing function are used to weightedly fuse multiple sub-convolution kernels (experts) to obtain the final convolution kernel, which is then convolved with the input data to obtain the output; the convolution kernel of the conditional convolution layer is parameterized as:

[0035] ;

[0036] in, is the convolution operation, Represents the routing function The weight size corresponding to the generated expert (sub-convolution kernel), For experts;

[0037] like Figure 5 As shown, the conditional squeeze and excitation (SE) module based on the channel attention mechanism includes a squeeze operation, an excitation operation and a channel vector multiplication; the conditional squeeze and excitation module can filter out the information of important channels at the channel level and suppress the information of redundant channels, thereby effectively increasing the representation ability of the network; the squeeze operation of the conditional squeeze and excitation module can be expressed as:

[0038] ;

[0039] in, is the channel of the feature map, For input data, and is the height and width of the feature map channel, is the global average of the input data;

[0040] The excitation operation is used to generate the channel-level weight, which is calculated according to the following formula:

[0041] ;

[0042] in, is the ReLU function, represents the Sigmoid function, and There are two kernels of size , a conditional convolutional layer with a step size of 1, setting Compression ratio is 16, that is The number of output channels is ;

[0043] Finally, the channel vector is multiplied using the weights For initial input Perform weight correction:

[0044] ;

[0045] in, is the multiplication operation between channels and vectors, Represents the output result of the conditional squeeze-excitation module;

[0046] The residual SE conditional convolution module includes a conditional squeeze-excitation module, two kernel sizes of , a conditional convolutional layer with a stride of 1, a ReLU function, and a residual connection, such as Figure 6 As shown; in the residual SE conditional convolution module, the input data After a kernel size of , a conditional convolutional layer with a stride of 1, a ReLU function, and a kernel size of , the conditional convolution layer with a step size of 1 and the conditional squeeze-excitation module are processed to obtain the result , Through the identity path and the initial input Add together to get the final output result;

[0047] Step 3: Input the training set images into the constructed underwater image enhancement model for training;

[0048] In terms of network loss, the L1 loss function and the perceptual loss function are used in the embodiment of the present invention:

[0049] The L1 loss function is: ;

[0050] in, , are respectively a clear underwater reference image and an underwater output image of the underwater image enhancement model, is the total number of pixels in the image, is the pixel index of the image, is the total number of color channels of the image;

[0051] The perceptual loss function is: ;

[0052] in, For the VGG network layer, , and They are the number of channels, height, and width of the feature map respectively;

[0053] The total loss function is: ;

[0054] in, is the weight balance coefficient, and the underwater image enhancement model is set to .

[0055] Step 4: Input the test set images into the trained underwater image enhancement model, and output high-quality underwater images after image enhancement;

[0056] The underwater image enhancement model processing image in step 2 specifically includes the following steps:

[0057] Step 2.1, get the input image I, pass the input image through a kernel size of 1 , the conditional convolution layer with a stride of 1 is projected into a feature map of any dimension;

[0058] Step 2.2, the feature map of any dimension is then input into the first layer of the encoder, which contains two residual SE conditional convolution modules and one kernel size , a convolution downsampling layer with a step size of 2; the residual SE conditional convolution module is used to learn feature information to obtain the first layer processing feature map of the encoder , the convolutional downsampling layer further performs the feature map Perform downsampling and feature dimension increase to obtain the downsampled feature map ;

[0059] Step 2.3: Repeat step 2.2 three times to obtain the second layer processing feature map of the encoder. , the third layer processes the feature map And the fourth layer processes the feature map , and the downsampled feature maps of each layer , and ;

[0060] Step 2.4: Set the feature map of the minimum spatial scale The input is sent to the bottleneck layer, which includes two residual SE conditional convolution modules; the bottleneck layer further learns the feature information to obtain the feature map ;

[0061] Step 2.5: Feature map Expand the feature map to get the upsampled map , and the same scale feature map from the encoder Splicing, get a spliced ​​feature map, and then use a kernel size of 1 , the conditional convolution layer with a step size of 1 processes the spliced ​​feature map to obtain the fused feature map ; A residual SE conditional convolution module in the fourth layer of the decoder extracts feature maps The feature information of the fourth layer is obtained ;

[0062] Step 2.6: Feature map Expand the feature map to get the upsampled map , and the same scale feature map from the encoder Splicing, get a spliced ​​feature map, and then use a kernel size of 1 , the conditional convolution layer with a step size of 1 processes the spliced ​​feature map to obtain the fused feature map ; A residual SE conditional convolution module in the third layer of the decoder extracts feature maps The feature information of the third layer is obtained ;

[0063] Step 2.7, step 2.6 are repeated twice to obtain the second layer processing feature map of the decoder , the first layer processes the feature map , the fused feature map , the fused feature map , upsampled feature map And upsampled feature maps ;

[0064] Step 2.8: Process feature map Through a kernel size of 1 , after convolution processing by the conditional convolution layer with a step size of 1, the final output image is obtained.

[0065] Finally, it should be understood that the above description of the preferred embodiment is relatively detailed, but it cannot be regarded as a limitation on the scope of patent protection of the present invention. Under the enlightenment of the present invention, ordinary technicians in this field can also make substitutions or modifications without departing from the scope of protection of the claims of the present invention, all of which fall within the scope of protection of the present invention. The scope of protection requested for the present invention shall be based on the attached claims.

Claims

1. An underwater image enhancement method based on conditional convolution and channel attention mechanism, characterized in that: The following steps are involved: Step 1: Obtain a distorted underwater image and its corresponding underwater reference image to form a training set and a test set; Step 2, construct an underwater image enhancement model based on conditional convolution and channel attention mechanism; the underwater image enhancement model takes the distorted underwater image as input and outputs the enhanced high-quality underwater image, and the underwater image enhancement model includes an encoder, a decoder, a bottleneck layer and a jump connection for fusing two feature maps of the same scale between the encoder and decoder; the encoder includes a residual SE conditional convolution module and a convolution downsampling layer, and the convolution downsampling layer is used to achieve downsampling of the feature map; the decoder includes a residual SE conditional convolution module and a deconvolution upsampling layer, and the deconvolution upsampling layer is used to achieve upsampling of the feature map; the jump connection includes a splicing operation and a conditional convolution layer, and the conditional convolution layer is used to achieve feature dimension reduction and feature fusion of the spliced ​​feature map between the encoder and decoder; the bottleneck layer includes a residual SE conditional convolution module; the encoder, the decoder and the residual SE conditional convolution module in the bottleneck layer are all used as basic processing units to process feature information; The residual SE conditional convolution module includes a conditional squeeze-excitation module based on a channel attention mechanism, a conditional convolution layer, a ReLU function and a residual connection; Step 3: input the training set images into the constructed underwater image enhancement model for training; Step 4: input the test set image into the trained underwater image enhancement model, and output the underwater enhanced result image after model processing.

2. According to claim 1, the underwater image enhancement method based on conditional convolution and channel attention mechanism is characterized in that: The network structure of the underwater image enhancement model specifically includes: The overall structure of the underwater image enhancement model is divided into 5 layers, including encoder, bottleneck layer, decoder and skip connection; The encoder includes four downsampling stages, each of which includes a kernel size of , a convolutional downsampling layer with a stride of 2 and 2 residual SE conditional convolution modules; The decoder includes four upsampling stages, each of which includes a kernel size of , a deconvolution upsampling layer with a stride of 2 and a residual SE conditional convolution module; The bottleneck layer includes 2 residual SE conditional convolution modules; There are 4 jump connections in total, each of which includes a concatenation operation and a kernel size of , a conditional convolutional layer with a stride of 1, and the skip connection is used to fuse the upsampled feature map of the decoder with the feature map from the encoder at the same scale.

3. The underwater image enhancement method based on conditional convolution and channel attention mechanism according to claim 2, characterized in that: In the conditional convolution layer, the input data is first used to generate the weights of each expert through the routing function, and then the group of weights is used to perform weighted fusion on all experts to obtain the final convolution kernel, which is then convolved with the input data to obtain the result.

4. The underwater image enhancement method based on conditional convolution and channel attention mechanism according to claim 3 is characterized in that: In the conditional squeeze-excitation module based on the channel attention mechanism, the input feature map first obtains global statistics through a global average pooling operation, and then undergoes a compression ratio of 16 and a step size of 1. Conditional convolutional layer, a ReLU function and a stride of 1 The conditional convolutional layer maps a set of weights for each channel, and then uses this set of weights to correct the input feature map to obtain the final output result of the conditional squeeze-excitation module.

5. The underwater image enhancement method based on conditional convolution and channel attention mechanism according to claim 4, characterized in that: In the residual SE conditional convolution module, the input feature map first passes through a kernel of size 1. , a conditional convolutional layer with a stride of 1, a ReLU function, and a kernel size of 1 , a conditional convolution layer with a step size of 1 and a conditional squeeze-excitation module are calculated to obtain an intermediate result, and then the intermediate result is added to the initial input feature map through a residual connection to obtain the final output result of the residual SE conditional convolution module.

6. The underwater image enhancement method based on conditional convolution and channel attention mechanism according to claim 5, characterized in that: The loss functions used in the underwater image enhancement model during training include L1 loss function and perceptual loss function; The L1 loss function is: ; in, , are respectively a clear underwater reference image and an underwater output image of the underwater image enhancement model, is the total number of pixels in the image, is the pixel index of the image, is the total number of color channels of the image; The perceptual loss function is: ; in, For the VGG network layer, , and They are the number of channels, height, and width of the feature map respectively; The total loss function is: ; in, is the weight balance coefficient, and the underwater image enhancement model is set to .

7. The underwater image enhancement method based on conditional convolution and channel attention mechanism according to claim 2, characterized in that: The underwater image enhancement model processing image in step 2 specifically includes the following steps: Step 2.1, get the input image I, pass the input image through a kernel size of 1 , the conditional convolution layer with a stride of 1 is projected into a feature map of any dimension; Step 2.2, the feature map of any dimension is then input into the first layer of the encoder, which contains two residual SE conditional convolution modules and one kernel size , a convolution downsampling layer with a step size of 2; the residual SE conditional convolution module is used to learn feature information to obtain the first layer processing feature map of the encoder , the convolutional downsampling layer further performs the feature map Perform downsampling and feature dimension increase to obtain the downsampled feature map ; Step 2.3: Repeat step 2.2 three times to obtain the second layer processing feature map of the encoder. , the third layer processes the feature map And the fourth layer processes the feature map , and the downsampled feature maps of each layer , and ; Step 2.4: Set the feature map of the minimum spatial scale The input is sent to the bottleneck layer, which includes two residual SE conditional convolution modules; the bottleneck layer further learns the feature information to obtain the feature map ; Step 2.5: Feature map Expand the feature map to get the upsampled map , and the same scale feature map from the encoder Splicing, get a spliced ​​feature map, and then use a kernel size of 1 , the conditional convolution layer with a step size of 1 processes the spliced ​​feature map to obtain the fused feature map ; A residual SE conditional convolution module in the fourth layer of the decoder extracts feature maps The feature information of the fourth layer is obtained ; Step 2.6: Feature map Expand the feature map to get the upsampled map , and the same scale feature map from the encoder Splicing, get a spliced ​​feature map, and then use a kernel size of 1 , the conditional convolution layer with a step size of 1 processes the spliced ​​feature map to obtain the fused feature map ; A residual SE conditional convolution module in the third layer of the decoder extracts feature maps The feature information of the third layer is obtained ; Step 2.7, step 2.6 are repeated twice to obtain the second layer processing feature map of the decoder , the first layer processes the feature map , the fused feature map , the fused feature map , upsampled feature map And upsampled feature maps ; Step 2.8: Process feature map Through a kernel size of 1 , after convolution processing by the conditional convolution layer with a step size of 1, the final output image is obtained.

Images