An underwater image enhancement method and system based on underwater light field Transformer

Abstract

The application provides an underwater image enhancement method based on an underwater light field Transformer, comprising the following steps: acquiring real paired underwater images and dividing an underwater image dataset, building a layered encoder-decoder enhancement network based on the underwater light field Transformer, constructing an L1 loss function and an SSIM loss function combination constraint mode to train the layered encoder-decoder enhancement network based on the underwater light field Transformer, obtaining a trained underwater image enhancement network, inputting a to-be-processed underwater image into the trained underwater image enhancement network, and realizing enhancement of the to-be-processed underwater image. The underwater light field information is used to eliminate the influence of underwater noise on the feature aggregation of the Transformer, the parallel attention mechanism is introduced to improve the local information interaction capability of the Transformer and improve the generalization capability of the model, the clarity and contrast of the original underwater image are greatly improved, and the method can play an important role in underwater image processing, underwater target detection and underwater visual tasks.

Description

Technical Field

[0001] This invention relates to the field of underwater image enhancement technology, and more particularly to an underwater image enhancement method and system based on an underwater light field Transformer. Background Technology

[0002] Underwater imaging presents a number of challenges due to the complex underwater environment. 1. Suspended particles and dissolved substances in the water cause light scattering and absorption, leading to blurring, reduced contrast, and color shifts. 2. Water refraction and scattering effects result in image blurring and loss of detail. Edges and textures become blurred, reducing image sharpness. 3. Various types of noise exist in the underwater environment, including ripple noise, atmospheric noise, and sensor noise, all of which degrade image quality. Marine resources are crucial for human livelihoods and the economy; however, unreliable underwater images limit the accuracy and efficiency of deep-sea research and resource exploration. Therefore, improving the quality of underwater images is essential for marine resource exploration. Summary of the Invention

[0003] To address the technical problems mentioned in the background section, this invention provides an underwater image restoration method based on wavelength compensation. This invention introduces underwater light field information to guide the Transformer algorithm in complex underwater environments, aiming to utilize only key windows for feature learning representation. This guides the model to model specific regions of the underwater image, thereby eliminating the impact of underwater noise on model performance. Furthermore, to address the insufficient local information interaction capability of the Transformer, based on the characteristics of convolutional local perception and parameter sharing, a parallel attention mechanism is used to fuse Transformer features as local information supplementation. Finally, the method is trained on a real-world underwater dataset to enhance the underwater image.

[0004] The technical means employed in this invention are as follows:

[0005] An underwater image enhancement method based on underwater light field Transformer includes the following steps:

[0006] Step 01: Obtain a real-pair underwater image dataset and divide the underwater image dataset into a training set, a validation set, and a test set;

[0007] Step 02: Construct a hierarchical encoder-decoder enhancement network based on underwater light field Transformer;

[0008] Step 03: Construct a combination of L1 loss function and SSIM loss function constraints to iteratively train the hierarchical encoder-decoder enhancement network based on underwater light field Transformer to obtain the trained underwater image enhancement network;

[0009] Step 04: Input the original underwater image to be enhanced into the trained underwater image enhancement network to achieve the enhancement processing of the original underwater image.

[0010] Furthermore, the specific steps for establishing the underwater image dataset in step S01 are as follows:

[0011] Based on the UIEB dataset, 800 images were randomly selected as the training set, 90 images as the validation set, and 60 images as the test set.

[0012] Furthermore, step S02, which describes constructing a hierarchical encoder-decoder network based on underwater light field Transformer, includes the following steps:

[0013] Step 21: Obtain the image input from the network and extract the underwater light field feature map from the input image;

[0014] Step 22: For the underwater light field feature map, first perform parameter weighting processing, then select k windows according to the parameter weights and generate an underwater light field mask;

[0015] Step 23: For the underwater light field mask, reshape the bonds and values ​​in the Transformer module, calculate self-attention, and fuse parallel convolutional attention to obtain the final features;

[0016] Step 24: Stack the Transformer modules and build a hierarchical encoder-decoder structure based on the U-Net framework;

[0017] Step 25: Construct an adaptive Transformer network structure to complete the construction of a hierarchical encoder-decoder enhancement network based on underwater light field.

[0018] Furthermore, step S21, which involves extracting an underwater light field feature map from the input image, includes the following steps:

[0019] Step 211: Extract underwater light field information and use a low-pass Gaussian filter to preserve the light field information of the underwater image;

[0020] Step 212: Extract the underwater light field feature map G from the underwater light field map using a 3×3 convolution and activation function. The formula is as follows:

[0021] G = Conv(Sigmoid((Gauss) σ (X))));

[0022] Where G represents the underwater light field feature map, Gauss represents the low-pass Gaussian filter, and X represents the underwater image.

[0023] Furthermore, step S22, generating the underwater light field window mask, includes the following steps:

[0024] Step 221: Divide the underwater light field feature map G into N windows and multiply them by the learnable parameter S for adaptive weighting;

[0025] G = G·S

[0026] Where G represents the underwater optical field feature map, and S represents the learnable parameters;

[0027] Step 222: Obtain the k windows with the largest weights using the Topk function, and replace their probabilities with 0s in the given window mask matrix M, which is all 1s, using the following formula;

[0028]

[0029] Where M represents a window mask with all values ​​of 1; index represents the index position of the selected k windows in M.

[0030] Furthermore, step S23, which introduces the underwater window mask into the Transformer, includes the following steps:

[0031] Step 231: Aggregate the window mask based on underwater light field information with the keys and values ​​in Self-Attention to obtain the reshaped keys and values;

[0032] Step 232: Perform matrix multiplication on the queried and reshaped keys to obtain the attention map;

[0033] Step 233: Normalize the attention map and perform matrix multiplication with the reshaped value to obtain the final result:

[0034]

[0035] Where Q represents the query, K represents the key, V represents the value, and d represents the coefficient.

[0036] Furthermore, step S23, which integrates the Transformer with a parallel attention mechanism, involves the following specific steps:

[0037] A parallel attention mechanism is introduced, in which the obtained feature maps are input into the Transformer and the Attention mechanism respectively, and a 3×3 convolution is used to fuse the Transformer features and the Attention features.

[0038] Y=Conv(Transformer(X)+Attention(X));

[0039] Where X represents the input and Y represents the output.

[0040] Furthermore, step 24, which constructs a hierarchical encoder-decoder structure based on the U-Net framework, includes a Transformer module and a parallel Attention module. The Transformer module includes a multi-head window self-attention mechanism based on underwater light field mask information and a multilayer perceptron. The Attention module includes a channel attention mechanism and a pixel attention mechanism.

[0041] This invention also includes an underwater image enhancement system based on an underwater light field Transformer, comprising:

[0042] Preprocessing module: Obtains a dataset of real-world paired underwater images and divides it into training, validation, and test sets;

[0043] Network module: Construct a hierarchical encoder-decoder enhancement network based on underwater light field Transformer;

[0044] Training module: This module constructs the loss function, takes paired underwater images as input, sets network parameters and optimizers, etc., and trains the hierarchical encoder-decoder network based on the underwater light field Transformer. We use the L1 loss function and the SSIM loss function for the loss function, and the combined loss function is described below:

[0045] L total =L l1 +L SSIM ;

[0046] Among them, L l1 Let L represent the L1 loss function, L SSIM Represents the SSIM loss function;

[0047] Enhancement module: Inputs the underwater image to be processed into the trained underwater image enhancement network for processing.

[0048] Compared with the prior art, the present invention has the following advantages:

[0049] 1. To address the problem of deep learning models struggling to accurately identify and handle underwater distortions due to a lack of prior information, this invention first processes underwater images using a low-pass Gaussian filter to obtain an underwater light field map. Then, features are extracted from this underwater light field map to obtain an underwater light field feature map. Adaptive weighting of this feature map is applied, and k irrelevant windows are selected to generate an underwater light field window mask. Finally, this underwater light field window mask guides the Transformer to perform feature aggregation, thereby eliminating the impact of underwater noise on model performance.

[0050] 2. This invention takes into account the computational and memory overhead of Transformer, and uses window self-attention and introduces a parallel attention mechanism. This approach not only reduces computational and memory overhead, but also enhances the local information modeling capability of Transformer and improves the performance of the model.

[0051] Based on the above reasons, this invention can be widely applied in fields such as underwater image processing. Attached Figure Description

[0052] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0053] Figure 1 This is a flowchart of the model of the present invention.

[0054] Figure 2 This image shows a comparison of the enhancement effects of this invention and other algorithms on diver images. (a) represents the initial image before restoration; (b) represents the result image processed using the GDCP (Generalization of the Dark Channel Prior for Single Image Restoration) method; (c) represents the result image processed using the WaterNet (An Underwater Image Enhancement Benchmark Dataset and Beyond) method; (d) represents the result image processed using the MLLE (Underwater Image Enhancement via Minimal Color Loss and Locally Adaptive Contrast Enhancement) method; (e) represents the restoration result image using this invention; and (f) represents a reference image of the real scene.

[0055] Figure 3This image shows a comparison of the enhancement effects of this invention and other algorithms on fish images. (a) represents the initial image before restoration; (b) represents the result image processed using the GDCP (Generalization of the Dark Channel Prior for Single Image Restoration) method; (c) represents the result image processed using the WaterNet (An Underwater Image Enhancement Benchmark Dataset and Beyond) method; (d) represents the result image processed using the MLLE (Underwater Image Enhancement via Minimal Color Loss and Locally Adaptive Contrast Enhancement) method; (e) represents the restoration result image using this invention; and (f) represents a reference image of the real scene.

[0056] Figure 4 This image shows a comparison of the enhancement effects of this invention and other algorithms on coral images. (a) represents the initial image before restoration; (b) represents the result image processed using the GDCP (Generalization of the Dark Channel Prior for Single Image Restoration) method; (c) represents the result image processed using the WaterNet (An Underwater Image Enhancement Benchmark Dataset and Beyond) method; (d) represents the result image processed using the MLLE (Underwater Image Enhancement via Minimal Color Loss and Locally Adaptive Contrast Enhancement) method; (e) represents the restoration result image using this invention; and (f) represents a reference image of the real scene. Detailed Implementation

[0057] To enable those skilled in the art to better understand the present invention, the technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort should fall within the scope of protection of the present invention.

[0058] It should be noted that the terms "first," "second," etc., in the specification, claims, and accompanying drawings of this invention are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that the embodiments of the invention described herein can be implemented in orders other than those illustrated or described herein. Furthermore, the terms "comprising" and "having," and any variations thereof, are intended to cover a non-exclusive inclusion; for example, a process, method, system, product, or apparatus that comprises a series of steps or units is not necessarily limited to those steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to such processes, methods, products, or apparatus.

[0059] like Figure 1 As shown, this invention provides an underwater image enhancement method based on an underwater light field transformer, comprising the following steps:

[0060] Step 01: Obtain a dataset of real-world paired underwater images and divide it into training, validation, and test sets;

[0061] Step 02: Construct a hierarchical encoder-decoder enhancement network based on underwater light field Transformer;

[0062] Step 03: Construct a combination of L1 loss function and SSIM loss function constraints to iteratively train the hierarchical encoder-decoder enhancement network based on underwater light field Transformer to obtain the trained underwater image enhancement network;

[0063] Step 04: Input the original underwater image into the trained underwater image enhancement network to enhance the original underwater image.

[0064] Based on the UIEB dataset, 800 images were selected as the training set, 90 images as the validation set, and 60 images as the test set.

[0065] Step 21: Obtain the image input from the network and extract the underwater light field feature map from the input image;

[0066] Step 22: For the underwater light field feature map, first perform parameter weighting processing, then select k windows according to the parameter weights and generate an underwater light field mask;

[0067] Step 23: For the underwater light field mask, reshape the bonds and values ​​in the Transformer module, calculate self-attention, and fuse parallel convolutional attention to obtain the final features;

[0068] Step 24: Stack the Transformer modules and build a hierarchical encoder-decoder structure based on the U-Net framework;

[0069] Step 25: Construct an adaptive Transformer network structure to complete the construction of a hierarchical encoder-decoder enhancement network based on underwater light field.

[0070] Step 211: Extract underwater light field information and use a low-pass Gaussian filter to preserve the light field information of the underwater image;

[0071] Step 212: Extract the underwater light field feature map G from the underwater light field map using a 3×3 convolution and activation function. The formula is as follows:

[0072] G = Conv(Sigmoid((Gauss) σ (X))));

[0073] Where G represents the underwater light field feature map, Gauss represents the low-pass Gaussian filter, and X represents the underwater image.

[0074] Step 221: Divide the underwater light field feature map G into N windows and multiply them by the learnable parameter S for adaptive weighting;

[0075] G = G·S

[0076] Where G represents the underwater optical field feature map, and S represents the learnable parameters;

[0077] Step 222: Obtain the k windows with the largest weights using the Topk function, and replace their probabilities with 0s in the given window mask matrix M, which is all 1s, using the following formula;

[0078]

[0079] Where M represents a window mask with all values ​​of 1; index represents the index position of the selected k windows in M.

[0080] Step 231: Aggregate the window mask based on underwater light field information with the keys and values ​​in Self-Attention to obtain the reshaped keys and values;

[0081] Step 232: Perform matrix multiplication on the queried and reshaped keys to obtain the attention map;

[0082] Step 233: Normalize the attention map and perform matrix multiplication with the reshaped value to obtain the final result:

[0083]

[0084] Where Q represents the query, K represents the key, V represents the value, and d represents the coefficient.

[0085] Step S23: Integrate the Transformer with a parallel attention mechanism: Introduce a parallel attention mechanism, inputting the obtained feature maps into the Transformer and the Attention mechanism respectively, and then using a 3×3 convolution to fuse the Transformer features and the Attention features.

[0086] Y=Conv(Transformer(X)+Attention(X));

[0087] Where X represents the input and Y represents the output.

[0088] Step 24: Construct a hierarchical encoder-decoder structure based on the U-Net framework, which includes a Transformer module and a parallel Attention module. The Transformer module includes a multi-head window self-attention mechanism based on underwater light field mask information and a multilayer perceptron. The Attention module includes a channel attention mechanism and a pixel attention mechanism.

[0089] As a preferred embodiment, this application also includes an underwater image enhancement system based on an underwater light field Transformer, comprising:

[0090] Preprocessing module: Obtains a dataset of real-world paired underwater images and divides it into training, validation, and test sets;

[0091] Network module: Construct a hierarchical encoder-decoder enhancement network based on underwater light field Transformer;

[0092] Training module: This module constructs the loss function, takes paired underwater images as input, sets network parameters and optimizers, etc., and trains the hierarchical encoder-decoder network based on the underwater light field Transformer. We use the L1 loss function and the SSIM loss function for the loss function, and the combined loss function is described below:

[0093] L total =L l1 +L SSIM ;

[0094] Among them, L l1Let L represent the L1 loss function, L SSIM Represents the SSIM loss function;

[0095] Enhancement module: Inputs the underwater image to be processed into the trained underwater image enhancement network for processing.

[0096] Example

[0097] like Figure 2 As shown, this invention provides a comparison of the enhancement effects of other algorithms on diver images. The experimental results show that all four algorithms restored the underwater image to some extent and improved its contrast. However, compared to the reference image, GDCP's restoration effect is poor, and the restored image exhibits severe color cast. While the Water algorithm alleviates the color cast problem, it suffers from overall dullness, low contrast, and loss of background detail. The MLLE algorithm improves contrast and preserves background detail, but the enhanced image suffers from distortion and blurring. The algorithm of this invention better addresses the color cast problem compared to other algorithms, improving the contrast between distant and near-field parts of the underwater image while preserving background detail. Therefore, this invention has a better enhancement effect, effectively solving the color cast problem in underwater images, improving global contrast, and preserving background detail, thus making it closer to the real-world reference image.

[0098] like Figure 3 As shown, this invention provides a comparison of the enhancement effects of other algorithms on fish school images. The experimental results show that all four algorithms restored the underwater image to some extent and improved its contrast. However, compared to the reference image, GDCP's restoration effect is poor, and the restored image exhibits severe color cast. While the Water algorithm alleviates the color cast problem, it suffers from overall dullness, low contrast, and loss of background detail. The MLLE algorithm improves contrast and preserves background detail, but the enhanced image suffers from distortion and blurring. The algorithm of this invention better addresses the color cast problem compared to other algorithms, improving the contrast between distant and near-field parts of the underwater image while preserving background detail. Therefore, this invention has a better enhancement effect, effectively solving the color cast problem in underwater images, improving global contrast, and preserving background detail, thus making it closer to the real-world reference image.

[0099] like Figure 4As shown, this invention provides a comparison of the enhancement effects of other algorithms on coral images. The experimental results show that all four algorithms restored the underwater image to some extent and improved its contrast. However, compared to the reference image, GDCP's restoration effect is poor, and the restored image exhibits severe color cast. While the Water algorithm alleviates the color cast problem, it suffers from overall dullness, low contrast, and loss of background detail. The MLLE algorithm improves contrast and preserves background detail, but the enhanced image suffers from distortion and blurring. The algorithm of this invention better addresses the color cast problem compared to other algorithms, improving the contrast between distant and near-field parts of the underwater image while preserving background detail. Therefore, this invention has a better enhancement effect, effectively solving the color cast problem in underwater images, improving global contrast, and preserving background detail, thus making it closer to the real-world reference image.

[0100] This embodiment compares the experimental results of different algorithms using two objective metrics: PSNR and SSIM. PSNR assesses the similarity between the original signal (e.g., unprocessed image or audio) and the processed signal by comparing their mean squared errors. SSIM considers three aspects of image information: brightness, contrast, and structure. Its calculation is based on this information and aims to simulate the human eye's perception of image quality. Data from Tables 1 and 2 shows that the PSNR and SSIM improvements of the GDCP, WaterNet, and MLLE algorithms are not significant compared to the original image, but the image enhanced by this method is better than all three methods. This invention uses underwater light field information to guide Transformer feature aggregation, introduces a parallel attention mechanism to improve the Transformer's local information interaction capability, and uses a loss function combining L1 and SSIM to correct the structural information of the generated image. Therefore, this invention significantly improves the PSNR and SSIM of the original image and outperforms other underwater image enhancement algorithms.

[0101] Table 1. Comparison of PSNR between the model of this invention and other advanced algorithms.

[0102]

[0103] Table 2. Comparison of SSIM results between the model of this invention and other advanced algorithms.

[0104]

[0105] The sequence numbers of the above embodiments of the present invention are for descriptive purposes only and do not represent the superiority or inferiority of the embodiments. In the above embodiments of the present invention, the descriptions of each embodiment have their own emphasis; parts not described in detail in a certain embodiment can be referred to in the relevant descriptions of other embodiments. It should be understood that the disclosed technical content in the several embodiments provided in this application can be implemented in other ways.

[0106] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some or all of the technical features; and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the scope of the technical solutions of the embodiments of the present invention.

Claims

1. An underwater image enhancement method based on underwater light field Transformer, characterized in that, Includes the following steps: Step 01: Obtain a real-pair underwater image dataset and divide the underwater image dataset into a training set, a validation set, and a test set; Step 02: Construct a hierarchical encoder-decoder enhancement network based on underwater light field Transformer; Step 02, constructing a hierarchical encoder-decoder network based on underwater light field Transformer includes the following steps: Step 21: Obtain the image input from the network and extract the underwater light field feature map from the input image; Step 22: For the underwater light field feature map, first perform parameter weighting processing, then select k windows according to the parameter weights and generate an underwater light field mask; Step 23: For the underwater light field mask, reshape the bonds and values ​​in the Transformer module, calculate self-attention, and fuse parallel convolutional attention to obtain the final features; Step 24: Stack the Transformer modules and build a hierarchical encoder-decoder structure based on the U-Net framework; Step 25: Construct an adaptive Transformer network structure to complete the construction of a hierarchical encoder-decoder enhancement network based on underwater light field; Step 03: Construct a combination of L1 loss function and SSIM loss function constraints to iteratively train the hierarchical encoder-decoder enhancement network based on underwater light field Transformer to obtain the trained underwater image enhancement network; Step 04: Input the original underwater image to be enhanced into the trained underwater image enhancement network to achieve the enhancement processing of the original underwater image.

2. The underwater image enhancement method based on underwater light field Transformer according to claim 1, characterized in that, In step 01, establishing the underwater image dataset includes the following steps: based on the UIEB dataset, arbitrarily select 800 images as the training set, 90 images as the validation set, and 60 images as the test set.

3. The underwater image enhancement method based on underwater light field Transformer according to claim 1, characterized in that, Step 21, extracting the underwater light field feature map from the input image, includes the following steps: Step 211: Extract underwater light field information and use a low-pass Gaussian filter to preserve the light field information of the underwater image; Step 212: Extract the underwater light field feature map G from the underwater light field map using a 3×3 convolution and activation function. The formula is as follows: ； in, G This represents a characteristic image of the underwater light field. Gauss This represents a low-pass Gaussian filter. X This is an underwater image.

4. The underwater image enhancement method based on underwater light field Transformer according to claim 1, characterized in that, In step 22, generating the underwater light field window mask includes the following steps: Step 221: Analyze the underwater light field characteristics. G Divided into N Each window is multiplied by a learnable parameter. S Perform adaptive weighting; in, G This represents a characteristic image of the underwater light field. S Indicates learnable parameters; Step 222: Obtain the k windows with the largest weights using the Topk function, and replace their probabilities with 0s in the given window mask matrix M, which is all 1s, using the following formula; ； in, M This represents a window mask where all values ​​are 1. index This indicates the index positions of the selected k windows in M.

5. The underwater image enhancement method based on underwater light field Transformer according to claim 1, characterized in that, In step 23, introducing the underwater window mask into the Transformer includes the following steps: Step 231: Aggregate the window mask based on underwater light field information with the keys and values ​​in Self-Attention to obtain the reshaped keys and values; Step 232: Perform matrix multiplication on the queried and reshaped keys to obtain the attention map; Step 233: Normalize the attention map and perform matrix multiplication with the reshaped value to obtain the final result: ； in, Q Indicates a query. K Indicates key, V Represents value, d Represents the coefficient.

6. An underwater image enhancement method based on an underwater light field Transformer according to any one of claims 1 or 5, characterized in that, In step 23, the Transformer is fused with a parallel convolutional attention mechanism, and a parallel Attention mechanism is introduced. The obtained feature maps are input into the Transformer and Attention respectively, and a 3×3 convolution is used to fuse the Transformer features and Attention features. ； Where X represents the input and Y represents the output.

7. The underwater image enhancement method based on underwater light field Transformer according to claim 1, characterized in that, In step 24, the hierarchical encoder-decoder structure based on the U-Net framework includes a Transformer module and a parallel Attention module. The Transformer module includes a multi-window self-attention mechanism based on underwater light field masking information and a multilayer perceptron, while the Attention module includes a channel attention mechanism and a pixel attention mechanism.

8. A system for underwater image enhancement based on underwater light field Transformer, applying the underwater image enhancement method based on underwater light field Transformer as described in any one of claims 1-7, characterized in that, include: Preprocessing module: Obtains a dataset of real-world paired underwater images and divides it into training, validation, and test sets; Network module: Construct a hierarchical encoder-decoder enhancement network based on underwater light field Transformer; Training Module: Constructs the loss function, takes paired underwater images as input, sets network parameters and optimizer, and trains a hierarchical encoder-decoder network based on the underwater light field Transformer. The loss function uses L1 loss and SSIM loss, and the combined loss function is expressed as follows: in, Describes the L1 loss function. Represents the SSIM loss function; Enhancement module: Inputs the underwater image to be processed into the trained underwater image enhancement network for processing.

Images