Underwater image enhancement method based on double-domain attention U-Net

Abstract

The invention provides an underwater image enhancement method based on a U-Net backbone network and fusing a position sensing module, a double-branch attention bridge and a DCT frequency domain enhancement module, and aims to solve the problems of low contrast ratio, color distortion, detail loss and the like caused by absorption and scattering effects in the light propagation process of an underwater image, improve the image quality and improve the image quality. And the visual perception effect is improved. The method specifically comprises the following four main steps: firstly, acquiring and preprocessing initial data sets of different underwater scenes; secondly, extracting spatial relation information of the image through a position sensing convolution module, and enhancing image details; thirdly, a double-branch attention bridge module is used for further integrating local details and a global structure, and the response capability of a key area is enhanced; and finally, in combination with RGB and NIR modal features, weighted fusion of frequency-space features is realized through a DCT fusion enhancement module, so that detail information of the image is recovered, color deviation is corrected, and the definition and visual effect of the underwater image are remarkably improved. According to the method, the enhancement performance of the underwater image in a complex environment can be improved, and powerful support is provided for applications such as underwater detection, target detection, object recognition and image segmentation.

Description

TECHNICAL FIELD

[0001] The present application relates to computer vision and image processing technology, and is an underwater image enhancement method based on a dual-domain attention U-Net. The method is mainly applied to the quality improvement of underwater images, especially in aspects such as detail recovery, color correction and texture enhancement in complex underwater environments, and has important practical application value. BACKGROUND

[0002] The ocean covers about 70% of the earth's surface and contains rich biological and mineral resources, playing an important role in ecological balance and human sustainable development. With the increasing demand for deep-sea exploration, resource development and underwater engineering operations, underwater image enhancement (UIE) technology has become an important support for autonomous underwater vehicles (AUVs) and remotely operated vehicles (ROVs) to achieve navigation, target recognition and environmental perception.

[0003] Due to the characteristics of underwater light propagation, underwater imaging generally suffers from quality degradation. Specifically, long-wave light (such as red light) quickly attenuates, resulting in the loss of warm colors; short-wave light (such as blue-green light) travels far, but is easily scattered, causing color distortion. At the same time, suspended particles, dissolved substances and plankton in the water further exacerbate scattering and absorption, causing color distortion, reduced contrast, blurred details and increased noise in the image, severely affecting the perception and decision-making ability of the vision system.

[0004] To address the above problems, existing methods include contrast enhancement, color correction, frequency domain detail recovery and image fusion, but most rely on a single prior, lack of adaptability to complex underwater environments, and have insufficient robustness. With the development of deep learning, end-to-end image enhancement methods have achieved some success in color restoration, dehazing and detail repair. However, due to the complexity of underwater imaging degradation mechanisms, existing methods still struggle to fully address color distortion, blurring and detail loss in various water qualities. Improving the color restoration, structural clarity and detail integrity of underwater images remains a technical challenge that needs to be addressed. High-quality images are crucial for tasks such as target detection, object recognition and image segmentation, and directly affect the intelligence level and efficiency of underwater equipment.

[0005] To address the above challenges, the present application proposes an underwater image enhancement method based on a U-Net backbone network, which integrates a position perception module, a dual-branch attention bridge and a DCT frequency domain enhancement module. The method aims to improve the image enhancement performance of the model in complex underwater environments, effectively address underwater image degradation problems, and improve the visual quality and detail performance of the image. SUMMARY

[0006] In view of this, the present application proposes an underwater image enhancement method based on a U-Net backbone network, which fuses a position perception module, a double-branch attention bridge and a DCT frequency domain enhancement module. The method can improve the detail information of degraded underwater images, correct color deviation and improve visual perception effect.

[0007] The technical solution adopted by the present application to solve the technical problems is as follows:

[0008] An underwater image enhancement method based on a U-Net backbone network, which fuses a position perception module, a double-branch attention bridge and a DCT frequency domain enhancement module, comprises:

[0009] Step 1) Obtain initial data set of different scenes underwater, and pre-process and divide the images in the data set;

[0010] Step 2) input the pre-processed image into the U-Net network, first pass through the position perception convolution module for multiple downsampling, and extract the feature information of the original underwater image;

[0011] Step 3) input the extracted feature information into the double-branch attention bridge module, and correct the feature information in fine granularity, enhance the perception effect of the model on important features such as texture and edge;

[0012] Step 4) further pass the upsampled feature map through the DCT frequency domain enhancement module, retain the details of the underwater image, reduce artifacts and blur, and improve the effect of underwater image enhancement;

[0013] Further, the step 1) specifically comprises the following steps:

[0014] Step 11) underwater data set division, divide the public underwater image data set into training set, validation set and test set according to the proportion of 8:1:1, and at the same time, provide no-reference image as challenging image to verify the generalization of the model.

[0015] Further, the step 2) specifically comprises the following steps:

[0016] Step 21) position perception information extraction

[0017] First, the relative position information of the input image in the horizontal direction and the vertical direction is extracted, and the standard convolution operation is combined to express the spatial relationship of the image. This operation helps to enhance the dependence between distant pixels and improve the global perception ability of the features.

[0018] Step 22) channel segmentation and convolution operation

[0019] The channel dimension of the image is split into multiple sub-blocks, and convolution operations are performed in the horizontal and vertical directions respectively. This process enables the convolutional network to extract spatial information with strong direction sensitivity, effectively capturing local structural features in the image.

[0020] Step 23) Horizontal direction path processing

[0021] In the horizontal direction path, direction modeling and deep convolution are performed after integrating position encoding. Through 1x1 convolution, the original features are fused to enhance the context information of multi-scale features and improve the spatial expression ability of the image.

[0022] Step 24) Vertical direction path processing

[0023] The features processed by the horizontal direction path are passed to the vertical path, and similar convolution operations are performed to extract structural features in the vertical direction, further enhancing the local detail performance of the image.

[0024] Step 25) Feature enhancement and activation

[0025] The features processed by the horizontal and vertical directions are standardized by batch normalization, and then a nonlinear activation function is used to further enhance the feature expression ability, optimizing the network's recognition and learning ability for different features.

[0026] Further, the step 3) specifically includes the following steps:

[0027] Step 31) Feature path construction

[0028] First, two feature paths are constructed in parallel through standard convolution and dilated convolution to achieve multi-scale coverage of the receptive field. This operation can simultaneously capture local details and global structural features of the image, improving the spatial perception ability of the network.

[0029] Step 32) Dimension reduction processing

[0030] After constructing the feature path, dimension reduction is performed through 1x1 convolution to reduce computational cost. This processing can reduce the complexity of the network while maintaining information expression, thereby improving the running efficiency of the model.

[0031] Step 33) Pooling operation and channel statistical feature extraction

[0032] Average pooling and max pooling operations are introduced to extract channel statistical features and obtain global information of the image. Through these two pooling operations, the model can capture image features at different scales, providing more rich feature information for subsequent processing.

[0033] Step 34) Introduction of spatial attention mechanism

[0034] In combination with the spatial attention strategy, the response capability of the model to key areas is improved. By calculating the spatial attention map, the model can dynamically adjust the attention degree of different areas, so as to better capture important areas in the image, such as texture, edge and other details.

[0035] Step 35) Channel attention path construction

[0036] On this basis, a channel attention path based on global pooling and layer-by-layer convolution is constructed. This path adjusts the feature channels through weighting, highlights the significant features, and further improves the model's ability in feature extraction.

[0037] Step 36) Output fusion and fine-grained correction

[0038] Finally, the weighted multi-path outputs are fused, and a spatial attention map is generated to make fine-grained corrections to the original feature map. This process enhances the model's perception of important targets, ultimately improving the quality of underwater image enhancement.

[0039] Further, the step 4) specifically comprises the following steps:

[0040] Step 41) Feature extraction and construction of query, key and value representation

[0041] First, features are extracted from RGB and NIR modalities through independent convolution paths. For the extracted features, query (Q), key (K) and value (V) representations are constructed to prepare for subsequent frequency domain processing and attention mechanism operations.

[0042] Step 42) Feature division and frequency domain mapping

[0043] The extracted features are divided into small blocks, and then two-dimensional discrete cosine transform (DCT) is applied to each small block to map it to the frequency domain. Through frequency domain conversion, the module can enhance the texture and structure information in the image, especially in detail capture and high-frequency feature enhancement.

[0044] Step 43) Similarity calculation between query and key

[0045] In the frequency domain, the module calculates the similarity between the query (Q) and the key (K). Based on the calculated similarity, the module performs weighted fusion on the value (V) feature through the attention mechanism, so as to realize efficient interaction and fusion of information.

[0046] Step 44) Inverse DCT transformation and space domain back

[0047] The frequency domain features after weighted fusion are converted back to the spatial domain by inverse DCT transformation to restore the spatial domain features of the image. In this process, a learnable parameter λ is introduced, and an initialization function is combined to adaptively balance the contribution of multi-modal information and ensure reasonable fusion of the two modalities in image restoration.

[0048] Step 45) Multi-path pooling and feature integration

[0049] Finally, the multi-path pooling operation is used to integrate the frequency domain and spatial domain features. Through linear mapping and pooling, the multi-modal features are further integrated, which significantly enhances the comprehensive modeling capability of the model for context information and local details, and improves the overall performance and detail perception capability of the image. BRIEF DESCRIPTION OF DRAWINGS

[0050] In order to make the purpose, technical scheme and advantages of the present application clearer, the present application will be further described in detail below in combination with the drawings, in which:

[0051] Figure 1 The flow chart of the method is shown in

[0052] Figure 2 The overall model framework of the method is shown in

[0053] Figure 3 The position-aware convolution structure proposed by the method is shown in

[0054] Figure 4 The double-branch attention bridging structure proposed by the method is shown in

[0055] Figure 5 The DCT fusion enhancement structure proposed by the method is shown in DETAILED DESCRIPTION

[0056] The preferred examples of the present application will be described in detail below in combination with the drawings.

[0057] The underwater image enhancement method based on the dual-domain attention U-Net provided by the present application is shown in the flow chart of the method Figure 1 The overall model framework of the method proposed by the present application is shown in Figure 2 The method includes the following steps:

[0058] Step 1) Divide the underwater image enhancement dataset;

[0059] Step 2) Build a position-aware convolution module;

[0060] Step 3) Build a double-branch attention bridging module;

[0061] Step 4) Build a DCT fusion enhancement module;

[0062] Step 5) training the underwater image enhancement network;

[0063] Step 6) testing the underwater image enhancement network;

[0064] Further, the specific content of step 5) includes the following steps:

[0065] Step 51) configure the corresponding environment, set the training parameters, and carry out the training task;

[0066] Step 52) the configured training environment is: Ubuntu 22.04, CUDA 11.8, Python 3.10.14 and Pytorch 2.3.0. The hardware devices include CPU: Intel(R) Xeon(R) CPU E5-2680v4@2.40GHz, memory 32GB; GPU: NVIDIA RTX 3090, memory 24GB.

[0067] Step 53) the configuration of the training parameters is: the input image size is 256x256, the batch size is set to 4, the learning rate is 0.0001, and the training round is 500.

[0068] Further, the specific content of step 6) includes the following steps:

[0069] Step 61) load the trained weight file into the underwater image enhancement network, select pictures from the test data set for enhancement, and return the underwater images with enhancement.

[0070] Finally, it should be pointed out that the above embodiments are only used to illustrate the technical solutions of the present application and are not limiting. Although the present application has been described in detail through the above embodiments, those skilled in the art should understand that various changes can be made in form and detail without departing from the scope defined by the claims of the present application.

Claims

1. An underwater image enhancement method based on a U-Net backbone network, integrating a position perception module, a dual-branch attention bridge, and a DCT frequency domain enhancement module. Characterized by: The following steps are involved: Step 1: Obtain an open-source dataset for underwater image enhancement from a public data platform. The collected images are divided into training, validation, and test sets in an 8:1:1 ratio, covering a variety of typical underwater scenes, including complex background environments, biological communities, and single-object images. For each degraded underwater image, a corresponding reference image pair is constructed to ensure complete supervision information for the training process. Step 2: Construct a position-aware convolution module: This module focuses on directional feature modeling, integrating relative position embedding with standard convolution operations to achieve joint modeling of global image relationships and local details. Specifically, the module first introduces relative position information in the horizontal and vertical directions to express spatial relationships and enhance the dependencies between distant pixels. The channel dimension is then divided into multiple sub-blocks, and convolution operations are performed in the horizontal and vertical directions to extract spatial information with directionality. In the horizontal path, after integrating the positional encoding, directional modeling and depthwise convolution are performed, and then fused with the initial features through 1×1 convolution to enhance the multi-scale feature context. The output is then passed to the vertical path, where similar operations are repeated to capture vertical structural features. Finally, batch normalization and nonlinear activation functions are used to enhance the overall expression capability. Step 3: Construct a dual-branch attention bridge module: This module integrates local details and global structure, and constructs two feature pathways in parallel through standard convolution and dilated convolution to achieve multi-scale coverage of the receptive field. 1×1 convolution is then used to reduce dimensionality to reduce computational costs. On this basis, the module introduces channel statistical features extracted by average pooling and maximum pooling, combined with a spatial attention strategy to improve the responsiveness of key areas. At the same time, a channel attention path based on global pooling and layer-by-layer convolution is constructed to achieve weighted adjustment of feature channels and highlight significant features. Finally, the weighted multi-path outputs are fused and a spatial attention map is generated, which performs fine-grained corrections on the original feature map to improve the model's perception of textures, edges, and important targets. Step 4: Construct a DCT fusion enhancement module: This module focuses on the effective synergy between the frequency and spatial domains, improving the model's ability to integrate multimodal information. Features are extracted from the RGB and NIR modalities through independent convolutional paths, and query (Q), key (K), and value (V) representations are constructed. The features are then divided into small blocks and mapped to the frequency domain using a two-dimensional discrete cosine transform (DCT) to enhance the capture of texture and structural information. In the frequency domain, the module calculates the similarity between the query and key and weights the fused value features based on an attention mechanism to achieve information interaction. The fusion result is transformed back to the spatial domain via an inverse DCT. A learnable parameter λ is introduced, and an initialization function is used to adaptively balance multimodal contributions. Finally, multi-path pooling and linear mapping are used to integrate the frequency and spatial features, significantly enhancing the ability to comprehensively model context and local details. Step 5, model training phase: using paired original and reference images as training samples, combining multiple metric functions such as pixel-level error, gradient difference, and structural similarity loss, a weighted combination method is used to supervise and optimize the network output. By continuously adjusting the weight ratio of each loss item, the network is guided to learn parameters with high-quality image restoration capabilities, thereby improving the overall enhancement effect. Step 6: Testing the underwater image enhancement network model: After training, the optimal model parameters are loaded and the image enhancement results are evaluated on a test set. Model performance is verified by comparing metrics such as structural similarity (SSIM) and peak signal-to-noise ratio (PSNR) between the before and after enhancement. Additionally, some unreferenced images are tested, combining image quality metrics such as color saturation and contrast to comprehensively assess the model's performance and generalization capabilities in complex underwater environments.

2. According to claim 1, an underwater image enhancement method based on a U-Net backbone network, integrating a position perception module, a dual-branch attention bridge, and a DCT frequency domain enhancement module, is characterized in that: The position-aware convolutional module constructed in step 2 processes underwater image features as follows: This module focuses on spatial relationship modeling, aiming to enhance the representation of structure and semantics in the image. First, a relative position embedding mechanism is introduced in the horizontal and vertical directions. By encoding the relative distances and semantic associations between different spatial locations, it improves the representation of global dependencies between remote pixels in the feature map. Next, the channel dimension is divided into several parts according to a certain ratio, and directional convolution operations are performed in the horizontal and vertical directions to obtain direction-sensitive features. In the horizontal path, after fusing the relative position encoding, a directional reconstruction operation and depthwise separable convolution are performed. Feature channels are compressed and fused using 1×1 convolution to enhance contextual awareness. This intermediate result is then passed to the vertical path, where the same processing is repeated to obtain structural features in the vertical dimension and achieve deep fusion of multi-directional features. Finally, the overall feature is further processed using batch normalization (BN) and activation functions (such as ReLU) to enhance the network's ability to represent and distinguish spatial details and image structure, thereby optimizing downstream image enhancement performance.

3. The underwater image enhancement method according to claim 1, which is based on a U-Net backbone network and integrates a position perception module, a dual-branch attention bridge, and a DCT frequency domain enhancement module, is characterized in that: The dual-branch attention bridge module constructed in step 3 processes underwater image features as follows: This module aims to accurately model and extract key region features by fusing different scales and attention mechanisms. Initially, the module constructs two parallel feature extraction paths: one is a standard convolution branch, which extracts local texture details in the image; the other uses a dilated convolution structure to expand the receptive field and capture a wider range of contextual information. Both branches perform channel compression through 1×1 convolution to reduce computational burden. Subsequently, based on the fusion of the two branch outputs, average pooling and max pooling operations are introduced to extract channel statistics at different scales. A spatial attention mechanism is combined to enhance activation responses at key locations. Furthermore, a channel attention path based on global pooling and layer-by-layer convolution is introduced to selectively retain or suppress features by learning the importance weights of different channels. Finally, a spatial attention map is generated by fusing the weighted features from the two branches. This is further used to guide fine-grained adjustments to the original feature map, significantly improving the model's responsiveness to salient image regions and texture details, and enhancing the discriminative and adaptable nature of feature representation.

4. The underwater image enhancement method according to claim 1, which is based on a U-Net backbone network and integrates a position perception module, a dual-branch attention bridge, and a DCT frequency domain enhancement module, is characterized in that: The DCT fusion enhancement module constructed in step 4 processes underwater image features as follows: This module improves the accuracy of structural texture modeling by introducing a joint processing mechanism for information in the frequency and spatial domains, making it particularly suitable for multimodal image fusion tasks. Initially, multiple independent convolutional paths are used to extract modal features from the RGB and NIR (near-infrared) images, respectively, and these features are constructed into three representations: query (Q), key (K), and value (V). Subsequently, the feature map is divided into several local regions, and each region is mapped to the frequency domain using a two-dimensional discrete cosine transform (2D DCT). This frequency domain transformation helps reveal texture intensity and repetitive structure within the local region. In the frequency domain, the module calculates the similarity between the query and key using an attention mechanism and guides a weighted fusion operation on the value features, thereby achieving effective interaction between the modalities. After fusion, the frequency domain features are restored back to the spatial domain using an inverse DCT, preserving their enhanced texture and structural properties. Furthermore, the module introduces a trainable λ parameter to automatically adjust the fusion weights between the different modalities based on network learning. To further improve the fusion effect, the module also uses multi-path pooling operations and linear projection mechanisms to integrate multi-scale contextual information, and ultimately outputs feature maps with greater discriminative power and structural consistency, thereby optimizing the quality of subsequent image enhancement.

Images