Underwater image enhancement method based on self-attention guidance and correlation feature compensation

Abstract

The application discloses an underwater image enhancement method based on self-attention guidance and associated feature compensation, relates to the technical field of underwater image enhancement, and is used for quality improvement of underwater images with low contrast, color cast, blur and missing texture details and comprises the following steps: high-frequency components and low-frequency components of the underwater image are obtained through frequency domain decomposition; for the high-frequency components, a texture detail feature compensation module is constructed to obtain a high-frequency enhancement result; for the low-frequency components, a color space feature compensation module is constructed to obtain a low-frequency enhancement result; a feature fusion optimization module is constructed, the high-frequency enhancement result and the low-frequency enhancement result are subjected to feature fusion, an ultimate enhancement result is obtained, a loss function is constructed, the ultimate enhancement result is optimized, and an ultimate enhancement image is output. The application splits the enhancement task into two subtasks of texture enhancement and color correction, processes the high-frequency information and the low-frequency information according to the division, and makes the output image have clearer structure and more stable overall visual effect.

Description

Technical Field

[0001] This invention relates to the field of image data processing technology, and more particularly to the field of underwater image enhancement technology, specifically to an underwater image enhancement method based on self-attention guidance and associated feature compensation. Background Technology

[0002] Underwater imagery plays a crucial role in tasks such as ocean observation, seabed resource exploration, underwater robotic operations, marine ranching monitoring, and underwater target identification. However, the underwater imaging environment differs significantly from the terrestrial imaging environment. Light propagating in water is affected by wavelength-dependent absorption and scattering; red light attenuates more rapidly, while blue and green light travel relatively longer distances. This often results in underwater images exhibiting degradation phenomena such as blue-green color cast, decreased contrast, uneven brightness, blurred edges, and weakened details and textures. Furthermore, suspended particles, plankton, sediment, and complex lighting conditions in the water cause forward and backscattering, leading to problems such as fogging, noise, or unclear structural boundaries in localized areas of the image.

[0003] Existing underwater image enhancement methods can be broadly categorized into three types. The first type is based on physical models. These methods typically establish an underwater imaging model and use dark channel priors, red channel priors, light attenuation priors, background light estimation, or transmission map estimation to infer a clear image. This type of method is effective when degradation conditions are relatively stable, but the turbidity, light intensity, shooting distance, and target material vary significantly in different areas of the underwater environment. A single physical prior cannot accurately describe complex mixed degradation, easily leading to inaccurate model parameter estimation, color overcompensation, or local enhancement failure. The second type is traditional enhancement methods that do not rely on physical models. These methods typically adjust pixel distribution directly through histogram equalization, color balancing, Retinex, gamma correction, contrast stretching, or multi-result fusion. These methods are simple to implement and computationally inexpensive, but they rely primarily on global statistical laws or manually designed rules for enhancement, making it difficult to simultaneously address color correction, detail preservation, and noise suppression. This often results in over-enhancement, abnormal saturation, texture edge erasure, or artificial artifacts. The third type is underwater image enhancement methods based on deep learning. These methods typically employ convolutional neural networks, encoder-decoder structures, multi-branch structures, or generative adversarial networks to learn the mapping relationship from degraded images to reference enhanced images from training samples. Deep learning methods can learn complex nonlinear mappings, but in real underwater scenes, it is difficult to obtain a large number of strictly paired degraded images and clear reference images, resulting in limited training constraint information. Furthermore, underwater degradation is usually caused by a combination of factors such as color shift, blurring, noise, and contrast loss. If a single network branch is used to simultaneously handle color restoration and texture enhancement tasks, the optimization objectives can easily interfere with each other.

[0004] While schemes that directly output enhanced images using a single image-to-image mapping network or a conventional encoder-decoder network can achieve end-to-end enhancement, they do not fully utilize the information differences carried by high-frequency and low-frequency components in the image's frequency domain. Generally, high-frequency components are more concentrated in representing edges, textures, structural details, and local noise, while low-frequency components are more concentrated in representing color, illumination, overall brightness, and slowly changing background information. Therefore, relying solely on a single mapping to uniformly process hybrid degradation can easily lead to problems such as insufficient color correction, inadequate texture compensation, or unstable details after enhancement.

[0005] Furthermore, different color spaces emphasize different aspects of underwater image degradation. RGB color space is convenient for imaging and display, but its three channels are highly correlated, making it susceptible to lighting and occlusion. HSV color space can more intuitively describe hue, saturation, and brightness. Lab color space can more closely approximate human visual perception and separate brightness and color information. Existing augmentation networks, without filtering and distilling features from multiple color spaces, may directly feed redundant or invalid color space features into the network, increasing the learning difficulty of the model.

[0006] Therefore, there is a need for an underwater image enhancement method that can first decouple the underwater image enhancement task into two sub-tasks: texture detail enhancement and color correction, then perform targeted enhancement through self-attention and correlation compensation mechanisms, and finally achieve high- and low-frequency fusion optimization. Summary of the Invention

[0007] The purpose of this invention is to provide an underwater image enhancement method based on self-attention guidance and associated feature compensation, so as to solve the problem that it is difficult to balance color correction and texture structure preservation when performing underwater image enhancement in the prior art.

[0008] To address the above objectives, this invention provides an underwater image enhancement method based on self-attention guidance and associated feature compensation, comprising: S1. Obtain the original image data, perform frequency domain decomposition on each original image in the original image data, and obtain the high-frequency components and low-frequency components of each original image; S2. Perform detail detection on the high-frequency components obtained in S1 to obtain the corresponding texture detail features. Perform color space conversion on the low-frequency components obtained in S1 to obtain the color space features of the corresponding color space. Detail detection includes edge detection and texture detection. Color spaces include RGB, Lab and HSV. S3. For high-frequency components, construct a texture detail feature compensation module. Input the high-frequency components obtained in S1 and the texture detail features obtained in S2 into the texture detail feature compensation module to obtain high-frequency enhancement results. S4. For low-frequency components, construct a color space feature compensation module. Input the low-frequency components obtained in S1 and the color space features obtained in S2 into the color space feature compensation module to obtain low-frequency enhancement results. The color space feature compensation module includes a self-attention threshold control module and a feature joint distillation module. S5. Construct a feature fusion optimization module to fuse the high-frequency enhancement results obtained in S3 and the low-frequency enhancement results obtained in S4 to obtain the final enhancement result; S6. Construct a loss function to optimize the texture detail feature compensation module, color space feature compensation module, and feature fusion optimization module, and construct an underwater image enhancement network model. The loss function includes high-frequency loss, low-frequency loss, and pixel loss. High-frequency loss is used to optimize the model parameters of the texture detail feature compensation module, low-frequency loss is used to optimize the model parameters of the color space feature compensation module, and pixel loss is used to optimize the model parameters of the feature fusion optimization module. S7. The underwater image enhancement network model built on S6 performs image enhancement processing on the underwater image to be enhanced and outputs the final enhanced image.

[0009] S3 includes: S3.1 Based on the high-frequency components obtained in S1, the texture detail feature compensation module generates pixel attention maps of each texture detail feature. Based on each pixel attention map, the corresponding texture detail features are weighted and filtered to generate multiple compensated texture feature maps and obtain multiple compensated texture features. S3.2, perform feature concatenation on the multiple compensation texture features obtained in S3.1, and extract features from the concatenated compensation texture features based on pixel attention and channel attention to obtain effective texture compensation features; S3.3, add the effective texture compensation features to the high-frequency components obtained in S1 to obtain the high-frequency enhancement result.

[0010] S4 includes: S4.1, extract features from the color space features of Lab and HSV respectively, obtain the channel features of each channel of Lab and the channel features of each channel of HSV, input the channel features of each channel of Lab and the channel features of each channel of HSV into the self-attention threshold control module, generate the self-attention feature map of the corresponding color space according to the complementary characteristics of the values ​​of each channel, and adaptively enhance the channel features of the corresponding color space based on the respective attention feature map; S4.2 Perform multi-scale downsampling on the color space features of the three color spaces RGB, Lab and HSV respectively to obtain color space features at different scales of the corresponding color space. Construct a multi-scale color space feature set for the corresponding color space based on the color space features of each color space. Generate a joint attention map at each scale based on the color space features of the three color spaces at each scale. S4.3, input the multi-scale color space features of each color space obtained in S4.2 into the feature joint distillation module, filter the color space features of Lab and HSV according to the joint attention map obtained in S4.2, and superimpose the filtered color space features into the RGB color space features to obtain the color distillation features of the corresponding scale; S4.4 performs low-frequency decoding on the color distillation features obtained in S4.3 to obtain low-frequency enhancement results.

[0011] High frequency enhancement results for: ; ; In the formula, Represents high-frequency components. Indicates effective texture compensation features. This represents the stitched compensated texture features, which is equivalent to the feature obtained by stitching all the compensated texture features along the channel dimension. Represents the pixel attention module. This indicates the channel attention module; The compensated texture features are: ; ; In the formula, Index of the number of texture detail features. Indicates the first Individual texture features, express pixel attention map, Indicates correspondence Compensation texture features, This represents element-wise multiplication. This indicates a feature concatenation operation. This represents the pixel attention module.

[0012] The self-attention feature map of Lab is as follows: ; ; ; In the formula, This represents the channel index; for Lab, The value is , This represents the brightness channel of Lab. This indicates the red and green opposing channels of Lab. This represents the blue and yellow opposing channels of Lab. Belongs to, Indicates the first of Lab Self-attention feature map of each channel, Indicates the first of Lab Channel characteristics of each channel, express Self-attention graph.

[0013] For the saturation and brightness channels of HSV, the self-attention feature map is as follows: ; ; ; ; In the formula, This represents the self-attention feature map of the saturation channel of HSV. The channel characteristics representing the saturation channel of HSV. express Self-attention graph; This represents the self-attention feature map of the brightness channel of HSV. This indicates the channel characteristics of the lightness channel in HSV. express Self-attention graph; For the hue channels of HSV, angular metrics are used to represent color information, and a self-attention feature map is constructed based on the angular metrics: ; ; In the formula, A self-attention feature map representing the hue channel of HSV. This represents a segmented self-attention feature map of the hue channel in HSV. This indicates the channel characteristics of the hue channel in HSV. Represents angular measurement.

[0014] In S4.3, the color distillation feature corresponding to the scale is: ; ; In the formula, Indicates scale index; Indicates the first Color distillation characteristics at various scales; Indicates the first Joint attention maps at multiple scales; , , They represent the first At each scale, the color space characteristics of RGB, Lab, and HSV are analyzed. This represents the activation function.

[0015] In S5, the feature fusion optimization module includes a long skip connection module, four convolutional layers, an activation function, and three cascaded feature attention blocks. Each feature attention block contains a recalibration structure to adjust the enhancement intensity of different channels and different spatial locations. The long skip connection module is used to implement residual connections, the convolutional layers are used to reduce the dimensionality of features, and the activation function is used to output the final enhanced image. The features from the input feature fusion optimization module are sequentially passed through the first feature attention block, the first convolutional layer, the second feature attention block, the second convolutional layer, the third feature attention block, and the third convolutional layer to obtain deep enhanced features. The deep enhanced features enter the long skip connection module, which adds the input features from the feature fusion optimization module to the deep enhanced features after the third convolutional layer element by element. The added features are then passed through the fourth convolutional layer and the activation function for output.

[0016] In S7, the loss function for: ; In the formula, Indicates high-frequency loss. Indicates low-frequency loss. Indicates pixel loss, , , They represent , , Loss weights; ; ; ; In the formula, Represents the number of pixel rows, Indicates the number of pixel columns. is the pixel coordinate index, indicating the . Line number Column of pixels, This indicates that the high-frequency enhancement results are in Pixel value at that location, This indicates that the high-frequency enhancement results are in Pixel reference value at that location, This indicates that the low-frequency enhancement results are in Pixel value at that location, This indicates that the low-frequency enhancement results are in Pixel reference value at that location, This indicates the final enhancement result in Pixel value at that location, This indicates the final enhancement result in Pixel reference value at that location, This indicates that the sum of squares is calculated over all pixel positions.

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

[0018] This invention decomposes the underwater image enhancement task into two relatively clear sub-tasks, high-frequency texture enhancement and low-frequency color correction, through frequency domain decomposition, thereby reducing the difficulty of a single network to process multiple degradation mixture maps at the same time.

[0019] This invention uses a multi-input texture feature association compensation module to associate and model external texture cues such as edge detection and texture detection with high-frequency components, and uses pixel attention and channel attention to filter effective texture compensation information, which can better restore edge and texture details.

[0020] This invention constructs a self-attention feature map based on the prior expression of different color space channels through a self-attention threshold control module, selectively enhancing the effective color information in Lab and HSV spaces, reducing interference from invalid and redundant features. At the same time, through a feature joint distillation module, feature screening and distillation are performed between RGB, Lab and HSV multiple color spaces, enabling the color correction branch to absorb the advantageous expression of different color spaces and improve the accuracy of color restoration.

[0021] This invention further fuses the enhanced high-frequency texture features and low-frequency color features through a feature fusion optimization module, so that the final output simultaneously takes into account color naturalness, contrast, edge sharpness and overall visual consistency. Attached Figure Description

[0022] Figure 1 The overall network structure diagram of the underwater image enhancement network model provided by this invention; Figure 2 for Figure 1 Network architecture diagram of the mid-texture detail feature compensation module; Figure 3 for Figure 1 Network architecture diagram of the color space feature compensation module; Figure 4The flowchart for generating a self-attention feature map of Lab provided by the present invention; Figure 5 A flowchart for obtaining color distillation features at a corresponding scale is provided for this invention; Figure 6 The feature fusion optimization module structure diagram provided by the present invention. Detailed Implementation

[0023] To make the objectives, technical solutions, and advantages of this invention clearer, the technical solutions of this invention are described clearly and completely below. Obviously, the described embodiments are only some, not all, of the embodiments of this invention. All other embodiments obtained by those skilled in the art based on the embodiments of this invention without creative effort are within the scope of protection of this invention.

[0024] Underwater image enhancement methods based on self-attention guidance and associated feature compensation include: S1. Obtain the original image data, perform frequency domain decomposition on each original image in the original image data, and obtain the high-frequency components and low-frequency components of each original image; S2. Perform detail detection on the high-frequency components obtained in S1 to obtain the corresponding texture detail features. Perform color space conversion on the low-frequency components obtained in S1 to obtain the color space features of the corresponding color space. Detail detection includes edge detection and texture detection. Color spaces include RGB, Lab and HSV. S3. For high-frequency components, construct a texture detail feature compensation module. Input the high-frequency components obtained in S1 and the texture detail features obtained in S2 into the texture detail feature compensation module to obtain high-frequency enhancement results. S4. For low-frequency components, construct a color space feature compensation module. Input the low-frequency components obtained in S1 and the color space features obtained in S2 into the color space feature compensation module to obtain low-frequency enhancement results. The color space feature compensation module includes a self-attention threshold control module and a feature joint distillation module. S5. Construct a feature fusion optimization module to fuse the high-frequency enhancement results obtained in S3 and the low-frequency enhancement results obtained in S4 to obtain the final enhancement result; S6. Construct a loss function to optimize the texture detail feature compensation module, color space feature compensation module, and feature fusion optimization module, and construct an underwater image enhancement network model. The loss function includes high-frequency loss, low-frequency loss, and pixel loss. High-frequency loss is used to optimize the model parameters of the texture detail feature compensation module, low-frequency loss is used to optimize the model parameters of the color space feature compensation module, and pixel loss is used to optimize the model parameters of the feature fusion optimization module. S7. The underwater image enhancement network model built on S6 performs image enhancement processing on the underwater image to be enhanced and outputs the final enhanced image.

[0025] S3 includes: S3.1 Based on the high-frequency components obtained in S1, the texture detail feature compensation module generates pixel attention maps of each texture detail feature. Based on each pixel attention map, the corresponding texture detail features are weighted and filtered to generate multiple compensated texture feature maps and obtain multiple compensated texture features. S3.2, perform feature concatenation on the multiple compensation texture features obtained in S3.1, and extract features from the concatenated compensation texture features based on pixel attention and channel attention to obtain effective texture compensation features; S3.3, add the effective texture compensation features to the high-frequency components obtained in S1 to obtain the high-frequency enhancement result.

[0026] S4 includes: S4.1, extract features from the color space features of Lab and HSV respectively, obtain the channel features of each channel of Lab and the channel features of each channel of HSV, input the channel features of each channel of Lab and the channel features of each channel of HSV into the self-attention threshold control module, generate the self-attention feature map of the corresponding color space according to the complementary characteristics of the values ​​of each channel, and adaptively enhance the channel features of the corresponding color space based on the respective attention feature map; S4.2 Perform multi-scale downsampling on the color space features of the three color spaces RGB, Lab and HSV respectively to obtain color space features at different scales of the corresponding color space. Construct a multi-scale color space feature set for the corresponding color space based on the color space features of each color space. Generate a joint attention map at each scale based on the color space features of the three color spaces at each scale. S4.3, input the multi-scale color space features of each color space obtained in S4.2 into the feature joint distillation module, filter the color space features of Lab and HSV according to the joint attention map obtained in S4.2, and superimpose the filtered color space features into the RGB color space features to obtain the color distillation features of the corresponding scale; S4.4 performs low-frequency decoding on the color distillation features obtained in S4.3 to obtain low-frequency enhancement results.

[0027] High frequency enhancement results for: ; ; In the formula, Represents high-frequency components. Indicates effective texture compensation features. This represents the stitched compensated texture features, which is equivalent to the feature obtained by stitching all the compensated texture features along the channel dimension. Represents the pixel attention module. This indicates the channel attention module; The compensated texture features are: ; ; In the formula, Index of the number of texture detail features. Indicates the first Individual texture features, express pixel attention map, Indicates correspondence Compensation texture features, This represents element-wise multiplication. This indicates a feature concatenation operation. This represents the pixel attention module.

[0028] The self-attention feature map of Lab is as follows: ; ; ; In the formula, This represents the channel index; for Lab, The value is , This represents the brightness channel of Lab. This indicates the red and green opposing channels of Lab. This represents the blue and yellow opposing channels of Lab. Belongs to, Indicates the first of Lab Self-attention feature map of each channel, Indicates the first of Lab Channel characteristics of each channel, express Self-attention graph.

[0029] For the saturation and brightness channels of HSV, the self-attention feature map is as follows: ; ; ; ; In the formula, This represents the self-attention feature map of the saturation channel of HSV. The channel characteristics representing the saturation channel of HSV. express Self-attention graph; This represents the self-attention feature map of the brightness channel of HSV. This indicates the channel characteristics of the lightness channel in HSV. express Self-attention graph; For the hue channels of HSV, angular metrics are used to represent color information, and a self-attention feature map is constructed based on the angular metrics: ; ; In the formula, A self-attention feature map representing the hue channel of HSV. This represents a segmented self-attention feature map of the hue channel in HSV. This indicates the channel characteristics of the hue channel in HSV. Represents angular measurement.

[0030] In S4.3, the color distillation feature corresponding to the scale is: ; ; In the formula, Indicates scale index; Indicates the first Color distillation characteristics at various scales; Indicates the first Joint attention maps at multiple scales; , , They represent the first At each scale, the color space characteristics of RGB, Lab, and HSV are analyzed. This represents the activation function.

[0031] In S5, the feature fusion optimization module includes a long skip connection module, four convolutional layers, an activation function, and three cascaded feature attention blocks. Each feature attention block contains a recalibration structure to adjust the enhancement intensity of different channels and different spatial locations. The long skip connection module is used to implement residual connections, the convolutional layers are used to reduce the dimensionality of features, and the activation function is used to output the final enhanced image. The features from the input feature fusion optimization module are sequentially passed through the first feature attention block, the first convolutional layer, the second feature attention block, the second convolutional layer, the third feature attention block, and the third convolutional layer to obtain deep enhanced features. The deep enhanced features enter the long skip connection module, which adds the input features from the feature fusion optimization module to the deep enhanced features after the third convolutional layer element by element. The added features are then passed through the fourth convolutional layer and the activation function for output.

[0032] In S7, the loss function for: ; In the formula, Indicates high-frequency loss. Indicates low-frequency loss. Indicates pixel loss, , , They represent , , Loss weights; ; ; ; In the formula, Represents the number of pixel rows, Indicates the number of pixel columns. is the pixel coordinate index, indicating the . Line number Column of pixels, This indicates that the high-frequency enhancement results are in Pixel value at that location, This indicates that the high-frequency enhancement results are in Pixel reference value at that location, This indicates that the low-frequency enhancement results are in Pixel value at that location, This indicates that the low-frequency enhancement results are in Pixel reference value at that location, This indicates the final enhancement result in Pixel value at that location, This indicates the final enhancement result in Pixel reference value at that location, This indicates that the sum of squares is calculated over all pixel positions.

[0033] The core of the underwater image enhancement method based on self-attention guidance and associated feature compensation is not to enhance the entire image through a single path, but to divide the processing according to the degradation differences of high-frequency and low-frequency information. Among them, the high-frequency branch focuses on solving the problems of edge blurring, texture weakening and loss of detail; the low-frequency branch focuses on solving the problems of blue-green color cast, uneven brightness and color expression distortion; the fusion branch is responsible for unifying and optimizing the two types of enhancement results, so that the output image has more natural colors, clearer structure and more stable overall visual effect.

[0034] In S1, existing frequency decomposition methods (such as Fourier transform, wavelet transform, spatial domain filtering, etc.) are used to decompose the original image into frequency components, extracting the high-frequency and low-frequency components. The core is to separate signals of different frequencies in the image. Taking Fourier transform as an example, the original image is first transformed from the spatial domain to the frequency domain. Then, low-pass and high-pass filters are applied to the output. Finally, inverse Fourier transform is applied to the low-pass and high-pass filters to obtain the filtered image in the spatial domain, thus obtaining the low-frequency and high-frequency components of the original image. Taking spatial domain filtering as an example, Gaussian blurring is directly applied to the original image, and the result is the low-frequency component. The high-frequency component can be obtained by subtracting the low-frequency component from the original image. The low-frequency component contains the main outline and color changes of the image, while the high-frequency component contains the details in the image, corresponding to rapidly changing areas such as edges, textures, and details.

[0035] In S2, detail detection of high-frequency components includes edge detection or texture detection of the high-frequency components obtained in S1. Detail detection can be performed directly on the original image instead of the high-frequency components. Edges are sets of pixels in an image where grayscale changes drastically. Edge detection can be performed on the high-frequency components using existing edge detection algorithms (such as those based on classical edge detection operators and those based on deep learning). Texture describes the spatial distribution pattern of pixel grayscale or color in an image; texture features at different scales can be extracted using methods such as bandpass filters and wavelet transforms. Classical edge detection operators include the Sobel operator, Canny operator, and Scharr operator. The Sobel operator detects edges by calculating the first derivatives (gradients) of the image in the horizontal and vertical directions. The Canny operator detects image edges based on a multi-stage optimization algorithm, including Gaussian filtering for noise reduction, calculating gradient magnitude and direction, non-maximum suppression (thinning edges), and double threshold detection (connecting edges). The Scharr operator is an improved version of the Sobel operator, using more accurate kernel coefficients and being more sensitive to gradient changes. Deep learning-based edge detection algorithms utilize end-to-end edge detection networks, fusing features from different levels through multi-scale output layers to learn edge information from global to local perspectives.

[0036] RGB is the most compatible color space with display hardware. Its values ​​are intuitive three-channel intensities, and its color range covers the color gamut of common display devices. In addition, RGB is a color model based on the principle of additive color mixing. The R channel is the red channel, the G channel is the green channel, and the B channel is the blue channel, corresponding to the three colors of light. Lab refers to the CIELAB color space, which covers all colors visible to the human eye and uniformly expresses color differences. HSV organizes colors in a way that is intuitively perceived by humans. H, S, and V are abbreviations for Hue, Saturation, and Value, respectively, corresponding to the Hue channel, Saturation channel, and Value channel of HSV.

[0037] In S4.1, for the hue channel of HSV, angular metrics are used to represent color information. The value range of the angular metric is from zero to 360 degrees, and the angular metric satisfies the following condition: A self-attention feature map is constructed based on the angle measurement.

[0038] In S6, a multi-level loss function is constructed during the training phase, and network parameters are jointly optimized. For each reference augmented image (clear image), the same frequency domain decomposition method as in S1 is used to obtain its high-frequency reference component and low-frequency reference component. The pixel reference value at the corresponding position (pixel) of the high-frequency augmentation result is obtained based on the high-frequency reference component, and the pixel reference value at the corresponding position (pixel) of the low-frequency augmentation result is obtained based on the low-frequency reference component. The low-frequency branch uses low-frequency loss to constrain the pixel error between the low-frequency augmentation result and the low-frequency reference. The low-frequency branch is used to strengthen color correction learning, and the low-frequency loss is used to separately supervise the low-frequency color correction result, ensuring that the low-frequency processing branch does not only rely on backpropagation of the final output but also directly learns the color restoration mapping during training. The high-frequency branch uses high-frequency loss to constrain the pixel error between the high-frequency augmentation result and the high-frequency reference. The high-frequency branch is used to strengthen edge and texture detail preservation, and the high-frequency loss is used to separately supervise the high-frequency texture augmentation result, enabling the high-frequency branch to more explicitly learn texture compensation and edge enhancement. The final output uses pixel loss to constrain the pixel error between the fused augmented result and the reference augmented image. The pixel loss is used to supervise the final output after FRM fusion, ensuring that the final augmented image is close to the reference augmented image at the entire pixel level. The total loss is obtained by summing the above losses according to their weights. It is used for end-to-end training of the underwater image enhancement network model SAFCNet.

[0039] Among them, the low-frequency loss is directed towards the low-frequency color correction branch, which mainly constrains the network to learn blue-green bias correction, brightness adjustment and overall color distribution restoration, so that the low-frequency enhancement result is as close as possible to the low-frequency components of the reference image.

[0040] The high-frequency loss branch focuses on high-frequency texture enhancement, primarily constraining the network's recovery of edge, contour, local texture, and detail information, ensuring that the high-frequency enhancement result maintains consistency with the high-frequency details of the reference image. Since the mean squared error is more sensitive to larger deviations, this loss helps reduce texture blurring and edge loss.

[0041] Pixel loss is applied to the final enhanced image output by the Feature Fusion Optimization (FRM) module to constrain the overall restoration effect of color, contrast, structure, and texture.

[0042] Final total loss Depend on , and The weights are combined to simultaneously constrain high-frequency branches, low-frequency branches, and the final output within a single training objective. Preferably, the loss weights... , , All values ​​can be set to 100. In different datasets or engineering application scenarios, the weights of the three losses can also be adjusted according to the emphasis on color correction, texture restoration, and overall visual quality. By jointly constructing the total loss from the low-frequency loss, high-frequency loss, and pixel loss, the problem of insufficient branch learning caused by only supervising the final result is avoided, so that the high-frequency enhancement branch, the low-frequency color branch, and the final fusion module all have clear optimization constraints.

[0043] To validate the underwater image enhancement network model SAFCNet, we compared it with existing image enhancement algorithms on the UIEB, UIEB-Challenging, RUIE, EUVP, and SQUID datasets. The comparison was based on SSIM (Structural Similarity Index), PSNR (Peak Signal-to-Noise Ratio), UIQM (Underwater Image Quality Measure), UCIQE (Underwater Color Image Quality Evaluation), MSE (Mean Squared Error), and ψ (Color Restoration Error). Higher values ​​for SSIM, PSNR, UIQM, and UCIQE indicate better enhancement quality, while lower values ​​for MSE and ψ indicate smaller errors or smaller color restoration errors. The UIEB dataset provides reference augmented images; therefore, SSIM, PSNR, and MSE are used for full-reference evaluation of underwater image augmentation network models and existing image augmentation algorithms. The full-reference quantization results of SAFCNet and the contrasting methods on the UIEB dataset are shown in Table 1. The UIEB-Challenging, RUIE, and EUVP datasets do not have corresponding reference augmented images; therefore, UIQM and UCIQE are used for no-reference evaluation of underwater image augmentation network models and existing image augmentation algorithms. The no-reference quantization results of SAFCNet and the contrasting methods on the UIEB-Challenging, RUIE, and EUVP datasets are shown in Table 2. The SQUID dataset consists of stereo images taken in environments with different water properties (such as clear and turbid). The images include standard color charts for color calibration, which are used to quantitatively evaluate the performance of underwater image restoration algorithms, especially image dehazing and color restoration algorithms. Therefore, on the SQUID dataset, ψ, UIQM, and UCIQE are used to evaluate the color restoration of underwater image enhancement network models and existing deep learning-based image enhancement algorithms. The color restoration quantification results of SAFCNet and deep learning-based comparison methods on the SQUID dataset are shown in Table 3.

[0044] Table 1. Full Reference Quantization Results; ;

[0045] Table 2 has no reference quantitative results; ;

[0046] In Tables 1 and 2, Method 1 is a Two-Step underwater image enhancement method, which automatically selects the optimal color compensation method through a two-step strategy and performs multi-scale fusion using weights such as contrast, saturation, and exposure to obtain the enhanced image; Method 2 is ULAP, which is based on a physical model of underwater light attenuation prior, estimating scene depth and transmittance by constructing a linear relationship between the depth map and the intensity difference of red, green, and blue light, and then recovering the image; Method 3 is CBF (Conventional Beamforming), mainly used in the field of sonar image processing; Method 4 is WaterNet, a gated fusion network that first preprocesses the input image using three methods: white balance, gamma correction, and histogram equalization, and then... The results are adaptively fused with the confidence map predicted by the network to obtain the final output. Method 5 is Ucolor, which is a multi-color space embedding network based on medium transport guidance. By combining the advantages of physical models and deep learning, it performs feature learning in multiple color spaces and uses the medium transport map of the physical model for guidance. It shows superior performance in both visual and quantitative metrics. Method 6 is UICoENet, which addresses the similar degradation characteristics of multiple images in the same scene. It uses a twin encoder and decoder architecture to process two related images simultaneously and performs joint learning by mining complementary information between the two images through a related feature matching unit. Method 7 is SAFCNet, an underwater image enhancement network model built by S6. As shown in Table 1, SAFCNet achieves the highest SSIM (0.888) and PSNR (23.520) on the UIEB dataset and the lowest MSE (0.420). Compared with UICoENet (0.875), which has the second highest SSIM, SAFCNet improves SSIM by 0.013; compared with Ucolor (20.742), which has the second highest PSNR, SAFCNet improves PSNR by 2.778. As shown in Table 2, in the absence of reference evaluation, the traditional method CBF scores higher on some UIQM metrics, but SAFCNet performs more stably overall among all deep learning methods; it achieves an absolute best value of 0.552 on the UCIQE metric of RUIE and 0.591 on the UCIQE metric of EUVP, which is only lower than ULAP's 0.600.

[0047] Table 3. Quantitative results of color restoration; ;

[0048] In Table 3, Method 1 is WaterNet, Method 2 is Ucolor, Method 3 is UICoENet, and Method 4 is SAFCNet, an underwater image enhancement network model built with S6. Methods 1 to 3 are existing deep learning-based image enhancement algorithms. As shown in Table 3, SAFCNet achieves the lowest ψ (16.889) and the highest UIQM (0.263) and UCIQE (0.535) on the SQUID dataset, indicating its advantages in both color correction accuracy and subjective quality-related metrics.

[0049] To verify the contribution of the four key modules—MTAC (Texture Detail Compensation Module), STCM (Self-Attention Threshold Control Module), JFDM (Joint Feature Distillation Module), and FRM (Feature Fusion Optimization Module)—to the final enhancement performance of the underwater image enhancement network model SAFCNet, ablation experiments were conducted on the underwater image enhancement network model. By removing one module at a time (keeping other parts unchanged), multiple ablation models were constructed, and the performance changes were evaluated on the same training settings and test set (UIEB dataset). The necessity of each module was quantified. The quantitative results of the ablation experiments for each ablation model are shown in Table 4, and the performance changes of each ablation model with a single module removed relative to the complete model are shown in Table 5.

[0050] Table 4. Quantization results of the ablation model on the UIEB dataset; ;

[0051] Table 5 shows the performance changes of each ablation model with individual modules removed relative to the complete model; ;

[0052] As shown in Tables 4 and 5, the complete SAFCNet outperforms each ablation model that removes a single module in terms of SSIM, PSNR, and MSE, indicating that all four core modules (MTAC, STCM, JFDM, and FRM) are indispensable and each contributes positively to the final enhancement result. Among these, removing JFDM resulted in the largest decrease in PSNR and the largest increase in MSE, indicating that JFDM has the greatest impact on pixel-level reconstruction error (PSNR / MSE) and is key to improving signal fidelity. Removing MTAC resulted in a significant decrease in SSIM, indicating that MTAC has the greatest impact on structural similarity (SSIM), demonstrating its unique role in preserving texture details and spatial structure.

[0053] On the UIEB dataset with reference, SAFCNet achieves the best results in SSIM, PSNR, and MSE, indicating its superior performance in structure preservation, pixel-level reconstruction error, and overall enhancement quality. On unreferenced datasets such as UIEB-Challenging, RUIE, and EUVP, SAFCNet demonstrates stable overall performance among deep learning methods, showcasing its ability to generalize to complex underwater degradation environments. In the SQUID color restoration evaluation, SAFCNet achieves both the lowest color angle error and the highest unreferenced quality index, indicating its good color correction capability. Ablation experiments show that all four modules—MTAC, STCM, JFDM, and FRM—are indispensable, and the complete model achieves the best overall performance.

[0054] Figure 1 This is a diagram showing the overall network structure of the underwater image enhancement network model. Figure 2 for Figure 1 Network architecture diagram of the mid-texture detail feature compensation module. Figure 3 for Figure 1 The network architecture diagram of the color space feature compensation module is as follows: Figures 1 to 3 As shown, the input underwater image (original image) is first fed into the frequency domain decomposition unit, where it is separated into high-frequency and low-frequency components. The high-frequency components reflect rapidly changing information such as image edges, textures, and structural details; the low-frequency components reflect slowly changing information such as image color, brightness, illumination, and large-scale background. Thus, the network decomposes a complex underwater image enhancement task into two sub-tasks: "texture detail enhancement" and "color correction." In the high-frequency branch, the input image is processed by edge detection or texture detection to obtain multiple texture feature maps. Each map, along with the high-frequency components, is fed into the texture detail feature compensation module (MTAC). The texture detail feature compensation module first generates and... Individual texture features Corresponding pixel attention image reuse right Element-wise weighting is performed to obtain the filtered texture compensation features. Multiple After concatenation, effective texture compensation features are obtained through pixel attention and channel attention filtering, and then added to the high-frequency components to obtain the high-frequency enhancement result. In the low-frequency branch, the low-frequency components are converted into Lab and HSV color space representations, respectively. The Lab space is beneficial for expressing luminance and color components, while the HSV space is beneficial for describing hue, saturation, and lightness. The network adaptively enhances each color space channel through the Self-Attention Threshold Control (STCM) module, and filters effective color compensation features among RGB, Lab, and HSV color spaces through the Feature Joint Distillation (JFDM) module. Finally, the low-frequency color enhancement result is obtained through decoding. (The text then abruptly shifts to a different topic: calculating the...) When generating joint attention maps at multiple scales, the Sigmoid function is used as the activation function. In the fusion output stage, the high-frequency enhancement results and low-frequency enhancement results are fed into the Feature Fusion Optimization (FRM) module. FRM further adjusts the relationship between high-frequency textures and low-frequency colors through cascaded feature attention blocks and residual connections, and outputs the final enhanced image through convolution and the Tanh activation function.

[0055] Figure 4 The flowchart for generating the self-attention feature map of Lab based on the self-attention threshold control module is as follows: Figure 4 As shown, the channel features of the input color space Self-attention map of the corresponding channel They jointly participate in pixel attention calculation to obtain a weight map used to control the degree of feature enhancement for that color channel, which is the self-attention feature map of the corresponding color space. Figure 4 In the text, 'c' represents the channel index, which can refer to the L, a, b channels in Lab color space, or the H, S, V channels in HSV color space. The module's goal is not to simply enhance all channels indiscriminately, but to automatically determine which regions need enhancement and which regions should be preserved or suppressed based on the degradation characteristics of each color channel. and First, feature concatenation is performed, followed by input into a pixel attention unit consisting of convolutional and activation layers. This unit generates a spatial weight map representing the importance of the current color channel at different spatial locations. Subsequently, the weight map is multiplied point-to-point with the original channel features to achieve adaptive enhancement of local regions within the channel. Finally, the module adds the enhanced information to the original input through long skip connections to obtain the output features. The key to this structure lies in "self-attention guidance" and "residual enhancement." (Self-attention map) The enhancement is derived from the current color channel itself, allowing the network to utilize the color distribution information within the channel to form an enhancement basis. Residual connections ensure that the enhancement process does not completely destroy the original features, but rather selectively compensates for them based on the original channel features. This strengthens areas conducive to color recovery while reducing the impact of abnormal color shifts, oversaturation, or unreliable areas on subsequent color correction. The Self-Attention Threshold Control (STCM) module is used at the front end of the Lab and HSV branches. The Lab space provides brightness and color difference information that better matches visual perception, while the HSV space provides hue, saturation, and lightness information. After preprocessing these color space channels through the STCM, the subsequent JFDM module receives not the original, potentially redundant or biased color features, but rather effective color compensation features filtered by self-attention.

[0056] Figure 5A flowchart for generating color distillation features at the corresponding scale based on the feature joint distillation module JFDM is shown below. Figure 5 As shown, the RGB, Lab, and HSV color space features at the corresponding scale are simultaneously used as input. The feature joint distillation module uses an attention mechanism to automatically determine which information from the Lab and HSV features is suitable for supplementing the RGB main color branch, generating color distillation features at the corresponding scale. Figure 3 JFDM1 (Feature Joint Distillation Module 1), JFDM2 (Feature Joint Distillation Module 2), and JFDM3 (Feature Joint Distillation Module 3) correspond to the fusion of color features at three different downsampling scales. The modules first concatenate the RGB, Lab, and HSV feature sets, then generate a gated attention map through convolution, activation, and pixel attention structures. This attention map can be understood as a feature selection switch, determining which supplementary information from Lab and HSV should be introduced into the RGB main branch. Subsequently, Lab and HSV features are multiplied point-to-point with this attention map to filter out effective color compensation components, which are then added point-to-point with the RGB features to form the output features after joint distillation. The design purpose of JFDM is to absorb the advantageous information of the Lab and HSV color spaces while maintaining the stability of the RGB backbone representation. RGB features retain the basic image structure and conventional color expression, Lab features provide color difference information closer to human visual perception, and HSV features provide more intuitive hue, saturation, and brightness information. Simply concatenating these features directly might introduce redundant or even erroneous color interference. JFDM, however, uses attention gating for filtering, allowing the network to automatically learn "which color space information to borrow, and where to borrow it." In SAFCNet, JFDM is embedded in the multi-scale encoding-decoding process of the low-frequency branch. Shallow scales focus more on local color and edge-related changes, while deeper scales focus more on global hue and wide-area illumination distribution. By performing joint distillation at multiple scales, the network can simultaneously improve local color details and overall color consistency, thus more effectively correcting common problems in underwater images such as blue-green cast, insufficient brightness, and abnormal saturation.

[0057] Figure 6 The structure diagram of the Feature Fusion Optimization (FRM) module is as follows: Figure 6As shown, the input features of the feature fusion optimization module pass through FAB1 (first feature attention block), a convolutional layer, FAB2 (second feature attention block), another convolutional layer, FAB3 (third feature attention block), and another convolutional layer sequentially to obtain deep enhanced features. Simultaneously, a long skip connection is set to add the shallow fused features (the input features of the feature fusion optimization module) to the deep enhanced features, and finally, convolution and Tanh activation are applied to obtain the output. The feature fusion optimization module can further optimize the balance between color and texture while preserving the input structural information. The role of FRM is to unify and integrate the enhanced high-frequency texture features and low-frequency color features, and further correct problems such as discontinuities in details, color inconsistencies, or insufficient local enhancement that may occur during the fusion process. Structurally, FRM consists of three cascaded feature attention blocks (FABs) and multiple convolutional layers, while retaining a long skip connection directly connecting the input to the subsequent layers. The input features are processed sequentially through a sequence of "first feature attention block—first convolutional layer—second feature attention block—second convolutional layer—third feature attention block—third convolutional layer," and then added point-to-point with the original input features from the long skip branch. This residual learning structure reduces training difficulty, making it easier for the module to learn "the parts that need adjustment," rather than regenerating the entire image content. Each feature attention block (FAB) in the diagram can be understood as a local feature re-selection unit. The FAB contains cascaded spatial attention and pixel attention modules, working in conjunction with skip connections to improve learning efficiency. Spatial attention helps determine which areas in the image need enhancement, while pixel attention further refines the details to the pixel level, allowing the network to fine-tune the details after color and texture fusion. At the end of the module, the fused features continue to pass through convolutional layers for channel compression, ultimately transforming into a 3-channel image output, and then using the Tanh activation function to generate the final underwater image enhancement result. Overall, FRM does not simply add high-frequency and low-frequency results together, but rather optimizes them through attention, convolution, and residual learning, so that the final result has both natural color representation and clear texture details.

[0058] The above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit it. 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. Such 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 self-attention guidance and associated feature compensation, characterized in that, include: S1. Obtain the original image data, perform frequency domain decomposition on each original image in the original image data, and obtain the high-frequency components and low-frequency components of each original image; S2. Perform detail detection on the high-frequency components obtained in S1 to obtain the corresponding texture detail features. Perform color space conversion on the low-frequency components obtained in S1 to obtain the color space features of the corresponding color space. Detail detection includes edge detection and texture detection. Color spaces include RGB, Lab and HSV. S3. For high-frequency components, construct a texture detail feature compensation module. Input the high-frequency components obtained in S1 and the texture detail features obtained in S2 into the texture detail feature compensation module to obtain high-frequency enhancement results. S4. For low-frequency components, a color space feature compensation module is constructed. The low-frequency components obtained in S1 and the color space features obtained in S2 are input into the color space feature compensation module to obtain the low-frequency enhancement result. The color space feature compensation module includes a self-attention threshold control module and a feature joint distillation module. S5. Construct a feature fusion optimization module to fuse the high-frequency enhancement results obtained in S3 and the low-frequency enhancement results obtained in S4 to obtain the final enhancement result; S6. Construct a loss function to optimize the texture detail feature compensation module, color space feature compensation module, and feature fusion optimization module, and construct an underwater image enhancement network model. The loss function includes high-frequency loss, low-frequency loss, and pixel loss. High-frequency loss is used to optimize the model parameters of the texture detail feature compensation module, low-frequency loss is used to optimize the model parameters of the color space feature compensation module, and pixel loss is used to optimize the model parameters of the feature fusion optimization module. S7. The underwater image enhancement network model built on S6 performs image enhancement processing on the underwater image to be enhanced and outputs the final enhanced image.

2. The underwater image enhancement method based on self-attention guidance and associated feature compensation according to claim 1, characterized in that, S3 include: S3.1, based on the high-frequency components obtained in S1, the texture detail feature compensation module generates pixel attention maps of each texture detail feature, and the corresponding texture detail features are weighted and filtered based on each pixel attention map to generate multiple compensated texture feature maps and obtain multiple compensated texture features. S3.2, perform feature concatenation on the multiple compensation texture features obtained in S3.1, and extract features from the concatenated compensation texture features based on pixel attention and channel attention to obtain effective texture compensation features; S3.3, add the effective texture compensation features to the high-frequency components obtained in S1 to obtain the high-frequency enhancement result.

3. The underwater image enhancement method based on self-attention guidance and associated feature compensation according to claim 2, characterized in that, S4 include: S4.1, extract features from the color space features of Lab and HSV respectively, obtain the channel features of each channel of Lab and the channel features of each channel of HSV, input the channel features of each channel of Lab and the channel features of each channel of HSV into the self-attention threshold control module, generate the self-attention feature map of the corresponding color space according to the complementary characteristics of the values ​​of each channel, and adaptively enhance the channel features of the corresponding color space based on the respective attention feature map; S4.2 Perform multi-scale downsampling on the color space features of the three color spaces RGB, Lab and HSV respectively to obtain color space features at different scales of the corresponding color space. Construct a multi-scale color space feature set for the corresponding color space based on the color space features of each color space. Generate a joint attention map at each scale based on the color space features of the three color spaces at each scale. S4.3, input the multi-scale color space features of each color space obtained in S4.2 into the feature joint distillation module, filter the color space features of Lab and HSV according to the joint attention map obtained in S4.2, and superimpose the filtered color space features into the RGB color space features to obtain the color distillation features of the corresponding scale; S4.4 performs low-frequency decoding on the color distillation features obtained in S4.3 to obtain low-frequency enhancement results.

4. The underwater image enhancement method based on self-attention guidance and associated feature compensation according to claim 3, characterized in that, High frequency enhancement results for: ； ； In the formula, Represents high-frequency components, Indicates effective texture compensation features. This represents the stitched compensated texture features, which is equivalent to the feature obtained by stitching all the compensated texture features along the channel dimension. Represents the pixel attention module. This indicates the channel attention module; The compensated texture features are: ； ； In the formula, Index of the number of texture detail features. Indicates the first Individual texture features, express pixel attention map, Indicates correspondence Compensation texture features, This represents element-wise multiplication. This indicates a feature concatenation operation. This represents the pixel attention module.

5. The underwater image enhancement method based on self-attention guidance and associated feature compensation according to claim 4, characterized in that, The self-attention feature map of Lab is as follows: ； ； ； In the formula, This represents the channel index; for Lab, The value is , This represents the brightness channel of Lab. This indicates the red and green opposing channels of Lab. This represents the blue and yellow opposing channels of Lab. Belongs to, Indicates the first of Lab Self-attention feature map of each channel, Indicates the first of Lab Channel characteristics of each channel, express Self-attention graph.

6. The underwater image enhancement method based on self-attention guidance and associated feature compensation according to claim 5, characterized in that, For the saturation and brightness channels of HSV, the self-attention feature map is as follows: ； ； ； ； In the formula, This represents the self-attention feature map of the saturation channel of HSV. The channel characteristics representing the saturation channel of HSV. express Self-attention graph; This represents the self-attention feature map of the brightness channel of HSV. This indicates the channel characteristics of the lightness channel in HSV. express Self-attention graph; For the hue channels of HSV, angular metrics are used to represent color information, and a self-attention feature map is constructed based on the angular metrics: ； ； In the formula, A self-attention feature map representing the hue channel of HSV. This represents a segmented self-attention feature map of the hue channel in HSV. The channel characteristics of the hue channel in HSV are represented. Represents angular measurement.

7. The underwater image enhancement method based on self-attention guidance and associated feature compensation according to claim 6, characterized in that, In S4.3, the color distillation feature corresponding to the scale is: ； ； In the formula, Indicates scale index; Indicates the first Color distillation characteristics at various scales; Indicates the first Joint attention maps at multiple scales; , , They represent the first At each scale, the color space characteristics of RGB, Lab, and HSV are analyzed. This represents the activation function.

8. The underwater image enhancement method based on self-attention guidance and associated feature compensation according to claim 7, characterized in that, In S5, the feature fusion optimization module includes a long skip connection module, four convolutional layers, an activation function, and three cascaded feature attention blocks. Each feature attention block contains a recalibration structure to adjust the enhancement intensity of different channels and different spatial locations. The long skip connection module is used to implement residual connections, the convolutional layers are used to reduce the dimensionality of features, and the activation function is used to output the final enhanced image. The features from the input feature fusion optimization module are sequentially passed through the first feature attention block, the first convolutional layer, the second feature attention block, the second convolutional layer, the third feature attention block, and the third convolutional layer to obtain deep enhanced features. The deep enhanced features enter the long skip connection module, which adds the input features from the feature fusion optimization module to the deep enhanced features after the third convolutional layer element by element. The added features are then passed through the fourth convolutional layer and the activation function for output.

9. The underwater image enhancement method based on self-attention guidance and associated feature compensation according to claim 8, characterized in that, In S7, the loss function for: ； In the formula, Indicates high-frequency loss. Indicates low-frequency loss. Indicates pixel loss, , , They represent , , Loss weights; ； ； ； In the formula, Represents the number of pixel rows, Indicates the number of pixel columns. is the pixel coordinate index, indicating the . Line number Column of pixels, This indicates that the high-frequency enhancement results are in Pixel value at that location, This indicates that the high-frequency enhancement results are in Pixel reference value at that location, This indicates that the low-frequency enhancement results are in Pixel value at that location, This indicates that the low-frequency enhancement results are in Pixel reference value at that location, This indicates the final enhancement result in Pixel value at that location, This indicates the final enhancement result in Pixel reference value at that location, This represents summing the squares of the pixel values ​​at all pixel locations.

Images