Underwater image enhancement method based on prior guidance network

Abstract

The invention discloses an underwater image enhancement method based on a prior guide network. The underwater image enhancement method comprises the following steps: S1, establishing a double-flow architecture network model based on a Transform U-shaped encoder-decoder network, a physical perception prior encoder and a multi-scale feature enhancement module; s2, taking an original underwater image X and a prior image Xp enhanced based on physical perception as inputs of a U-shaped network and a prior encoder respectively, and performing iterative training on a double-flow architecture network model; carrying out back propagation on network parameters by taking loss function convergence stability as a training cut-off condition, and obtaining an ideal double-current architecture model; and S3, collecting an actual underwater image, inputting the actual underwater image into the ideal double-flow architecture model, and automatically outputting an enhanced image Y. By using the double-flow architecture network model of the method, the underwater image enhancement effect can be improved, the image details and the local contrast can be enhanced, and more accurate and natural colors can be recovered.

Description

Technical Field

[0001] The present invention relates to the field of computer vision technology, and in particular to an underwater image enhancement method based on a priori guided network. Background Art

[0002] With the development of marine resources, underwater imaging is crucial in many fields. However, the complex underwater environment and poor lighting conditions lead to problems in underwater images, such as blurred features, light scattering, color imbalance, and low contrast. Existing underwater image enhancement methods are divided into traditional methods and deep learning-based methods. Traditional methods often rely on visual priors and utilize techniques such as histogram transformation to improve image visual quality. However, these methods ignore the physical processes of underwater imaging, are prone to over- or under-enhancement, and have poor versatility. While physical model-based methods simulate the underwater imaging process, they are limited by the complexity of the scene and the uncertainty of the priors, making it difficult to simulate it realistically. Deep learning-based methods, such as those based on convolutional neural networks (CNNs) and generative adversarial networks (GANs), while effective, suffer from difficulties in capturing long-range dependencies and global contextual information, and suffer from poor generalization. Transformer-based methods can simulate global features, but the generation process is limited by the input, and the training parameters are numerous and computationally complex. Summary of the Invention

[0003] In order to solve the problems of insufficient feature extraction, difficulty in processing global information, susceptibility to noise and artifacts, and poor robustness in existing underwater image enhancement methods, this paper proposes an underwater image enhancement method based on a priori guided network, which can improve the underwater image enhancement effect, enhance image details and local contrast, and restore more accurate and natural colors.

[0004] The technical solution adopted in the present invention is:

[0005] The underwater image enhancement method based on the prior guided network includes the following steps:

[0006] S1, a dual-stream architecture network model is established based on the Transformer-based U-shaped encoder-decoder network, the physical perception prior encoder, and the multi-scale feature enhancement module;

[0007] S2, the original underwater image X and the prior image X enhanced based on physical perception p The two-stream architecture model is iteratively trained as the input of the U-shaped network and the prior encoder respectively. The network parameters are back-propagated with the loss function convergence and stability as the training cutoff condition to obtain the ideal two-stream architecture model.

[0008] S3. Collect actual underwater images and input them into the ideal dual-stream architecture model, which automatically outputs the enhanced image Y.

[0009] In the dual-flow architecture network model, the U-shaped encoder and the prior encoder both adopt a two-stage architecture to retain shallow and deep features; then the features extracted by the two encoders are spliced respectively to obtain F l and F h An input multi-scale feature enhancement module is used to fuse features by different scale partial convolution, output enhanced features F and fuse with the U-shaped network jump connection, as the new jump connection of the U-shaped network decoder; finally the enhanced image Y and the transmission graph T are output by the U-shaped network decoder of the Transformer.

[0010] In the U-shaped network, a window self-attention mechanism is used to construct a Transformer module to reduce the amount of calculation and capture global relationships.

[0011] Given a feature map First, project L to Q, K, V through a linear layer, then group according to the preset window size; the calculation method of self-attention is:

[0012]

[0013] Where Q=[q1,q2,...,q n ]∈R d×n is the query matrix, K=[k1,k2,...,k n ]∈R d×n is the key matrix, V=[v1,v2,...,v n ]∈R d×n is the value matrix, Softmax is a normalization operation on each row, is a scaling factor, and B is used to balance the magnitude of dot product scores.

[0014] An additional convolution is performed on V to obtain spatial aggregation, denoted as:

[0015] Q,K,V=Attention(Q,K,V)+Conv(V)

[0016] Where represents V without window division, and Conv(·) represents deep convolution.

[0017] In the U-shaped network, a feature fusion block is introduced to integrate multi-branch features.

[0018] The prior encoder first extracts features from the prior result using a convolution layer and a dense convolution with residual structure; then uses dense convolution to further extract features from the image; each dense convolution contains five residual blocks, and the residual block is described as:

[0019] f=Conv3×3 (ReLU(Conv 3×3 (X p )))

[0020] Out=αf(x)+x

[0021] where α is set to 0.1.

[0022] In the two-stage architecture, the extraction steps of the U-shaped encoder and the prior encoder are as follows:

[0023] Firstly, the following are extracted respectively:

[0024] where G Enc-u represents the feature extraction process of the U-shaped encoder, G Enc represents the feature extraction process of the prior encoder; i represents the features extracted from the original image, and j represents the features extracted from the prior result; i s , j s represent the shallow features extracted from the original image and the prior result respectively, i d , j d represent the deep features extracted from the original image and the prior result respectively.

[0025] Finally, the features of the same level from different encoders are spliced, denoted as: F l =Concat(j d ,i d ); F h =Concat(j s ,i s ).

[0026] The multi-scale feature enhancement module network captures the image structure from fine-grained texture to a larger range by combining the details and structural information through different scale partial convolution layers, enhances the contrast and clarity, and reduces the noise and interference in the image; the specific processing steps are as follows:

[0027] The input F l or F h is processed respectively to obtain:

[0028]

[0029] where respectively represent the results of image features obtained by different size convolution kernels, and x is the input F l or F h .

[0030] Then, the final F L and F H are obtained by using convolution layer splicing.

[0031] The enhanced feature F=(F L ,F H ) is supplemented by the skip connection for the decoder feature;

[0032] Finally, the enhanced image Y and the transmission graph T are obtained by the decoder: Y,T=G Dec (i n ,F).

[0033] The loss function adopts the average error loss L Avg , the light and dark perception loss L DBP and the structural similarity loss L ssim , and the total loss is: Loss=L Avg +L DBP +L ssim .

[0034] Compared with the prior art, the beneficial effects of the present application are:

[0035] 1、The present application integrates the advantages of the original network and the physical perception priori network based on the priori feature into the traditional framework, provides a new framework for data-driven and information-rich underwater image enhancement, and significantly improves the color contrast of underwater images through the multi-scale feature enhancement module, and provides a clearer and more natural underwater visual experience.

[0036] 2、Experiments show that the method of the present application is superior to the most advanced technology in both quantitative and qualitative aspects, and has excellent performance on paired data sets and is also competitive on unpaired data sets.

[0037] 3、The dual-flow architecture network model of the method of the present application can improve the underwater image enhancement effect, enhance the image details and local contrast, and restore more accurate and natural colors. BRIEF DESCRIPTION OF DRAWINGS

[0038] Figure 1 is a schematic diagram of the overall network structure.

[0039] Figure 2 is a schematic diagram of the window self-attention mechanism.

[0040] Figure 3 is a schematic diagram of the priori encoder structure.

[0041] Figure 4 is a network structure diagram of the multi-scale feature enhancement module.

[0042] Figure 5 is an effect data graph. DETAILED DESCRIPTION

[0043] In order to make the above objectives, characteristics and advantages of the present application more obvious and easy to understand, the specific implementation method of the present application is described in detail below. In the following description, a large number of specific details are set forth in order to facilitate a thorough understanding of the present application. However, the present application can be implemented in many other ways different from those described herein, and those skilled in the art can make similar improvements without departing from the spirit of the present application, so the present application is not limited to the specific implementation disclosed below.

[0044] Unless otherwise defined, all technical and scientific terms used herein have the same meaning as commonly understood by one of ordinary skill in the art to which the present application belongs. The terms used in the specification of the application herein are only for the purpose of describing specific embodiments and are not intended to limit the present application.

[0045] As Figure 1 is a schematic diagram of the overall network structure of the present application, showing the connection relationship and data flow of the Transformer-based U-shaped network, the prior encoder and the multi-scale feature enhancement module. The innovation of the present application lies in: the development of the scale feature enhancement module, realizing the fusion of features of different scales; and the improved application of the Transformer module. The points to be protected include the flow of the entire underwater image enhancement method, the specific design and mutual relationship of each module, and the combination method of the loss function.

[0046] The present application proposes a dual-flow architecture combining a Transformer-based U-shaped encoder-decoder network and a priori network based on physical perception. The original underwater image X and the priori image X enhanced based on physical perception are respectively taken as the input of the U-shaped network and the priori encoder. The U-shaped encoder and the priori encoder both adopt a two-stage architecture, retaining shallow and deep features. Then the features extracted by the two encoders are spliced respectively, and the obtained F p and F l are input into the multi-scale feature enhancement module. The multi-scale feature enhancement module fuses features through different scale partial convolutions, and the enhanced features F are fused with the U-shaped network jump connection, serving as a new jump connection of the decoder. Finally, the enhanced image Y and the transmission graph T are obtained by the decoder. h

[0047] In the U-shaped network, the Transformer module is constructed by using the window self-attention mechanism to reduce the computational amount and capture the global relationship; the feature fusion block is introduced to integrate multi-branch features. The priori encoder extracts features from the priori result by using convolution layers and dense convolution with residual structure. The loss function adopts a combination of mean error loss, light and dark perception loss and structural similarity loss for training. The specific module description will be expanded as follows:

[0048] As Figure 2 ​is a window self-attention mechanism schematic diagram, which shows the calculation process and structure of the mechanism. The application adopts a window self-attention mechanism to construct a Transformer module, given a feature map First, project L to Q, K, V through a linear layer, and then group them according to the preset window size. Therefore, the calculation method of self-attention is:

[0049]

[0050] where Q = [q1, q2,..., q n ] ∈ R d×n is the query matrix, K = [k1, k2,..., k n ] ∈ R d×n is the key matrix, V = [v1, v2,..., v n ] ∈ R d×n is the value matrix, Softmax is a normalization operation on each row, is a scaling factor, and B is used to balance the magnitude of dot product scores. Unlike using only attention mechanism to aggregate information within the window, the application performs additional convolution on V, so the spatial aggregation scheme of the application is called window self-attention mechanism, which can be represented as:

[0051] Q, K, V = Attention (Q, K, V) + Conv (V)

[0052] where represents V without window division, and Conv(·) represents deep convolution. Through the spatial aggregation method of the window self-attention mechanism, global relationships can be captured as much as possible under the premise of reducing the amount of calculation, and the information within the field can be supplemented through convolution.

[0053] As Figure 3 is a priori encoder structure schematic diagram, which embodies the connection and feature extraction process of the convolution layer and dense convolution in the priori encoder. In order to solve the problem that the images taken in the underwater environment are often affected by multiple degradation factors, including light absorption and scattering, interference of suspended particles and color distortion, etc. The application believes that integrating priori knowledge into the model is an effective strategy. In this study, the application uses the priori result of MLLE as auxiliary information. This priori result is particularly aimed at the common color cast and blurring problems in underwater images. By introducing color correction and sharpening processing, the quality of underwater images can be significantly improved, and the image blur caused by turbid water can be reduced. Therefore, by integrating these priori knowledge, the application believes that the model can more effectively adapt to various complex conditions underwater, restore more accurate and natural colors, and exhibit better enhancement effect.

[0054] To extract features from the prior results, the present application designs a prior encoder, as shown in Figure 3 The prior encoder uses convolutional layers and dense convolution with residual structure to extract features from the prior results, and then uses dense convolution to further extract features from the image. Each dense convolution contains five residual blocks. The residual block can be described as:

[0055] f = Conv 3×3 (ReLU(Conv 3×3 (X p )))

[0056] Out = a f(x) + x

[0057] where a is set to 0.1.

[0058] The architecture of this encoder is similar to the encoder architecture of the U-shaped network, which adopts a two-stage architecture to retain both shallow and deep features. The extraction process of the two decoders is summarized as follows:

[0059]

[0060] where G Enc-u represents the feature extraction process of the U-shaped encoder, and G Enc represents the feature extraction process of the prior encoder. i represents the features extracted from the original image, and j represents the features extracted from the prior results. i s , j s represent the shallow features extracted from the original image and the prior results, respectively, and i d , j d represent the deep features extracted from the original image and the prior results, respectively. Finally, the features of the same level from different encoders are spliced, which can be represented as:

[0061] F l = Concat(j d ,i d )

[0062] F h = Concat(j s ,i s )

[0063] where F l represents the spliced shallow features, and F h represents the spliced deep features.

[0064] As Figure 4, is a multi-scale feature enhancement module network structure diagram, which shows the process of extracting effective features at different scales. The application introduces a multi-scale enhancement network, which enhances the contrast and clarity by fusing feature images of different scales, and reduces noise and interference in the image. The module can capture image structures from fine-grained textures to larger ranges by combining detailed and structural information through different scale partial convolution layers. The original intention of partial convolution is to process images with occlusion or incomplete areas. In the underwater environment, due to the uneven propagation of light, the image may appear local area insufficient or overexposed. The network using partial convolution can adaptively model and compensate for these irregular illuminations. The three branch convolution kernel sizes are 3x3, 5x5 and 7x7 respectively. Such design allows the network to learn useful features in different areas, which allows the module to effectively highlight the tiny details of the image while preserving the integrity of larger structures. The specific process is as follows:

[0065]

[0066] Respectively represent the results of image features obtained by different size convolution kernels, x is the input F l Or F h , using convolution layer splicing to obtain the final F L And F H , wherein F L Indicates the enhanced shallow features of multi-scale, F H Indicates the enhanced deep features of multi-scale. The enhanced features F=(F L ,F H ) are used to supplement the decoder features through jump connection. Finally, the enhanced image Y and the transmission graph T are obtained through the decoder.

[0067] Y,T=G Dec (i n ,F)

[0068] Loss function part, the application uses average error loss, light perception loss and consistency loss for training. In order to train the underwater image enhancement model more comprehensively, ensure the overall improvement of image quality, and solve various challenges brought by underwater special conditions.

[0069] Average error loss

[0070] The average error loss function of the application is designed based on the concept of error norm minimization, which combines the robustness of L1 loss on data outliers and the sensitivity of L2 loss to various outliers. Its definition is:

[0071] L Avg =0.8L1+0.2L2

[0072] Where the L1 component accounts for 0.8 of the total weight, providing a more balanced gradient that helps improve generalization by avoiding overemphasis on outlier data points that can affect the L2 loss. On the other hand, the L2 component has a weight of 0.2, ensuring that the model gives sufficient attention to squared errors, promoting convergence by imposing a harsher penalty on completely wrong predictions, which cannot be done with L1 loss alone.

[0073] Lightness-aware loss

[0074] Based on the dark channel prior, which states that the darkest pixels in underwater natural images typically contain very low intensity values due to light absorption and scattering in underwater environments. The invention utilizes this property to design a dark channel prior loss function as part of the network optimization, promoting the network's ability to reduce color fading and enhance image clarity under underwater conditions. The loss function expression is:

[0075]

[0076] Where t and denote the dark channel prior of the original image and the estimated transmittance of the enhancement result, respectively. L represents a Laplacian-like matrix, and λ is a hyperparameter.

[0077] To overcome the problem of overly dark images that may result from using only the dark channel prior, the invention also adds a bright channel prior loss. This loss is based on the bright channel prior, which states that the brightest pixels in underwater images reflect the global lighting conditions. The loss function expression is:

[0078]

[0079] Where t and denote the transmittance estimates from the BCP and the network of the invention, respectively. By jointly using the dark channel and bright channel losses, the network can more comprehensively enhance underwater images and improve visual clarity while avoiding the problems of excessive darkening or overexposure.

[0080] Finally, a perceptual loss function is used to enhance details. This loss function uses a pre-trained VGG-16 network to extract feature representations at specific layers and calculates the mean squared error (MSE) between them. The loss function expression is:

[0081]

[0082] Where X is the original image, Y is the enhanced image, φ i (X) and φ i (Y) represent their feature representations at the i-th layer in the pre-trained neural network, and N is the number of feature layers. The final lightness-aware loss is:

[0083] L DBP =0.002L DCP +0.03L BCP +0.3L per

[0084] Structural similarity loss

[0085] Structural Similarity Loss Structural Similarity Index (SSIM) [44, 45] is a metric that quantifies the structural similarity between two images based on brightness, contrast, and structural information. The main purpose of incorporating this metric into the loss function is to encourage the network to generate enhanced images that are closer to the ground truth in these three aspects. The loss function is expressed as:

[0086]

[0087] where μ X and μ Y represents the mean of images X and Y, σ XY represents the covariance of images X and Y, and Represents the variance of images X and Y, c1 and c2 are two variables that maintain stability.

[0088] The total loss is written as:

[0089] Loss = L Avg +L DBP +L ssim

[0090] The overall process of the present invention is described as follows: First, the main encoder (Transformer U-type network) receives the original underwater image, and extracts multi-scale features through reflection-filled window segmentation and window self-attention modules. The prior encoder receives the prior image generated by MMEL, and extracts shallow and deep features through residual convolution blocks. The features (shallow and deep) of the two encoders at the same level are then spliced ​​to form fused features. The fused shallow and deep features are processed with 3×3, 5×5, and 7×7 convolution kernels, respectively, to extract local information of different scales. The multi-scale information is spliced ​​along the channel dimension, and the channel is compressed by the convolution layer to generate enhanced features. The enhanced features are then input into the decoder as jump connections to supplement the detail information. The decoder is based on a U-type structure, gradually upsampling and combining the multi-scale features of the jump connection. The final output is the enhanced image J and the transmission map T.

[0091] like Figure 5 As shown in the figure, it is an effect data diagram. It can be seen that our method has achieved pleasing results on the two public datasets UIEB and EUVP, and the enhanced images are close to the ideal images.

[0092] The above describes the preferred embodiments of the present application. It should be noted that, for those skilled in the art, several improvements and refinements can be made without departing from the principles of the present application, and these improvements and refinements should be considered as falling within the scope of the present application.

Claims

1. An underwater image enhancement method based on a priori guided network, characterized in that: The following steps are involved: S1, a dual-stream architecture network model is established based on the Transformer-based U-shaped encoder-decoder network, the physical perception prior encoder, and the multi-scale feature enhancement module; S2, the original underwater image X and the prior image X enhanced based on physical perception p The two-stream architecture model is iteratively trained as the input of the U-shaped network and the prior encoder respectively. The network parameters are back-propagated with the loss function convergence and stability as the training cutoff condition to obtain the ideal two-stream architecture model. S3. Collect actual underwater images and input them into the ideal dual-stream architecture model, which automatically outputs the enhanced image Y.

2. The underwater image enhancement method based on a priori guided network according to claim 1, characterized in that: In the dual-stream architecture network model, both the U-shaped encoder and the prior encoder adopt a two-stage architecture to retain shallow and deep features; the features extracted by the two encoders are then concatenated to obtain the F l With F h Input a multi-scale feature enhancement module; the multi-scale feature enhancement module is used to fuse features through partial convolutions of different scales, output the enhanced features F and fuse them with the U-type network jump connection, which serves as a new jump connection for the decoder of the U-type network; finally, the decoder of the Transformer's U-type network outputs the enhanced image Y and the transmission map T.

3. The underwater image enhancement method based on a priori guided network according to claim 1, characterized in that: In the U-shaped network, a window self-attention mechanism is used to construct a Transformer module to reduce the amount of computation and capture global relationships; Given a feature map First, L is projected to Q, K, V through a linear layer, and then grouped according to the preset window size; the calculation method for self-attention is: where Q = [q1,q2,…,q n ]∈R d×n is the query matrix, K=[k1,k2,...,k n ]∈R d×n is the key matrix, V=[v1,v2,…,v n ]∈R d×n Is a value matrix, Softmax is a normalization operation for each row, is a scaling factor, B, used to balance the magnitude of the dot product score; Performing additional convolution on V results in spatial aggregation, expressed as: Q,K,V=Attention(Q,K,V)+Conv(V) in represents V without window partitioning, and Conv(·) represents depthwise convolution.

4. The underwater image enhancement method based on a priori guided network according to claim 1, characterized in that: In the U-shaped network, a feature fusion block is introduced to integrate multi-branch features.

5. The underwater image enhancement method based on a priori guided network according to claim 1, characterized in that: The prior encoder first extracts features from the prior results using convolutional layers and dense convolution with a residual structure; and then further extracts features from the image using dense convolution; Each dense convolution contains five residual blocks, which are described as: f=Conv 3×3 (ReLU(Conv 3×3 (X p ))) Out=αf(x)+x where α is set to 0.

1.

6. The underwater image enhancement method based on a priori guided network according to any one of claims 1, 4, and 5, characterized in that: In the two-stage architecture, the extraction steps of the U-shaped encoder and the a priori encoder are as follows: First extract separately: Among them G Enc-u represents the feature extraction process of the U-type encoder, G Enc Represents the feature extraction process of the prior encoder; i represents the features extracted from the original image, and j represents the features extracted from the prior result; i s 、j s Represent the shallow features extracted from the original image and the prior results, i d 、j d Represent the deep features extracted from the original image and the prior results respectively; Finally, the features of the same level of different encoders are concatenated and expressed as: F l =Concat(j d ,i d );F h =Concat(j s ,i s ).

7. The underwater image enhancement method based on a priori guided network according to claim 6, characterized in that: The multi-scale feature enhancement module network combines details and structural information through partial convolutional layers of different scales to capture image structures from fine-grained textures to larger scales, enhance contrast and clarity, and reduce noise and interference in the image. The specific processing steps are as follows: For input F l or F h Process them separately to get: F L =(Concat(F l 3×3 ,F l 5×5 ,F l 7×7 )) in They represent the results of image features obtained by convolution kernels of different sizes, and x is the input F l or F h ; Then, convolutional layers are used to concatenate the layers to get the final F L and F H ; For the enhanced feature F=(F L ,F H ) is used to supplement the decoder features through skip connections; Finally, the decoder obtains the enhanced image Y and the transmission map T: Y, T = G Dec (i n ,F).

8. The underwater image enhancement method based on a priori guided network according to claim 1, characterized in that: The loss function adopts the average error loss L Avg , light and dark perception loss L DBP and structural similarity loss L ssim Combined, the total loss is: Loss = L Avg +L DBP +L ssim .

Images