An underwater image enhancement method based on attenuation-aware color correction and spatial-frequency cross-domain fusion

Abstract

The application provides an underwater image enhancement method based on attenuation-aware color correction and space-frequency cross-domain fusion, and relates to the field of underwater image enhancement, and comprises the following steps: obtaining an original underwater image, performing linear stretching and iterative compensation according to channel average intensity to obtain a color transfer image; taking a maximum attenuation map generated by normalizing the original image as a space weight, weighting and fusing the color transfer image and the original image pixel by pixel to obtain a color correction image; inputting the color correction image into an enhancement network, extracting and fusing spatial domain local texture features and frequency domain global structure features by using a space-frequency cross attention module; and enhancing feature expression through a multi-scale channel-space attention module to reconstruct and output an enhanced image. The application effectively solves the problems of blue-green color deviation, uneven correction and misalignment of space-frequency feature semantics of the underwater image.

Description

Technical Field

[0001] This invention relates to the technical field of underwater image enhancement, and more particularly to an underwater image enhancement method based on attenuation-sensory color correction and spatial-frequency cross-domain fusion. Background Technology

[0002] With the continuous development of tasks such as marine resource exploration, underwater target identification, marine engineering operations, and visual navigation for underwater robots, underwater imagery has become an important carrier for acquiring underwater environmental information. However, water selectively absorbs light of different wavelengths, and suspended particles cause scattering and noise interference, resulting in degradation phenomena in underwater images, such as blue-green tint, decreased contrast, blurred edges, and loss of texture details. Degraded underwater images not only reduce the quality of human visual observation but also affect the accuracy of subsequent visual tasks such as target detection, scene understanding, and equipment maintenance.

[0003] Existing underwater image enhancement methods mainly include traditional image processing methods, physics-based methods, and deep learning-based methods. Traditional image processing methods typically improve visual effects through histogram equalization, color correction, and multi-scale fusion, offering the advantage of low computational cost. However, they lack explicit modeling of underwater optical degradation mechanisms, easily leading to color distortion or local over-enhancement. Physics-based methods attempt to estimate parameters such as transmittance, background light, or depth through underwater imaging models and manual prior estimation, possessing a certain degree of physical interpretability. However, their performance is highly dependent on prior assumptions, and they lack robustness in complex water conditions, non-uniform lighting, and highly turbid environments.

[0004] In recent years, deep learning methods have been able to automatically learn the nonlinear mapping relationship between degraded images and clear images through convolutional neural networks, generative models, or semi-supervised learning frameworks, achieving good results in color correction, contrast enhancement, and detail restoration. However, most methods still mainly perform feature extraction and pixel mapping in the spatial domain, making insufficient use of information in the frequency domain that describes global structure, energy distribution, and high-frequency details. Although a few space-frequency combination methods introduce Fourier transform or wavelet transform, they often use static fusion methods such as splicing, element-wise addition, or fixed gating, making it difficult to achieve dynamic interaction between spatial and frequency features according to the degree of degradation in different regions. Summary of the Invention

[0005] To address the technical problems of color shift, low contrast, and blurred details in complex underwater scenes as described in the background art, this invention provides an underwater image enhancement method based on attenuation-aware color correction and spatial-frequency cross-domain fusion. This invention primarily utilizes an attenuation-aware color correction module to redefine channels, compensate colors, and guide fusion with a maximum attenuation map on the input underwater image. Then, an encoder-decoder enhancement network embedding a spatial-frequency cross-attention module and a multi-scale channel-spatial attention module is used to adaptively fuse local texture features in the spatial domain with global structural features in the frequency domain, thereby outputting an enhanced image with natural colors, clear structure, and rich details.

[0006] The technical means employed in this invention are as follows: An underwater image enhancement method based on attenuation-aware color correction and spatial-frequency cross-domain fusion includes: Acquire the original underwater image and normalize the original underwater image to obtain the normalized image; The average intensity of the red, green, and blue color channels of the original underwater image is obtained, and the average intensity is defined as large color channel, medium color channel, and small color channel in descending order of the mean value. A color loss is constructed based on the average intensity value, and the color loss is used to process the large, medium, and small color channels respectively to obtain a color transfer image; Based on the normalized image, the channel attenuation values ​​of the red, green, and blue color channels of the original image are obtained. The maximum attenuation value of the red, green, and blue color channels is obtained by selecting the maximum value of the channel attenuation value pixel by pixel. The maximum attenuation map is then constructed using the maximum attenuation value. Using the maximum attenuation map as the spatial weight, the color transfer image and the original underwater image are weighted and fused at each pixel location using the red, green, and blue color channels respectively to obtain a color correction image; The fused and enhanced feature representation is upsampled and reconstructed by convolution, a joint loss function is constructed, and the network is trained to obtain the trained network. The image to be enhanced is input into the trained network to obtain the enhanced image.

[0007] The encoder-decoder enhancement network includes a spatial-frequency cross-attention module. The spatial-frequency cross-attention module extracts the spatial domain local texture features and the frequency domain global structure features of the feature map, respectively. First, a spatial domain query matrix is ​​generated based on the spatial domain local texture features. Then, a frequency domain key matrix and a value matrix are generated based on the frequency domain global structure features. Finally, cross-attention calculation is performed based on the spatial domain query matrix and the frequency domain key matrix and value matrix to obtain the fused and enhanced feature representation. The spatial domain branch of the spatial-frequency cross-attention module includes a multi-scale channel-spatial attention module; the multi-scale features are obtained by multi-scale convolution of the input spatial domain local texture features through the multi-scale channel-spatial attention module, and the multi-scale features are weighted by both channel and spatial attention.

[0008] Furthermore, the underwater image enhancement dataset includes: the original underwater image and a reference image corresponding to the original underwater image.

[0009] Furthermore, the average intensity of the three color channels is calculated as follows:

[0010] in, Indicates channel The average strength; and These represent the height and width of the input image, respectively. Indicates color channels; Indicates color channel At pixel position Pixel value at; It represents the three color channels: red, green, and blue.

[0011] Furthermore, the color loss is expressed as:

[0012] in, Indicates color loss; , and These represent the average values ​​of the large, medium, and small color channels, respectively. The color loss is used to measure the difference in mean values ​​between different channels and serves as the basis for stopping the iterative color compensation process.

[0013] Furthermore, linear stretching is applied to the large color channel, and compensation is applied to the medium and small color channels, as expressed respectively:

[0014]

[0015]

[0016] in, , and These represent the compensated large color channel, medium color channel, and small color channel, respectively. Indicates pixel position; and These represent the minimum and maximum pixel values ​​of the major color channels, respectively. , and These represent the average intensity values ​​of the large, medium, and small color channels, respectively. and This represents the minimum and maximum stretch values. The above compensation process is performed iteratively until color loss occurs. Less than the preset threshold This yields a preliminary color transfer image.

[0017] Furthermore, the maximum attenuation map is a weight matrix, in which... The element is defined as The formula for calculating the maximum attenuation value and the maximum attenuation map is as follows:

[0018] in, express Maximum positional attenuation value; Indicates color channels; The parameter that controls the intensity of the decay response; This means normalizing the pixel values ​​of the input image to... ; It represents the three color channels: red, green, and blue.

[0019] Furthermore, pixel-wise weighted fusion is represented as:

[0020] Represents the pixel value of the corresponding channel in the color-transferred image; This represents the pixel value of the corresponding channel in the original image; This represents the pixel value after color correction; express Maximum positional attenuation value; This represents the red, green, and blue color channels. In areas of significant attenuation, a larger [color channel]... The color transfer result is given higher weight, thereby enhancing color compensation; in areas with weaker attenuation, more original image information is preserved to reduce color overcorrection and maintain the naturalness of local structures.

[0021] Furthermore, the space-frequency cross-attention module includes: a spatial domain feature extraction branch, a frequency domain feature enhancement branch, and a cross-attention fusion unit. The spatial domain feature extraction branch performs local convolution, normalization, and depthwise separable convolution on the input features to generate a query matrix; The frequency domain feature enhancement branch performs a fast Fourier transform on the global structural features in the frequency domain to separate the amplitude spectrum and phase spectrum. It then enhances the amplitude and phase information of the frequency domain feature enhancement branch by connecting a standard two-dimensional convolutional layer and a Leaky ReLU activation function in sequence. The frequency domain enhanced features of the frequency domain feature enhancement branch are obtained by inverse fast Fourier transform, and the frequency domain enhanced features are mapped to obtain the key matrix and value matrix in the frequency domain.

[0022] Furthermore, the cross-attention fusion unit is represented as:

[0023] in, Represents the spatial domain query matrix. and Represent the frequency domain key matrix and value matrix, respectively. Indicates the scaling factor. Indicates the characteristics of space-frequency fusion; express Activation function; express The transpose of .

[0024] The multi-scale channel-spatial attention module includes a multi-scale feature extraction unit, a channel attention unit, and a spatial attention unit. The multi-scale feature extraction unit captures multi-scale texture information through different convolutional kernels or different receptive field branches. The channel attention unit obtains channel descriptors and generates channel weights through global average pooling and global max pooling. The spatial attention unit generates region weights based on spatial location responses to highlight important structural and texture regions in underwater images.

[0025] Compared with the prior art, the present invention has the following advantages: 1. This invention introduces an attenuation-aware color correction module into the U-Net enhancement network, which can dynamically redefine and iteratively compensate the RGB channels according to the different light attenuation levels of different color channels in the underwater environment. This allows severely degraded areas to receive stronger compensation, while areas with weaker degradation retain more original information, effectively alleviating the problems of blue-green bias and uneven color correction.

[0026] 2. This invention introduces a spatial-frequency cross-attention module into the U-Net enhanced network. Compared with the method of performing convolution modeling only in the spatial domain, it can use the amplitude spectrum and phase spectrum in the frequency domain to describe the global structure, energy distribution and high-frequency details. Furthermore, it achieves dynamic interaction between local texture features in the spatial domain and global structural features in the frequency domain through cross-attention, thereby alleviating the problem of semantic misalignment of spatial-frequency features.

[0027] 3. This invention introduces a multi-scale channel-spatial attention module by adding a spatial domain branch to the spatial frequency cross-attention module. By combining multi-scale feature extraction, channel weight recalibration, and spatial region weighting, the network's ability to perceive key channels, important regions, and details at different scales is improved.

[0028] In summary, the technical solution of this invention can simultaneously improve the naturalness of image color, the clarity of structure, and the ability to retain details in complex underwater environments, providing high-quality image input for subsequent tasks such as underwater target detection, environmental perception, and robot visual navigation. Therefore, the technical solution of this invention solves the problems of blue-green tint, decreased contrast, blurred edges, and loss of texture details commonly found in existing underwater images.

[0029] Based on the above reasons, this invention can be widely applied in general image data processing or generation fields. Attached Figure Description

[0030] 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.

[0031] Figure 1 This is a schematic diagram of the overall structure of the AC-SFNet underwater image enhancement network in Embodiment 1 of the present invention.

[0032] Figure 2 This is a logic block diagram of the color correction module in Embodiment 1 of the present invention.

[0033] Figure 3 This is a comparison diagram of the color correction stage effect in Embodiment 1 of the present invention; wherein, (a) row represents the original image; (b) row represents the color-transferred image; and (c) row represents the color-corrected image.

[0034] Figure 4 This is a structural diagram of the space-frequency cross-attention module in Embodiment 1 of the present invention.

[0035] Figure 5 This is a structural diagram of the multi-scale channel-spatial attention module in Embodiment 1 of the present invention.

[0036] Figure 6 This is a comparison chart of the enhancement effects of different methods on a reference dataset in Embodiment 2 of the present invention; wherein, (a) row represents the test set Tset-UIEB; (b) row represents the test set Tset-EUVP; and (c) row represents the test set Test-LSUI.

[0037] Figure 7 This is a comparison chart of the enhancement effects of different methods on a no-reference dataset in Embodiment 2 of the present invention; wherein, (a) row represents the test set Tset-C60; (b) row represents the test set U45; and (c) row represents the test set UCCS.

[0038] Figure 8 This is a visualization of the ablation experiment results of the model structure in Embodiment 3 of the present invention; wherein, Figure 8 .1 represents the original image; Figure 8 .2 indicates the result of the ablation color correction module; Figure 8 .3 indicates the result of ablation of the spatial frequency cross-attention module; Figure 8 .4 indicates the result of ablation of the multi-scale channel-spatial attention module; Figure 8 .5 indicates the result of the complete model; Figure 8 .6 represents a reference image.

[0039] Figure 9 This is a visualization of the loss function ablation experiment results in Embodiment 3 of the present invention; wherein, Figure 9 .1 represents the original image; Figure 9 .2 indicates ablation The result; Figure 9 .3 indicates ablation The result; Figure 9 .4 indicates ablation The result; Figure 9 .5 indicates ablation The result; Figure 9 .6 represents the result of the full loss function; Figure 9 .7 indicates a reference image. Detailed Implementation

[0040] 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.

[0041] 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.

[0042] This invention provides an underwater image enhancement method based on attenuation-aware color correction and spatial-frequency cross-domain fusion, comprising an attenuation-aware color correction stage and a spatial-frequency cross-domain fusion enhancement stage. The attenuation-aware color correction stage redefines channels, compensates for colors, and guides fusion with a maximum attenuation map based on underwater light attenuation patterns in the original underwater image. The spatial-frequency cross-domain fusion enhancement stage fuses local texture features in the spatial domain and global structural features in the frequency domain within the encoder-decoder network using a spatial-frequency cross-attention module and a multi-scale channel-spatial attention module, thereby generating an enhanced underwater image.

[0043] In one optional embodiment, the processing flow of the present invention includes: acquiring the original underwater image and performing normalization processing; acquiring the average intensity of the red, green, and blue color channels of the original underwater image; constructing a color loss and acquiring a color transfer image; constructing a maximum attenuation map using the maximum attenuation value; performing weighted fusion pixel by pixel and acquiring a color-corrected image; constructing a spatial-frequency cross-attention module and inputting the color-corrected image into the spatial-frequency fusion enhancement network for image enhancement; constructing a multi-scale channel-spatial attention module; constructing a joint loss function and training the network to perform underwater image enhancement inference.

[0044] Example 1 like Figure 1 As shown, this invention provides an underwater image enhancement method based on attenuation-aware color correction and spatial-frequency cross-domain fusion, comprising the following steps: Step 1: Acquire the original underwater image and perform normalization processing.

[0045] We collected several publicly available underwater image augmentation datasets, which included original underwater images and corresponding reference or artificially augmented images. We then uniformly converted the underwater images to... Three-channel color images are used, and the original and reference images are subjected to the same size normalization, random cropping, horizontal flipping, or other data augmentation operations to obtain underwater image sample pairs for training.

[0046] During the training phase, the training samples are adjusted to a preset resolution, such as 256×256 pixels. During the testing phase, the images to be enhanced are input into the trained network model for inference. The underwater image dataset covers different water colors, turbidity levels, lighting conditions, and degradation degrees to improve the network's generalization ability in real underwater scenarios.

[0047] Step 2: Obtain the average intensity of the red, green, and blue color channels of the original underwater image.

[0048] Given an input underwater image First, normalize its pixel values ​​to Perform standardization processing. Let the height of the input image be... Width is Color channels Belongs to the red channel Green Channel and the blue channel The average intensity of each of the three color channels is calculated to determine the degree of attenuation of different color channels in the current underwater image.

[0049] The average intensity of a color channel is expressed as:

[0050] in, express Channel average intensity; and These represent the height and width of the input image, respectively. Indicates color channels; Indicates color channel At pixel position Pixel value at; This represents the red, green, and blue color channels. Based on the relationship between the average intensities of the three channels, the channel with the largest average value is redefined as the major color channel. Redefine the channel with the median mean as the middle color channel. The channel with the smallest mean is redefined as the smaller color channel. .

[0051] Step 3: Construct color loss and obtain color transfer image.

[0052] Because different wavelengths of light attenuate at varying degrees in underwater environments, some color channels are often overemphasized while others are lacking. To reduce the mean differences between different color channels, this embodiment constructs a color loss function based on the differences between the mean of the large color channel and the mean of the medium and small color channels. The color loss can be expressed as:

[0053] in, Indicates color loss; , and These represent the average values ​​of the large, medium, and small color channels, respectively.

[0054] Based on the aforementioned color loss, a linear stretching is first performed on the large color channel to expand its dynamic range. Then, based on the minimum and maximum pixel values ​​of the large color channel, the channel is mapped to a preset output range, for example... or normalized range Subsequently, based on the mean differences between the large color channel and the medium and small color channels, compensation is applied to the medium and small color channels to make the average intensity of different channels more consistent. Linear stretching is performed on the large color channel, and compensation is applied to the medium and small color channels, as expressed respectively:

[0055]

[0056]

[0057] in, , and These represent the compensated large color channel, medium color channel, and small color channel, respectively. Indicates pixel position; and These represent the minimum and maximum pixel values ​​of the major color channels, respectively. , and These represent the average intensity values ​​of the large, medium, and small color channels, respectively. and This represents the minimum and maximum stretch values. The above compensation process is performed iteratively until color loss occurs. Less than the preset threshold This process yields a preliminary color-transferred image. By employing this method, the blue-green bias can be initially reduced before entering the deep network, and the severely degraded color components in the underwater image can be compensated for.

[0058] Step 4: Construct a maximum attenuation map using the maximum attenuation value.

[0059] To further address spatially non-uniform color degradation in underwater images, this embodiment constructs an attenuation map for each color channel based on the normalized original image. The channel attenuation map describes the local attenuation response of that color channel at a corresponding location. After calculating the attenuation maps for the red, green, and blue color channels, the maximum attenuation value among the three channels is taken for each pixel location, denoted as . This yields the maximum attenuation map. The maximum attenuation map is used to approximately characterize the degree of color degradation at different spatial locations. When the attenuation value of a certain region is large, it indicates that the color degradation in that region is more severe and stronger color compensation is needed. When the attenuation value of a certain region is small, it indicates that the region retains more original information and is not suitable for overcorrection.

[0060] In a preferred embodiment, the maximum attenuation map is defined as a weight matrix, in which... The element is called the maximum decay value. The formula for calculating the maximum attenuation value is:

[0061] in, express Maximum positional attenuation value; Indicates color channels; The parameter that controls the intensity of the decay response; This means normalizing the pixel values ​​of the input image to... ; This represents the three color channels: red, green, and blue. These are parameters used to control the intensity of the decay response. Experimental verification has shown that... The effect is best when the value is 0.7.

[0062] Step 5: Perform weighted fusion pixel by pixel and obtain the color-corrected image.

[0063] Using the maximum attenuation map as spatial weights, the color-transferred image and the original underwater image are fused pixel-by-pixel with weighted fusion to obtain the final color-corrected image. For each color channel... The fusion result can be expressed as:

[0064] in, This represents the pixel value of the corresponding channel in the color-transferred image. This represents the pixel value of the corresponding channel in the original image. This represents the pixel value after color correction. express Maximum positional attenuation value; It represents the three color channels: red, green, and blue.

[0065] By using maximum attenuation map-guided fusion, severely attenuated areas receive higher color transfer image weights, thus enhancing color compensation; areas with less attenuation retain more original image information, thereby avoiding overcorrection and local color distortion. This process enables color correction results to reduce blue-green bias while maintaining more natural regional color gradations.

[0066] Step 6: Construct a spatial-frequency cross-attention module and input the color-corrected image into the spatial-frequency fusion enhancement network for image enhancement.

[0067] The color-corrected image obtained in step 5 is input into the spatial-frequency fusion enhancement network. The spatial-frequency fusion enhancement network adopts an encoder-decoder structure, including at least one downsampling encoding stage, at least one upsampling decoding stage, and skip connections between the encoder and decoder. The encoder is used to extract texture, edge, and structural features at different levels; the decoder is used to progressively restore the image spatial resolution and reconstruct the enhanced image.

[0068] Spatial-frequency cross-attention modules are embedded in multiple feature levels of the encoder and decoder, enabling the network to extract not only local texture features in the spatial domain but also explicitly incorporate global structure and high-frequency detail information in the frequency domain. Skip connections are used to pass shallow spatial details to the decoder to reduce detail loss caused by downsampling. Finally, the decoder outputs an enhanced underwater image with the same size as the input.

[0069] The spatial-frequency cross-attention module includes a spatial domain branch, a frequency domain branch, and a cross-attention fusion unit. The input feature map is first fed into the spatial domain branch and the frequency domain branch respectively to obtain local features in the spatial domain and global features in the frequency domain.

[0070] In the spatial domain branch, the input features first pass through an improved Inception feature extraction structure and a multi-scale channel-spatial attention module to capture texture, edge, and region structures at different scales. Subsequently, the spatial domain features undergo layer normalization, regular convolution, and depthwise separable convolution to obtain the spatial domain query matrix. The query matrix is ​​used to represent the selection requirements of frequency domain information for spatial location and local texture.

[0071] In the frequency domain branch, a Fast Fourier Transform is performed on the input features to transform the feature map from the spatial domain to the frequency domain, yielding the amplitude spectrum and phase spectrum. The amplitude spectrum primarily reflects the overall intensity, frequency energy, and contrast distribution of the image; the phase spectrum mainly preserves spatial structure, edge positions, and target contour information.

[0072] Learnable convolutional layers were applied to both the amplitude and phase spectra to enhance them, compensating for spectral distribution anomalies caused by underwater scattering and absorption. The enhanced amplitude and phase spectra were then recombined and recovered as spatial frequency enhancement features using an inverse fast Fourier transform. These frequency enhancement features were further mapped to obtain the bond matrix. Sum matrix .

[0073] The cross-attention fusion unit queries the matrix based on the spatial domain. With frequency domain key matrix The correlation between them, adaptively select useful structural information in the frequency domain, and adjust the frequency domain value matrix. Perform weighted aggregation. The cross-attention output can be represented as:

[0074] in, Represents the spatial domain query matrix. and Represent the frequency domain key matrix and value matrix, respectively. Indicates the scaling factor. Indicates the characteristics of space-frequency fusion. express Activation function; express The transpose of .

[0075] Through the aforementioned cross-attention mechanism, local texture features in the spatial domain can selectively aggregate global structural information in the frequency domain based on the current image content, thereby alleviating the semantic misalignment problem caused by simple concatenation or element-wise addition. The cross-attention output is then residually connected to the module input to preserve the original features and stabilize network training.

[0076] Step 7: Construct a multi-scale channel-spatial attention module.

[0077] The multi-scale channel-spatial attention module is used to enhance the feature representation of key channels and important regions in the spatial domain branch. This module includes a multi-scale feature extraction unit, a channel attention unit, and a spatial attention unit.

[0078] In the multi-scale feature extraction unit, a parallel branching structure is used to process the input features. Different branches can use convolutional kernels of different sizes, different receptive fields, or pooling operations to simultaneously capture fine-grained edge textures, mid-scale region structures, and a wider range of contextual information. The outputs of each branch are concatenated or fused along the channel dimension to form a multi-scale feature representation.

[0079] In the channel attention unit, global average pooling and global max pooling are performed on the multi-scale features to obtain channel statistical descriptors. Channel weights are then generated through convolutional or fully connected mappings and applied to the input features to highlight effective channels related to color restoration, edge preservation, and texture enhancement, while suppressing redundant channel responses.

[0080] In the spatial attention unit, spatial region weights are generated based on the strength of the response of the feature map at spatial locations. These weights are used to highlight target boundaries, coral textures, fish outlines, seabed structures, and other important areas in underwater images, while suppressing background noise, murky areas, and invalid textures. Through the joint weighting of channel attention and spatial attention, the network can better adapt to the non-uniform degradation features of underwater images.

[0081] Step 8: Construct the joint loss function and train the network.

[0082] To ensure the network simultaneously achieves high pixel reconstruction accuracy, structural consistency, visual perception quality, and frequency distribution consistency, this embodiment employs a joint loss function to train the space-frequency fusion enhancement network. The joint loss function includes pixel reconstruction loss, structural loss, perceptual loss, and frequency domain loss.

[0083] Among them, pixel reconstruction loss is used to constrain the differences between the enhanced image and the reference image at the pixel level; structural loss is used to constrain the consistency between the enhanced image and the reference image in terms of structural similarity; perceptual loss uses a pre-trained visual feature extraction network to constrain the differences in high-level semantics and texture perception; and frequency domain loss uses fast Fourier transform to constrain the differences between the enhanced image and the reference image in amplitude spectrum and phase spectrum, respectively.

[0084] In one alternative implementation, the joint loss function can be expressed as: ; , , , These represent the weights of pixel reconstruction loss, structural loss, perceptual loss, and frequency domain loss, respectively. Preferably, these weights can be set to 1.0, 0.15, 0.6, and 0.1, respectively.

[0085] The Adam optimization algorithm is used during training. Optionally, the initial learning rate is set to... Batch size set to 8, training epochs set to 120; adjust all training images to... Pixels; convolution weights are initialized using a Gaussian distribution. After training, network parameters with smaller loss values ​​during validation or training are selected as the final model parameters.

[0086] Step 9: Obtain the enhanced image from the trained network.

[0087] The original underwater image to be enhanced is input into the trained underwater image enhancement model. First, the attenuation-aware color correction module outputs a color-corrected image, and then the spatial-frequency fusion enhancement network performs structure and detail restoration. The network output is constrained by pixel range to obtain the final enhanced image. The final enhanced image has more natural colors, higher contrast, clearer edge structure, and richer texture details compared to the original underwater image.

[0088] Example 2 like Figure 6 and Figure 7 As shown in Example 1, the present invention also provides an underwater image enhancement method based on attenuation-aware color correction and spatial-frequency cross-domain fusion to illustrate the present invention in the underwater image enhancement training and testing process. The comparison diagram of the enhancement effects of different methods on reference datasets and non-reference datasets is shown.

[0089] Step 1: Collect and segment underwater image data. 800 pairs of paired images from the UIEB underwater image augmentation benchmark dataset are selected as the training set, and the remaining 90 pairs are used as the reference-based test set Test-UIEB. Further, 90 pairs of paired images are selected from the EUVP and LSUI datasets respectively to construct Test-EUVP and Test-LSUI, used to test cross-dataset generalization ability. Simultaneously, real-world, unreferenced underwater image datasets such as Test-C60, U45, and UCCS are selected to test the visual augmentation effect in real-world scenarios.

[0090] Step 2: Preprocess the training images. Adjust the original underwater images and reference images in the training set to be uniformly... The pixels are processed and normalized, randomly cropped, flipped, or otherwise enhanced as needed for network training. For each input image, attenuation-aware color correction is first performed to generate a color-corrected image, which is then used as input to the spatial-frequency fusion enhancement network.

[0091] Step 3: Train the AC-SFNet underwater image enhancement network. During training, the color correction module first calculates the channel mean, color loss, and maximum attenuation map of the input image to generate a color-compensated image. Then, the U-Net encoder extracts multi-layer features, embedding spatial-frequency cross-attention modules in multiple layers to fuse spatial and frequency domain features. Simultaneously, a multi-scale channel spatial attention module enhances key channels and important spatial regions. Finally, the decoder restores the image resolution and outputs the enhancement result.

[0092] Step 4: Set training parameters. In this embodiment, the network training epochs are set to 120 epochs, the optimizer is Adam, and the initial learning rate is set to... The batch size is set to 8. Training can be performed using a single or multiple GPUs, depending on the hardware requirements. After training, save the model parameters with the best enhancement or lowest loss value for subsequent testing.

[0093] Step 5: Conduct reference testing. Input the original underwater images from Test-UIEB, Test-EUVP, and Test-LSUI into the trained model to obtain augmented images. The differences between the augmented results and the reference images are evaluated using Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), and Mean Squared Error (MSE). Higher PSNR and SSIM values ​​indicate better reconstruction quality and structural consistency of the augmented image, while lower MSE values ​​indicate smaller pixel errors.

[0094] Step 6: Conduct no-reference testing. Input real underwater images from Test-C60, U45, and UCCS into the trained model to obtain enhanced images. The perceived quality of the real underwater images is evaluated using the Underwater Image Quality Evaluation Metric (UIQM) and the Underwater Color Image Quality Evaluation Metric (UCIQE). Higher UIQM and UCIQE values ​​indicate better overall quality in terms of color, sharpness, and contrast.

[0095] Step 7: Compare with other underwater image enhancement methods. Traditional prior methods, supervised deep learning methods, semi-supervised methods, generative methods, and space-frequency co-enhancement methods are selected as comparison methods, such as IBLA, Water-Net, FUnIE-GAN, Semi-UIR, FDCE-Net, CDF-UIE, and SFG-Net. The enhancement performance of the proposed method is comprehensively evaluated using reference metrics, no reference metrics, and visualization results.

[0096] Step 8: Perform visualization analysis. With a reference test set, observe whether the enhanced image can reduce the blue-green bias and restore structural details such as divers, fish, corals, and seabed textures. In real underwater images without a reference, observe whether the enhancement result has problems such as oversaturation, color drift, local over-enhancement, or structural damage. If the enhancement result is superior to or close to the comparison method in terms of color naturalness, background transparency, target boundary clarity, and texture sharpness, it indicates that the method of this invention has a good practical enhancement effect.

[0097] Example 3 like Figure 8 and Figure 9As shown in Example 1, the present invention also provides visualization results of model structure ablation experiments and loss function ablation experiments to verify the effectiveness of each component module of the present invention.

[0098] To verify the effectiveness of the attenuation-aware color correction module, the spatial-frequency cross-attention module, and the multi-scale channel-spatial attention module, several variant models were constructed under the same training data, training parameters, and loss function. The first variant removed the attenuation-aware color correction module and directly input the original underwater image into the spatial-frequency fusion enhancement network; the second variant removed the spatial-frequency cross-attention module and used only the spatial domain enhancement network; the third variant removed the multi-scale channel-spatial attention module; the complete model includes the color correction module, the spatial-frequency cross-attention module, and the multi-scale channel-spatial attention module simultaneously.

[0099] like Figure 8 As shown, where, Figure 8 .1 represents the original image; Figure 8 .2 indicates the result of the ablation color correction module; Figure 8 .3 indicates the result of ablation of the spatial frequency cross-attention module; Figure 8 .4 indicates the result of ablation of the multi-scale channel-spatial attention module; Figure 8 .5 indicates the result of the complete model; Figure 8 .6 represents the reference image. The comparison shows that removing the attenuation-aware color correction module easily leaves residual blue-green bias and inconsistent regional colors in the enhanced image, indicating that the color correction module can improve the color distribution of the input image and provide a more stable color foundation for subsequent networks. Removing the spatial-frequency cross-attention module significantly weakens the target contour, texture levels, and local contrast, indicating that the spatial-frequency cross-attention module plays an important role in global structure recovery and local detail preservation. Removing the multi-scale channel-spatial attention module, although color correction still exists, the recovery of fish body boundaries, rock textures, and fine structures is insufficient, indicating that the multi-scale channel-spatial attention module can enhance the expressive ability of degradation features at different scales.

[0100] Furthermore, such as Figure 9 As shown, where, Figure 9 .1 represents the original image; Figure 9 .2 indicates ablation The result; Figure 9 .3 indicates ablation The result; Figure 9 .4 indicates ablation The result; Figure 9 .5 indicates ablation The result; Figure 9 .6 represents the result of the full loss function; Figure 9.7 represents the reference image. Ablation analysis was performed on the joint loss function. Pixel reconstruction loss, structural loss, perceptual loss, and frequency domain loss were removed separately, while maintaining the network structure, training data, and other training parameters. Experiments show that pixel reconstruction loss ensures basic consistency between the enhancement result and the reference image; structural loss helps maintain image structure; perceptual loss improves texture appearance and visual consistency; and frequency domain loss constrains the frequency distribution of the image at the amplitude and phase spectrum levels, thus supplementing the spatial domain reconstruction constraints.

[0101] In summary, this invention achieves a complete processing flow from color compensation and structure restoration to detail enhancement through the synergistic effects of attenuation-perceptual color correction, spatial-frequency cross-attention, and multi-scale channel-spatial attention. This method is applicable to underwater image enhancement scenarios with blue-green bias, low contrast, blurred edges, missing textures, and spatial non-uniform degradation, and can provide high-quality input for subsequent visual tasks such as underwater target detection, marine environment perception, underwater robot navigation, and marine engineering inspection.

[0102] 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.

[0103] In the above embodiments of the present invention, the descriptions of each embodiment have different focuses. For parts not described in detail in a certain embodiment, please refer to the relevant descriptions of other embodiments.

[0104] 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 attenuation-aware color correction and spatial-frequency cross-domain fusion, characterized in that, Includes the following steps: Step 1: Acquire the original underwater image and normalize it to obtain the normalized image; Step 2: Obtain the average intensity of the red, green and blue color channels of the original underwater image, and define the average intensity as large color channel, medium color channel and small color channel in descending order of the mean value; Step 3: Construct a color loss based on the average intensity value, and process the large, medium and small color channels respectively using the color loss to obtain a color transfer image; Step 4: Obtain the channel attenuation values ​​of the red, green, and blue color channels of the original image based on the normalized image, and obtain the maximum attenuation value of the red, green, and blue color channels by selecting the maximum value of the channel attenuation value pixel by pixel for the red, green, and blue color channels, and then construct the maximum attenuation map based on the maximum attenuation value. Step 5: Using the maximum attenuation map as the spatial weight, perform weighted fusion of the red, green, and blue color channels at each pixel position on the color transfer image and the original underwater image to obtain a color-corrected image; Step 6: Input the color-corrected image into the encoder-decoder enhancement network to obtain the fused and enhanced feature representation; The encoder-decoder enhancement network includes a space-frequency cross-attention module. The space-frequency cross-attention module extracts the spatial domain local texture features and the frequency domain global structure features of the feature map respectively. First, a spatial domain query matrix is ​​generated based on the spatial domain local texture features. Then, a frequency domain key matrix and a value matrix are generated based on the frequency domain global structure features. Finally, cross-attention calculation is performed based on the spatial domain query matrix, the frequency domain key matrix, and the value matrix to obtain the fused and enhanced feature representation. Step 7: The spatial domain branch of the space-frequency cross-attention module includes: multi-scale channel-spatial attention module; The multi-scale features are obtained by multi-scale convolution extraction of the input spatial domain local texture features through the multi-scale channel-spatial attention module, and the multi-scale features are weighted by dual channel and spatial attention. Step 8: Upsample and reconstruct the fused and enhanced feature representations, construct a joint loss function to complete network training, and obtain a converged target model; Step 9: Input the image to be enhanced into the acquired target model to obtain the enhanced image.

2. The underwater image enhancement method based on attenuation-sensory color correction and spatial-frequency cross-domain fusion according to claim 1, characterized in that, In step 1, the underwater image enhancement dataset includes: the original underwater image and the enhanced reference image corresponding to the original underwater image.

3. The underwater image enhancement method based on attenuation-sensory color correction and spatial-frequency cross-domain fusion according to claim 1, characterized in that, In step 2, the formula for calculating the average intensity mean of any channel is: in, Indicates channel The average strength; and These represent the height and width of the input image, respectively. Indicates color channels; Indicates color channel At pixel position Pixel value at; It represents the three color channels: red, green, and blue.

4. The underwater image enhancement method based on attenuation-sensory color correction and spatial-frequency cross-domain fusion according to claim 1, characterized in that, In step 3, the formula for calculating the color loss is: in, Indicates color loss; , and These represent the average intensity values ​​of the large, medium, and small color channels, respectively.

5. The underwater image enhancement method based on attenuation-aware color correction and spatial-frequency cross-domain fusion according to claim 1, characterized in that, In step 3, linear stretching is performed on the large color channel, and compensation is applied to the medium and small color channels, as expressed as follows: in, , and These represent the compensated large color channel, medium color channel, and small color channel, respectively. Indicates pixel position; and These represent the minimum and maximum pixel values ​​of the major color channels, respectively. , and These represent the average intensity values ​​of the large, medium, and small color channels, respectively. and This indicates the minimum and maximum stretch values.

6. The underwater image enhancement method based on attenuation-sensory color correction and spatial-frequency cross-domain fusion according to claim 1, characterized in that, In step 4, the maximum attenuation map is used to approximately characterize the degree of color degradation at different spatial locations. The maximum attenuation map is a weight matrix, in which... The element is defined as The maximum attenuation value is calculated using the following formula: in, express Maximum positional attenuation value; Indicates color channels; The parameter that controls the intensity of the decay response; This means normalizing the pixel values ​​of the input image to... ; It represents the three color channels: red, green, and blue.

7. The underwater image enhancement method based on attenuation-aware color correction and spatial-frequency cross-domain fusion according to claim 1, characterized in that, In step 5, the pixel-wise weighted fusion is represented as: in, This represents the pixel value of the corresponding channel in the color-transferred image. This represents the pixel value of the corresponding channel in the original image. This represents the pixel value after color correction. express Maximum positional attenuation value It represents the three color channels: red, green, and blue.

8. The underwater image enhancement method based on attenuation-sensory color correction and spatial-frequency cross-domain fusion according to claim 1, characterized in that, In step 6, the spatial-frequency cross-attention module includes: a spatial domain feature extraction branch, a frequency domain feature enhancement branch, and a cross-attention fusion unit; the spatial domain feature extraction branch performs local convolution, normalization, and depth-separable convolution processing on local texture features in the spatial domain to obtain a spatial domain query matrix. The space-frequency cross-attention module is represented as follows: in, Represents a spatial domain query matrix; and These represent the frequency domain key matrix and value matrix, respectively. Indicates the scaling factor; Indicates the characteristics of space-frequency fusion; express Activation function; express Transpose of; The frequency domain feature enhancement branch performs a fast Fourier transform on the global structural features in the frequency domain to separate the amplitude spectrum and phase spectrum. It then enhances the amplitude and phase information of the frequency domain feature enhancement branch by connecting a standard two-dimensional convolutional layer and a Leaky ReLU activation function in sequence. The frequency domain enhanced features of the frequency domain feature enhancement branch are obtained by inverse fast Fourier transform, and the frequency domain enhanced features are mapped to obtain the key matrix and value matrix in the frequency domain.

9. The underwater image enhancement method based on attenuation-aware color correction and spatial-frequency cross-domain fusion according to claim 1, characterized in that, In step 7, the multi-scale channel-spatial attention module includes: a multi-scale feature extraction unit, a channel attention unit, and a spatial attention unit; the multi-scale feature extraction unit obtains multi-scale texture information through different convolutional kernels or different receptive field branches, the channel attention unit obtains channel descriptors and generates channel weights through global average pooling and global max pooling, and the spatial attention unit generates region weights based on spatial location responses.

Images