A lightweight real-time underwater image enhancement method based on axial fusion and dual-domain attention

Abstract

The application discloses a lightweight real-time underwater image enhancement method based on axial fusion and dual-domain attention, and belongs to the technical field of image data processing. The method is based on an encoder-decoder architecture, and through lightweight adaptive frequency-space attention and hybrid attention, the synergistic enhancement of frequency domain global features and spatial domain detail features is realized, and the adaptive weighted fusion of local features and axial features is completed by using an axial fusion block, so that underwater light scattering can be effectively inhibited, color deviation can be corrected, and image detail restoration precision can be improved. The application adopts grouping convolution and modular design to significantly reduce model parameter quantity and calculation overhead, and combines two-stage skip connection to guarantee cross-level feature transmission integrity, so that the lightweight and real-time performance can be realized while the enhancement effect is ensured, and the application can be well adapted to power-limited platforms such as underwater robots, and provides efficient and reliable image enhancement support for underwater visual perception tasks.

Description

Technical Field

[0001] This invention belongs to the field of image data processing technology, specifically relating to a lightweight real-time underwater image enhancement method based on axial fusion and dual-domain attention. Background Technology

[0002] Underwater image enhancement is a core technology at the intersection of computer vision and marine engineering. It forms the basis for tasks such as underwater target detection, navigation, and resource exploration, and its performance directly determines the perception and decision-making capabilities of underwater vision systems. However, the underwater imaging environment is complex and harsh: the selective absorption of different wavelengths of light by water causes severe color casts in images; light scattering caused by suspended particles and plankton leads to image blurring and significantly reduced contrast; and the lack of light in deep water areas and the susceptibility of auxiliary lighting to problems such as light spots and overexposure. These multiple degradations severely affect the extraction of effective information from underwater images, making the development of efficient and robust underwater image enhancement technologies crucial to overcoming the bottlenecks in underwater vision applications.

[0003] Existing underwater image enhancement methods are mainly divided into two categories: non-deep learning and deep learning. Non-deep learning methods include techniques based on physical models and those based on prior knowledge. The former relies on the underwater light propagation laws to construct a degradation model and recover the image by solving for parameters, while the latter adjusts visual attributes based on image statistical characteristics or traditional image processing methods. However, these methods have limited adaptability to complex underwater environments and suffer from drawbacks such as difficulty in parameter estimation, weak generalization, and susceptibility to interference from environmental factors such as water quality and depth. Deep learning methods, with their powerful nonlinear mapping capabilities, can autonomously learn the complex mapping relationship between degraded and clear images, performing better in complex scene processing. Models such as convolutional neural networks, generative adversarial networks, and Transformers have been widely applied in this field. However, current mainstream deep learning models generally adopt complex network architectures in pursuit of enhanced accuracy, resulting in problems such as a large number of parameters, high computational overhead, and low inference efficiency. This makes it difficult to meet the real-time processing and deployment requirements of computing-constrained platforms such as underwater robots and portable detection devices, and it is impossible to achieve a coordinated balance between lightweight design, real-time performance, and enhanced performance.

[0004] Therefore, developing underwater image enhancement technology that combines lightweight structure, high real-time performance, and superior enhancement effects, while ensuring color correction, detail restoration, and descattering performance, and adapting to the deployment requirements of underwater resource-constrained hardware, has significant theoretical and engineering practical value. This technology can provide high-quality real-time visual feedback for underwater mobile devices, improving the efficiency and safety of underwater exploration, operations, and monitoring tasks. It breaks through the hardware deployment bottleneck of underwater vision technology, provides technical support for marine resource development and underwater environmental monitoring, and promotes the implementation and large-scale application of underwater intelligent vision systems. Summary of the Invention

[0005] To address the challenges of existing underwater image enhancement methods that struggle to simultaneously achieve lightweight, real-time, and high-quality enhancement, and which fail to effectively suppress underwater scattering, correct color cast, and adapt to computationally limited platforms, this invention provides a lightweight real-time underwater image enhancement method (AFB-UNet) based on axial fusion and dual-domain attention. This method significantly reduces model computation and parameter count, achieving a high-efficiency balance between enhancement performance, model lightweightness, and real-time processing. The method designs an axial fusion module (AFB) as the main building block of AFB-UNet to ensure computational efficiency. To address the issue of computational accuracy being neglected in pursuit of speed in depthwise separable convolution, an attention mechanism is introduced after the depthwise separable convolution to separately select and weightedly fuse key features in the channel and spatial domains, enhancing the feature enhancement capabilities of AFB-UNet. Furthermore, to address the issue of global detail loss during attention selection, multi-scale pooling and global histogram stretching feature correction strategies are employed to compensate for lost global details and achieve dynamic feature optimization.

[0006] To achieve the above objectives, the present invention adopts the following technical solution:

[0007] A lightweight real-time underwater image enhancement method based on axial fusion and dual-domain attention includes the following steps:

[0008] S1. Perform visual feature preprocessing on the original underwater image. Initialize the features by depthwise separable convolution and batch normalization to obtain the preprocessed feature map.

[0009] S2. Perform two-level multi-scale coding feature extraction on the preprocessed feature map. The first-level encoder uses frequency-spatial attention (LA_FSA) and spatial attention to extract frequency domain global features and spatial domain detail features in sequence. The weighted fusion of local branch features and axial branch features is completed by the axial fusion block (AFB). The first-level encoder encoded features and corresponding skip connection features are downsampled and output. The second-level encoder uses hybrid attention to dynamically fuse channel and spatial information, adjust the axial fusion convolution size to expand the receptive field, and downsample and output the second-level encoder encoded features and corresponding skip connection features.

[0010] S3. Input the encoded features of the second-level encoder into the bottleneck layer, and perform deep feature enhancement through channel attention and large kernel axial fusion block to obtain the enhanced features of the bottleneck layer.

[0011] S4. Perform two-level multi-scale decoding feature restoration on the bottleneck layer enhancement features. Combine the skip connection features of the decoder and complete feature upsampling and detail fusion through the corresponding attention mechanism and axial fusion block to obtain the decoded output features.

[0012] S5. Map the decoded output features to the number of channels of the original underwater image, and output the final enhanced underwater image through an activation function.

[0013] Furthermore, the visual feature preprocessing in step S1 uses depthwise separable convolutional layers to reduce computational cost while initially extracting features, laying the foundation for subsequent processing. Step S1 specifically includes:

[0014] Raw underwater images Pixel-level feature decomposition and channel dimension mapping are performed sequentially using 3×3 convolution and 1×1 convolution, followed by batch normalization to stabilize the feature distribution, resulting in a preprocessed feature map. This lays a solid foundation for subsequent encoding and decoding feature extraction. The specific formula is as follows:

[0015] Formula 1: ,

[0016] in, This represents the preprocessed feature map. This represents the batch normalization function. Represents a 1×1 convolution. This represents a 3×3 convolution, where X represents the original underwater image. This indicates grouping by channel dimension.

[0017] The underwater image enhancement target can be defined as: the original underwater image (R represents an RGB image, and 3 represents the number of channels) ), through network mapping function To generate the final enhanced image ,Should It can approximate the real, clear underwater images in the calibration standard file of the dataset in terms of visual effects (detail, color, contrast) and quantitative indicators. The enhancement process is as follows:

[0018] Formula 2: ,

[0019] in, This represents the final enhanced image. This represents a true and clear underwater image. This represents the learnable parameters of the network. Represents network mapping function ( ) represents the loss function. This indicates that a minimization optimization is performed on all learnable parameters θ in the network. Represents the original input features. .

[0020] Furthermore, in step S2, multi-scale encoding feature extraction is performed, and a hierarchical extraction strategy is applied to the preprocessed feature map: basic-scale encoding by the first-level encoder and high-order-scale encoding by the second-level encoder. This addresses the differentiated needs of feature representation at different scales in stages, avoiding the global and local feature imbalance caused by single-scale encoding. Specifically:

[0021] The multi-scale coding feature extraction module can adapt to the expression rules of image features from "basic details" to "high-level abstraction". First, the first-level encoder captures frequency domain global features and spatial domain detail features at the basic scale. Then, the second-level encoder extracts hybrid attention features and wide receptive field axial features at the high-level scale. Each stage focuses on the core feature extraction target of the corresponding scale, achieving progressive feature encoding from "fine-grained details" to "high-dimensional abstraction". The specific process is as follows: First-level coding feature extraction: The task is to extract frequency domain global features and spatial domain detail features at the basic scale, and fuse axial dimension basic features. Therefore, feature dimension transformation is first completed through grouped convolution and 1×1 convolution to ensure feature dimension adaptation while reducing computational load. Then, the frequency attention module (LA_FSA) is connected to extract frequency domain features. LA-FSA can focus on global frequency domain rules and avoid interference from spatial details on frequency domain features. The spatial attention module accurately captures spatial details at the basic scale. The axial fusion block (AFB) fuses height and width dimension features through axial convolution. The specific process is as follows:

[0022] (1) First-stage encoder

[0023] The core objective of this layer is to capture the global regularity in the frequency domain and the fine-grained details in the spatial domain at the basic scale. Through the link of "dimensional transformation → frequency domain-spatial decoupling → axial fusion → residual downsampling", it achieves accurate encoding of features at the basic scale. The specific process and formulas are as follows:

[0024] Formula 3: ,

[0025] Formula 4: ,

[0026] in, This represents the initial features of the first encoder. Represents a 1×1 convolution. This represents a 3×3 convolution. This represents the preprocessed feature map. This indicates grouping by channel dimension. This represents the first-level frequency domain enhancement feature. This represents the Sigmoid activation function. This represents the global average pooling function. This represents a two-dimensional real-valued Fast Fourier Transform, where |·| represents the amplitude in the frequency domain. Represents the ReLU activation function. These are learnable parameters, and their range is (0,1).

[0027] The preprocessed feature map is used as the input feature map. The feature dimension is transformed through group convolution and 1×1 convolution to obtain the initial features of the first encoder. Subsequently, LA_FSA is used to extract and enhance global frequency domain information (such as overall brightness distribution, color shift trends, and low-frequency contour features) from the basic scale features, avoiding interference from spatial details (such as edges and textures) on the global frequency domain regularity. This lays the foundation for the subsequent spatial attention module to focus on local details, and improves the initial features. The LA_FSA module processes the data as follows: First, it maps the spatial features to the frequency domain using a two-dimensional real-valued Fast Fourier Transform (rfft2). After extracting the frequency domain amplitude features, it performs global adaptive average pooling to compress and obtain the global frequency domain energy statistics for each channel. Then, it generates frequency domain attention weights through 1×1 convolution and sigmoid activation, and then... Multiplication enhances important channel features in the frequency domain (such as channels with significant color shifts) and weakens worthless noise channel features, outputting the first-level frequency domain enhancement feature. This feature preserves the global optical regularity of the image and filters out the interference of local details on the global features.

[0028] Formula 5: ,

[0029] in, This indicates the first-level airspace detail enhancement feature. This represents the first-level frequency domain enhancement feature. This represents the Sigmoid activation function. This represents a 3×3 convolution. Indicates channel splicing. This represents the global average pooling function. This represents the max pooling function.

[0030] Calculate the first-level frequency domain enhancement features The global mean and global maximum features are concatenated along the channel dimension. Spatial correlation information of the concatenated features is extracted through 3×3 convolution, and then spatial attention weights are generated by Sigmoid activation. Multiplication enhances detailed regions such as edges and textures in the feature map while suppressing regions lacking detail, such as the background, and outputs the first-level spatial detail enhancement feature. This completes the decoupling and extraction of "frequency domain-spatial domain" features, preserving global frequency domain regularities while enhancing local spatial domain details.

[0031] Formula 6: ,

[0032] Formula 7: ,

[0033] Formula 8: ,

[0034] in, This represents the first-level local branch feature. This indicates the first-level airspace detail enhancement feature. This indicates the first-order axial branch feature. This represents a 1×3 convolution. This represents a 3×1 convolution. This indicates the first-level axial fusion feature. , This represents learnable weighted parameters.

[0035] Enhanced features for first-level airspace details The system sequentially performs 3×1 height-axis convolutions and 1×3 width-axis convolutions (orthogonal axis convolutions, both grouped convolutions) to separate and extract axial feature associations in the height and width dimensions (such as horizontal / vertical contour features). This compensates for the insufficient axial feature capture caused by the isotropic receptive field of traditional convolutions, and outputs the first-level axial branch features. .right Perform 3×3 grouped convolution to extract local spatial features, and then integrate channel information through 1×1 convolution to output the first-level local branch features. This branch ensures the diversity of feature extraction and avoids the loss of feature information caused by single-axis convolution. Learning parameters are introduced. Balancing the contribution of axial features and local features from convolutional features, and the contribution of first-level spatial detail enhancement features. Residual connections are performed, and axially fused enhancement features are superimposed while preserving the original feature structure to avoid information loss during feature fusion, finally yielding the first-level axially fused features. .

[0036] Formula 9: ,

[0037] in, This represents the encoding features of the first-level encoder. Represents the max pooling function. Represents a 1×1 convolution. This represents the batch normalization function. This indicates the first-level axial fusion feature.

[0038] Finally, the first-level axial fusion features After batch normalization and 1×1 convolution to adjust the channels, max pooling is performed. The core function is to halve the feature space resolution, realize downsampling and abstraction of basic scale features, and output the features encoded by the first-level encoder. It provides low-resolution input to the second-level encoder, while the features before pooling are retained as skip connection features for the decoder to use.

[0039] (2) Second-stage encoder

[0040] The second-stage encoder largely follows the core chain of the first stage: "convolutional transformation → attention extraction → axial fusion → normalization → residual downsampling". The core differences are in two aspects: first, the frequency and spatial attention of the first stage are replaced with hybrid attention. Dynamically fuse channel attention results through gating mechanism With spatial attention results First, it adapts to the need for capturing global channel associations in high-order scale features, compensating for the lack of dimensional information in single attention. Second, it adjusts the axial convolution kernel from 3×3 to 5×5 to expand the axial receptive field, adapting to the characteristics of reduced resolution and wider pixel association range in high-order scale features, and more fully capturing large-scale axial dimensional features. Finally, the second-level axial fusion feature obtained in this layer is used to... Perform downsampling and output the encoded features of the second-level encoder. Meanwhile, the features before pooling will be retained as skip connection features for the decoder to use.

[0041] Formula 10: ( ),

[0042] Formula 11: ,

[0043] Formula 12: ,

[0044] Formula 13: ,

[0045] Formula 14: ,

[0046] Formula 15: ,

[0047] Formula 16: ,

[0048] Formula 17: ,

[0049] in, This represents the initial features of the second-level encoder. This represents the encoding features of the first-level encoder. This represents the channel attention enhancement feature of the second-level encoder. This represents the spatial attention enhancement features of the second-level encoder. This represents the hybrid attention output features of the second-level encoder. This represents the second-level local branch feature. This indicates a second-order axial branch feature. This indicates the second-level axial fusion feature. This represents the encoding feature of the second-level encoder, where `gate` represents gating, and its value range is... , and These are learnable parameters.

[0050] Furthermore, in step S3, the bottleneck layer feature enhancement encodes features from the second-level encoder. By inputting the bottleneck layer and maintaining the spatial resolution throughout, the bottleneck channel attention enhancement feature is first obtained by strengthening the optical feature correlation of the channel dimension through the channel attention module. Then, with the help of the 7x7 large core AFB, Extracting more spatial detail features yields Then and The features are concatenated and fused along the channel dimension. A 1×1 convolution is used to restore the number of feature channels after concatenation to complete the dimension matching. Finally, the adjusted fused features are combined with the features encoded by the second-level encoder. Residual connections enable deep enhancement of both global color and local details, outputting the final enhanced features of the bottleneck layer. This provides high-purity core abstract features for subsequent feature upsampling and restoration in the decoder.

[0051] Formula 18: ,

[0052] Formula 19: ,

[0053] Formula 20: ,

[0054] Formula 21: ,

[0055] in, This indicates the local branching characteristics of the bottleneck. This represents the spatial attention features of the second-level encoder. Indicates the axial branching characteristic of the bottleneck. This represents the encoding features of the second-level encoder.

[0056] Indicates the axial fusion characteristics of the bottleneck. This indicates the final enhancement feature of the bottleneck layer. This indicates a bottleneck channel attention enhancement feature.

[0057] Furthermore, step S4, the feature restoration of the decoding layer, uses the bottleneck layer enhanced features as the core input, combined with the encoder's two-level skip connection features, and executes a hierarchical restoration strategy of "first-level medium-scale restoration → second-level original-scale restoration". This stage-wise completes the feature details at different scales, achieving accurate restoration from high-order abstract features to original image scale features, avoiding detail loss and feature distortion caused by single-scale restoration. Specifically:

[0058] The decoding layer feature restoration module can adapt to the restoration pattern of image features from "high-level abstraction" to "basic details". First, the bottleneck layer features are restored to mesoscale features through the first-level decoder, and the skip connection features of the second-level encoder are fused to complete the mesoscale details. Then, the second-level decoder restores the features to the original scale, and the skip connection features of the first-level encoder are fused to complete the fine-grained details. Each stage focuses on the core restoration target of the corresponding scale, realizing a progressive restoration from "high-dimensional abstraction" to "original scale features". The specific process is as follows:

[0059] (1) First-level decoder

[0060] The core objective of this layer is to ultimately enhance the features of the high-level abstraction obtained from the bottleneck layer. Transposed convolution Perform upsampling and restoration, and output the upsampled features of the first-stage decoder. This feature is a mesoscale feature, fused with skip connection features from the second-level encoder. Complete the mesoscale details to obtain the first-level decoder fused skip connection features By following the chain of "upsampling scale restoration → feature fusion → detail enhancement → dimensionality integration", a smooth transition from high-scale features to mesoscale features is achieved. The formula is as follows:

[0061] Formula 22: ,

[0062] Formula 23: ,

[0063] in, This represents the upsampled features of the first-stage decoder. This represents a 2×2 transposed convolution. This indicates the final enhancement feature of the bottleneck layer. This indicates that the first-level decoder fuses skip connection features. This represents the skip connection feature of the second-level encoder. This indicates channel splicing.

[0064] By integrating the hybrid attention module and adopting the gating fusion logic of the hybrid attention module in the second-level encoder, the parameter scale is simplified to adapt to the decoding layer's detail enhancement requirements. It dynamically balances global correlations between channels and spatial scale details, strengthens effective information in the stitched features, suppresses redundant interference, and outputs the hybrid attention-enhanced features of the first-level decoder. Then, through the AFB block, which employs 5×5 axial convolution, the first-level decoder axial fusion features under the mesoscale feature map are obtained. The specific formula is as follows:

[0065] Formula 24: ( ),

[0066] Formula 25: ,

[0067] Formula 26: ,

[0068] Formula 27: ,

[0069] Formula 28: ,

[0070] Formula 29: ,

[0071] Formula 30: ,

[0072] in, This represents the initial features of the first-level decoder. This indicates that the first-level decoder fuses skip connection features. This represents the channel attention enhancement features of the first-level decoder. This represents the spatial attention enhancement features of the first-level decoder. This represents the hybrid attention output features of the first-level decoder. This represents the local branch features of the first-level decoder. This indicates the axial branch feature of the first-level decoder. This indicates the axial fusion feature of the first-level decoder.

[0073] (2) Second-level decoder

[0074] The second-level decoder generally follows the core chain of the first-level decoder: "upsampling → skip fusion → attention extraction → axial fusion → normalization". The core differences are reflected in three aspects: First, the skip connection features with the first-level encoder. The first method involves fusion; the second method replaces the first-level hybrid attention with frequency and spatial attention (LA_FSA) to obtain the results. Third, the axial convolution kernel is adjusted from 5×5 to 3×3 to reduce the axial receptive field, adapting to the low-scale feature resolution and capturing more detailed features. Finally, the second-level decoder axial fusion feature output from this layer is obtained. The specific formula is as follows:

[0075] Formula 31: ,

[0076] Formula 32: ,

[0077] Formula 33: ,

[0078] Formula 34: ,

[0079] Formula 35: ,

[0080] Formula 36: ,

[0081] in, This represents the upsampled features of the second-stage decoder. This represents the skip connection feature of the first-level encoder. This indicates that the second-level decoder fuses skip connection features. This indicates the second-level frequency domain enhancement feature. This represents the local branch features of the second-level decoder. This indicates the axial branch feature of the second-level decoder. This indicates the axial fusion feature of the second-level decoder.

[0082] Furthermore, in step S5, the output layer receives the axial fusion features from the second-level decoder. The core objective is to perform final purification and dimensionality adaptation on the original scale features. First, the image is processed using depthwise separable convolutions consisting of 3×3 grouped convolutions combined with 1×1 convolutions. Then, it is mapped using a Sigmoid activation function to output the final enhanced underwater image. The specific process and formula are as follows:

[0083] Formula 37: ,

[0084] in, This indicates the final enhanced underwater image. This indicates the axial fusion feature of the second-level decoder.

[0085] The beneficial effects of this invention are:

[0086] (1) The present invention adopts a dual-domain attention mechanism consisting of lightweight adaptive frequency-space attention and hybrid attention, which can synergistically enhance underwater image features in the spatial domain and frequency domain, effectively suppress underwater light scattering and color shift problems, and improve image detail and color reproduction accuracy.

[0087] (2) The present invention uses the dual-branch parallel structure of the axial fusion block to adaptively weight and fuse local features and axial features, thereby reducing the amount of convolutional computation while ensuring the integrity of feature expression, so that the model has both strong feature extraction capability and lightweight characteristics.

[0088] (3) Based on the encoder-decoder structure and combined with a two-level skip connection mechanism, this invention realizes the cross-layer transmission and fusion of the encoder's low-level detailed features and the decoder's high-level semantic features, thus avoiding the loss of details and structural distortion during image enhancement.

[0089] (4) The present invention adopts a modular and progressive encoding and decoding design, which can realize real-time enhancement processing of underwater images while maintaining the compact structure of the model and low computational complexity, and can better adapt to the deployment needs of underwater mobile platforms with limited computing power. Attached Figure Description

[0090] Figure 1 This is a detailed framework diagram of the model of the present invention;

[0091] Figure 2 This is a schematic diagram of the data flow processing of the axial fusion module in step S2 of the present invention.

[0092] Figure 3 This is a schematic diagram illustrating the working principle of frequency-spatial domain collaborative attention in steps S2 and S4 of the present invention.

[0093] Figure 4 This is a schematic diagram illustrating the working principle of hybrid attention in steps S2 and S4 of the present invention. Detailed Implementation

[0094] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments, but the present invention is not limited thereto.

[0095] The detailed framework diagram of the model of this invention is as follows: Figure 1 As shown. The lightweight real-time underwater image enhancement method based on axial fusion and dual-domain attention provided by this invention includes the following steps:

[0096] In step S1, visual feature preprocessing uses depthwise separable convolutional layers to reduce computational cost while initially extracting features, laying the foundation for subsequent processing. Specifically:

[0097] Raw underwater images Visual feature preprocessing is performed: first, pixel-level feature decomposition and channel dimension mapping are completed through depthwise separable convolution; then, batch normalization is used to stabilize the feature distribution to obtain the preprocessed feature map. This lays a solid foundation for subsequent encoding and decoding feature extraction. The specific formula is as follows:

[0098] Formula 1: ,

[0099] in, This represents the preprocessed feature map. This represents the batch normalization function. Represents a 1×1 convolution. This represents a 3×3 convolution, where X represents the original underwater image. This indicates grouping by channel dimension.

[0100] The underwater image enhancement target can be defined as: the original underwater image (3-channel RGB image, image size) ), through network mapping function To generate the final enhanced image ,Should It can approximate the real, clear underwater images in the calibration standard file of the dataset in terms of visual effects (detail, color, contrast) and quantitative indicators. The enhancement process is as follows:

[0101] Formula 2: ,

[0102] in, This represents the final enhanced image. This represents a true and clear underwater image. This represents the learnable parameters of the network. Represents network mapping function ( ) represents the loss function. This indicates that a minimization optimization is performed on all learnable parameters θ in the network. Represents the original input features. .

[0103] In step S2, multi-scale encoding feature extraction employs a hierarchical extraction strategy of "first-level basic scale encoding → second-level higher-order scale encoding" on the preprocessed features. This addresses the differentiated needs of feature representation at different scales in stages, avoiding global-local feature imbalance caused by single-scale encoding. Specifically:

[0104] The multi-scale coding feature extraction module can adapt to the representation rules of image features from "basic details" to "high-level abstraction". First, the first-level encoder captures frequency domain global features and spatial domain detail features at the basic scale. Then, the second-level encoder extracts hybrid attention features and wide receptive field axial features at the high-level scale. Each stage focuses on the core feature extraction target of the corresponding scale, achieving progressive feature encoding from "fine-grained details" to "high-dimensional abstraction". The specific process is as follows: First-level coding feature extraction: The task is to extract frequency domain global features and spatial domain detail features at the basic scale, and fuse axial dimension basic features. Therefore, feature dimension transformation is first completed through grouped convolution and 1×1 convolution to ensure feature dimension adaptation while reducing computational load. Then, the frequency attention module (LA-FSA) is connected to extract frequency domain features. LA-FSA can focus on global frequency domain rules and avoid interference from spatial details on frequency domain features. The spatial attention module accurately captures spatial details at the basic scale. The axial fusion block (AFB) fuses height and width dimension features through axial convolution. The specific process is as follows:

[0105] (1) First-stage encoder

[0106] The core objective of this layer is to capture the global regularity in the frequency domain and the fine-grained details in the spatial domain at the basic scale. Through the link of "dimensional transformation → frequency domain-spatial decoupling → axial fusion → residual downsampling", it achieves accurate encoding of features at the basic scale. The specific process and formulas are as follows:

[0107] Formula 3: ,

[0108] Formula 4:

[0109] in, This represents the initial features of the first encoder. Represents a 1×1 convolution. This represents a 3×3 convolution. This represents the preprocessed feature map. This indicates grouping by channel dimension. This represents the first-level frequency domain enhancement feature. This represents the Sigmoid activation function. This represents the global average pooling function. This represents a two-dimensional real-valued Fast Fourier Transform, where |·| represents the amplitude in the frequency domain. Represents the ReLU activation function. is a learnable parameter, and its range is (0,1).

[0110] The preprocessed feature map is used as the input feature map. The feature dimension is transformed through grouped convolution and 1×1 convolution to obtain the initial feature map of the first encoder. ; then via LA_FSA (such as Figure 3 As shown, global frequency domain information (such as overall brightness distribution, color shift trends, and low-frequency contour features) is extracted and enhanced from basic scale features. This avoids interference from spatial details (such as edges and textures) on global frequency domain patterns, laying the foundation for the subsequent spatial attention module to focus on local details and improving the initial features. The LA_FSA module processes the data as follows: First, it maps the spatial features to the frequency domain using a two-dimensional real-valued Fast Fourier Transform (rfft2). After extracting the frequency domain amplitude features, it performs global adaptive average pooling to compress and obtain the global frequency domain energy statistics for each channel. Then, it generates frequency domain attention weights through 1×1 convolution and sigmoid activation, and then... Multiplication enhances important channel features in the frequency domain (such as channels with significant color shifts) and weakens worthless noise channel features, outputting the first-level frequency domain enhancement feature. This feature preserves the global optical regularity of the image and filters out the interference of local details on the global features.

[0111] Formula 5: ,

[0112] in, This indicates the first-level airspace detail enhancement feature. This represents the first-level frequency domain enhancement feature. This represents the Sigmoid activation function. This represents a 3×3 convolution. Indicates channel splicing. This represents the global average pooling function. This represents the max pooling function.

[0113] Calculate the first-level frequency domain enhancement features The global mean and global maximum features are concatenated along the channel dimension. Spatial correlation information of the concatenated features is extracted through 3×3 convolution, and then spatial attention weights are generated by Sigmoid activation. Multiplication enhances detailed regions such as edges and textures in the feature map while suppressing regions lacking detail, such as the background, and outputs the first-level spatial detail enhancement feature. This completes the decoupling and extraction of "frequency domain-spatial domain" features, preserving global frequency domain regularity while enhancing local spatial domain details.

[0114] Formula 6: ,

[0115] Formula 7: ,

[0116] Formula 8: ,

[0117] in, This represents the first-level local branch feature. This indicates the first-level airspace detail enhancement feature. This indicates the first-order axial branching feature. This represents a 1×3 convolution. This represents a 3×1 convolution. This indicates the first-level axial fusion feature. , This represents learnable weighted parameters.

[0118] Enhanced features for first-level airspace details Execute the AFB module (e.g.) Figure 2 As shown, it sequentially performs 3×1 height-axis convolution and 1×3 width-axis convolution (orthogonal axial convolution, both grouped convolutions), separating and extracting axial feature associations in the height and width dimensions (such as horizontal / vertical contour features), compensating for the insufficient axial feature capture caused by the isotropic receptive field of traditional convolution, and outputting the first-level axial branch features. .right Perform 3×3 grouped convolution to extract local spatial features, then integrate channel information through 1×1 convolution to output the first-level local branch features. This branch ensures the diversity of feature extraction and avoids the loss of feature information caused by single-axis convolution. Learning parameters are introduced. Balancing the contribution of axial features and local features from convolutional features, and the first-level spatial detail enhancement features. Residual connections are performed, and axially fused enhancement features are superimposed while preserving the original feature structure to avoid information loss during feature fusion, finally yielding the first-level axially fused features. .

[0119] Formula 9: ,

[0120] in, This represents the encoding features of the first-level encoder. Represents the max pooling function. Represents a 1×1 convolution. This represents the batch normalization function. This indicates the first-level axial fusion feature.

[0121] Finally, the first-level axial fusion features After batch normalization and 1×1 convolution to adjust the channels, max pooling is performed. The core function is to halve the feature space resolution, realize downsampling and abstraction of basic scale features, and output the features encoded by the first-level encoder. It provides low-resolution input to the second-level encoder, while the features before pooling are retained as skip connection features for the decoder to use.

[0122] (2) Second-stage encoder

[0123] The second-stage encoder largely follows the core chain of the first stage: "convolutional transformation → attention extraction → axial fusion → normalization → residual downsampling". The core differences are in two aspects: first, the frequency + spatial attention of the first stage is replaced with hybrid attention. (like Figure 4 As shown), the channel attention results are dynamically fused through a gating mechanism. With spatial attention results First, it adapts to the need for capturing global channel associations in high-order scale features, compensating for the lack of dimensional information in single attention. Second, it adjusts the axial convolution kernel from 3×3 to 5×5 to expand the axial receptive field, adapting to the characteristics of reduced resolution and wider pixel association range in high-order scale features, and more fully capturing large-scale axial dimensional features. Finally, the second-level axial fusion feature obtained in this layer is used to... Perform downsampling and output the encoded features of the second-level encoder. Meanwhile, the features before pooling will be retained as skip connection features for the decoder to use.

[0124] Formula 10: ( ),

[0125] Formula 11: ,

[0126] Formula 12: ,

[0127] Formula 13: ,

[0128] Formula 14: ,

[0129] Formula 15: ,

[0130] Formula 16: ,

[0131] Formula 17: ,

[0132] in, This represents the initial features of the second-level encoder. This represents the encoding features of the first-level encoder. This represents the channel attention enhancement feature of the second-level encoder. This represents the spatial attention enhancement features of the second-level encoder. This represents the hybrid attention output features of the second-level encoder. This represents the second-level local branch feature. This indicates a second-order axial branch feature. This indicates the second-level axial fusion feature. This represents the encoding feature of the second-level encoder, where `gate` represents gating, and its value range is... , and These are learnable parameters.

[0133] In step S3, the bottleneck layer feature enhancement encodes features from the second-level encoder. By inputting the bottleneck layer and maintaining the spatial resolution throughout, the bottleneck channel attention enhancement feature is first obtained by strengthening the optical feature correlation of the channel dimension through the channel attention module. Then, with the help of the 7x7 large core AFB, Extracting more spatial detail features yields Then and The features are concatenated and fused along the channel dimension. A 1×1 convolution is used to restore the number of feature channels after concatenation to complete the dimension matching. Finally, the adjusted fused features are combined with the features encoded by the second-level encoder. Residual connections enable deep enhancement of both global color and local details, outputting the final enhanced features of the bottleneck layer. This provides high-purity core abstract features for subsequent feature upsampling and restoration in the decoder.

[0134] Formula 18: ,

[0135] Formula 19: ,

[0136] Formula 20: ,

[0137] Formula 21: ,

[0138] in, This indicates the local branching characteristics of the bottleneck. This represents the spatial attention features of the second-level encoder. Indicates the axial branching characteristic of the bottleneck. This represents the encoding features of the second-level encoder.

[0139] Indicates the axial fusion characteristics of the bottleneck. This indicates the final enhancement feature of the bottleneck layer. This indicates a bottleneck channel attention enhancement feature.

[0140] Step S4, feature restoration of the decoding layer, uses the bottleneck layer enhanced features as the core input, combined with the encoder's two-level skip connection features, and executes a layered restoration strategy of "first-level medium-scale restoration → second-level original-scale restoration". This staged restoration of feature details at different scales achieves accurate restoration from high-order abstract features to original image-scale features, avoiding detail loss and feature distortion caused by single-scale restoration. Specifically:

[0141] The decoding layer feature restoration module can adapt to the restoration pattern of image features from "high-level abstraction" to "basic details". First, the bottleneck layer features are restored to mesoscale features through the first-level decoder, and the skip connection features of the second-level encoder are fused to complete the mesoscale details. Then, the second-level decoder restores the features to the original scale, and the skip connection features of the first-level encoder are fused to complete the fine-grained details. Each stage focuses on the core restoration target of the corresponding scale, realizing a progressive restoration from "high-dimensional abstraction" to "original scale features". The specific process is as follows:

[0142] (1) First-level decoder

[0143] The core objective of this layer is to ultimately enhance the features of the high-level abstraction obtained from the bottleneck layer. Transposed convolution Perform upsampling and restoration, and output the upsampled features of the first-stage decoder. This feature is a mesoscale feature, fused with skip connection features from the second-level encoder. Feature completion of mesoscale details yields first-level decoder fusion skip connection features By following the chain of "upsampling scale restoration → feature fusion → detail enhancement → dimensionality integration", a smooth transition from high-scale features to mesoscale features is achieved. The formula is as follows:

[0144] Formula 22: ,

[0145] Formula 23: ,

[0146] in, This represents the upsampled features of the first-stage decoder. This represents a 2×2 transposed convolution. This indicates the final enhancement feature of the bottleneck layer. This indicates that the first-level decoder fuses skip connection features. This represents the skip connection feature of the second-level encoder. This indicates channel splicing.

[0147] By integrating the hybrid attention module and adopting the gating fusion logic of the hybrid attention module in the second-level encoder, the parameter scale is simplified to adapt to the decoding layer's detail enhancement requirements. It dynamically balances global correlations between channels and spatial scale details, strengthens effective information in the stitched features, suppresses redundant interference, and outputs the hybrid attention-enhanced features of the first-level decoder. Then, through the AFB block, which employs 5×5 axial convolution, the first-level decoder axial fusion features under the mesoscale feature map are obtained. The specific formula is as follows:

[0148] Formula 24: ( ),

[0149] Formula 25: ,

[0150] Formula 26: ,

[0151] Formula 27: ,

[0152] Formula 28: ,

[0153] Formula 29: ,

[0154] Formula 30: ,

[0155] in, This represents the initial features of the first-level decoder. This indicates that the first-level decoder fuses skip connection features. This represents the channel attention enhancement features of the first-level decoder. This represents the spatial attention enhancement features of the first-level decoder. This represents the hybrid attention output features of the first-level decoder. This represents the local branch features of the first-level decoder. This indicates the axial branch feature of the first-level decoder. This indicates the axial fusion feature of the first-level decoder.

[0156] (2) Second-level decoder

[0157] The second-level decoder generally follows the core chain of the first-level decoder: "upsampling → skip fusion → attention extraction → axial fusion → normalization". The core differences are reflected in three aspects: First, the skip connection feature with the first-level decoder. The first method involves fusion; the second method replaces the first-level hybrid attention with frequency and spatial attention (LA_FSA) to obtain the results. Third, the axial convolution kernel is adjusted from 5×5 to 3×3 to reduce the axial receptive field, adapting to the low-scale feature resolution and capturing more detailed features. Finally, the second-level decoder axial fusion feature output from this layer is obtained. The specific formula is as follows:

[0158] Formula 31: ,

[0159] Formula 32: ,

[0160] Formula 33: ,

[0161] Formula 34: ,

[0162] Formula 35: ,

[0163] Formula 36: ,

[0164] in, This represents the upsampled features of the second-stage decoder. This represents the skip connection feature of the first-level encoder. This indicates that the second-level decoder fuses skip connection features. This indicates the second-level frequency domain enhancement feature. This represents the local branch features of the second-level decoder. This indicates the axial branch feature of the second-level decoder. This indicates the axial fusion feature of the second-level decoder.

[0165] Step S5: The output layer receives the axial fusion features from the second-level decoder. The core objective is to perform final purification and dimensionality adaptation on the original scale features. First, it uses depthwise separable convolutions consisting of 3×3 grouped convolutions combined with 1×1 convolutions. Then, it uses a sigmoid activation function to map the output, resulting in a final enhanced underwater image with the same dimensions as the input image. The specific formula is as follows:

[0166] Formula 37: ,

[0167] in, This indicates the final enhanced underwater image. This indicates the axial fusion feature of the second-level decoder.

[0168] This invention effectively reduces the number of model parameters and computational overhead by introducing depthwise separable convolution and grouped convolution structures. It enhances the model's ability to represent and enhance underwater degradation features by incorporating a dual-domain attention mechanism that collaborates with the frequency and spatial domains, along with the synergistic effect of axial fusion convolution. Furthermore, it ensures image detail restoration and structural integrity by combining an encoder-decoder structure and a skip connection mechanism. This invention guarantees the objectivity and realism of the enhancement results, achieving stable and efficient real-time underwater image enhancement processing at common image resolutions such as 720P and 1080P. It is particularly suitable for image enhancement tasks in complex underwater scenarios such as marine resource exploration, underwater edge machinery navigation, and underwater archaeology. It not only provides a novel and efficient solution for real-time underwater image enhancement but also offers a fundamental and practical solution for downstream tasks such as target detection.

Claims

1. A lightweight real-time underwater image enhancement method based on axial fusion and dual-domain attention, characterized in that, Includes the following steps: S1. Perform visual feature preprocessing on the original underwater image. Feature initialization is completed through depthwise separable convolution and batch normalization to obtain the preprocessed feature map. S2. Perform two-level multi-scale coding feature extraction on the preprocessed feature map. The first-level encoder uses frequency-space attention and spatial attention to extract frequency domain global features and spatial domain detail features in sequence. The weighted fusion of local branch features and axial branch features is completed by the axial fusion block. The downsampled output of the first-level encoder encoded features and corresponding skip connection features is then output. The second-stage encoder uses hybrid attention to dynamically fuse channel and spatial information, adjust the axial fusion convolution size to expand the receptive field, and downsample the output of the second-stage encoder's encoded features and corresponding skip connection features. S3. Input the encoded features of the second-level encoder into the bottleneck layer, and perform deep feature enhancement through channel attention and large kernel axial fusion block to obtain the enhanced features of the bottleneck layer. S4. Perform two-level multi-scale decoding feature restoration on the bottleneck layer enhancement features. Combine the skip connection features of the decoder and complete feature upsampling and detail fusion through the corresponding attention mechanism and axial fusion block to obtain the decoded output features. S5. Map the decoded output features to the number of channels of the original underwater image, and output the final enhanced underwater image through an activation function.

2. The method according to claim 1, characterized in that, Step S1 is as follows: Raw underwater images Pixel-level feature decomposition and channel dimension mapping are performed sequentially using 3×3 convolution and 1×1 convolution, followed by batch normalization to stabilize the feature distribution, resulting in a preprocessed feature map. The specific formula is as follows: Official 1: , in, This represents the preprocessed feature map. This represents the batch normalization function. Represents a 1×1 convolution. This represents a 3×3 convolution, where X represents the original underwater image. This indicates grouping by channel dimension.

3. The method according to claim 1, characterized in that, In step S2, the preprocessed feature map is used as the input feature map. The feature dimension is transformed through group convolution and 1×1 convolution to obtain the initial features of the first encoder. , for initial features The LA_FSA module processes the data as follows: First, it maps the spatial features to the frequency domain using a two-dimensional real-valued Fast Fourier Transform. After extracting the frequency domain amplitude features, it performs global adaptive average pooling to compress and obtain the global frequency domain energy statistics for each channel. Then, it generates frequency domain attention weights through 1×1 convolution and Sigmoid activation, and then... Multiply to output the first-level frequency domain enhancement feature. The specific formula is as follows: Official 3: , Official 4: , in, This represents the initial features of the first encoder. Represents a 1×1 convolution. This represents a 3×3 convolution. This represents the preprocessed feature map. This indicates grouping by channel dimension. This represents the first-level frequency domain enhancement feature. This represents the Sigmoid activation function. This represents the global average pooling function. This represents the two-dimensional real-valued fast Fourier transform. This indicates taking the amplitude in the frequency domain. Represents the ReLU activation function. These are learnable parameters, and their range is (0,1).

4. The method according to claim 3, characterized in that, First-level frequency domain enhancement features The global mean and global maximum features are concatenated along the channel dimension. Spatial correlation information of the concatenated features is extracted through 3×3 convolution, and then spatial attention weights are generated by Sigmoid activation. Multiply to output the first-level spatial detail enhancement features. The specific formula is as follows: Official 5: , in, This indicates the first-level airspace detail enhancement feature. This represents the first-level frequency domain enhancement feature. This represents the Sigmoid activation function. This represents a 3×3 convolution. Indicates channel splicing. This represents the global average pooling function. This represents the max pooling function.

5. The method according to claim 4, characterized in that, Enhanced features for first-level airspace details Perform a 3×1 height-axis convolution followed by a 1×3 width-axis convolution sequentially to output the first-level axial branch features. ;right Perform 3×3 grouped convolution to extract local spatial features, and then integrate channel information through 1×1 convolution to output the first-level local branch features. And introduce learning parameters With first-level airspace detail enhancement features Perform residual connection to obtain the first-level axial fusion feature. For the first-level axial fusion features After batch normalization, 1×1 convolution with channel adjustment, and then max pooling, the first-stage encoder's encoded features are output. Meanwhile, the features retained before pooling are used as skip connection features for the decoder, and the specific formula is as follows: Official 6: , Official 7: , Official 8: , Official 9: , in, This represents the first-level local branch feature. This indicates the first-level airspace detail enhancement feature. This indicates the first-order axial branch feature. This represents a 1×3 convolution. This represents a 3×1 convolution. This indicates the first-level axial fusion feature. , This represents learnable weighted parameters. This represents the encoding features of the first-level encoder. Represents the max pooling function. Represents a 1×1 convolution. This represents the batch normalization function.

6. The method according to claim 1, characterized in that, In step S2, the second-level encoder receives the first-level encoded features and performs dimensionality transformation. It employs hybrid attention with dynamic fusion of channel and spatial features, adjusting the convolutional kernel size of the axial fusion block to 5×5 to expand the receptive field. After batch normalization, 1×1 convolution, and max pooling, downsampling is completed, outputting the second-level encoder encoded features. Simultaneously, the features before pooling are retained as skip connection features for use by the decoder. The specific formula is as follows: Equation 10: ( ), Official 11: , Official 12: , Official 13: , Official 14: , Official 15: , Official 16: , Official 17: , in, This represents the initial features of the second-level encoder. This represents the encoding features of the first-level encoder. Represents a 1×1 convolution. This represents the channel attention enhancement feature of the second-level encoder. This represents the Sigmoid activation function. This represents the spatial attention enhancement features of the second-level encoder. This represents a 3×3 convolution. Indicates channel splicing. This represents the global average pooling function. Represents the max pooling function. This represents the hybrid attention output features of the second-level encoder. This represents the second-level local branch feature. This indicates a second-order axial branch feature. This represents a 1×5 convolution. This represents a 5×1 convolution. This represents the batch normalization function. This indicates the second-level axial fusion feature. This represents the encoding feature of the second-level encoder, where `gate` represents gating, and its value range is... , and These are learnable parameters.

7. The method according to claim 1, characterized in that, In step S3, the second-level encoder encodes the features. By inputting the bottleneck layer and maintaining the spatial resolution throughout, the bottleneck channel attention enhancement feature is first obtained by strengthening the optical feature correlation of the channel dimension through the channel attention module. Then, with the help of the 7x7 large core AFB, Extracting more spatial detail features yields Then and The features are concatenated and fused along the channel dimension. The number of feature channels in the concatenated feature is restored by a 1×1 convolution. Finally, the adjusted fused features are combined with the features encoded by the second-level encoder. Residual connections, output bottleneck layer, and final enhanced features The specific formula is as follows: Official 18: , Official 19: , Official 20: , Official 21: , in, This indicates the local branching characteristics of the bottleneck. Represents a 7×7 convolution. Represents a 1×1 convolution. This represents the spatial attention features of the second-level encoder. This indicates grouping by channel dimension. Indicates the axial branching characteristic of the bottleneck. This represents a 1×7 convolution. This represents a 7×1 convolution. This represents the encoding features of the second-level encoder. Indicates the axial fusion characteristics of the bottleneck. This indicates the final enhancement feature of the bottleneck layer. This indicates a bottleneck channel attention enhancement feature.

8. The method according to claim 1, characterized in that, In step S4, the bottleneck layer finally enhances the features. Transposed convolution Perform upsampling and restoration, and output the upsampled features of the first-stage decoder. , fusing skip connection features from the second-level encoder Complete the mesoscale details to obtain the first-level decoder fused skip connection features The system integrates a hybrid attention module, adopts the gating fusion logic of the hybrid attention module from the second-level encoder, and outputs the hybrid attention enhancement features of the first-level decoder. Then, through the AFB block, which uses 5×5 axial convolution, the axial fusion features of the first-level decoder are obtained. The specific formula is as follows: Official 22: , Official 23: , Official 24: ( ), Official 25: , Official 26: , Official 27: , Official 28: , Official 29: , Official 30: , in, This represents the upsampled features of the first-stage decoder. This represents the batch normalization function. Represents a 1×1 convolution. This represents a 2×2 transposed convolution. This indicates the final enhancement feature of the bottleneck layer. This indicates that the first-level decoder fuses skip connection features. This represents the skip connection feature of the second-level encoder. Represents the ReLU activation function. Indicates channel splicing. This represents the initial features of the first-level decoder. This indicates that the first-level decoder fuses skip connection features. This represents the channel attention enhancement features of the first-level decoder. This represents the spatial attention enhancement features of the first-level decoder. This represents a 3×3 convolution. This represents the Sigmoid activation function. This represents the global average pooling function. Represents the max pooling function. This represents the hybrid attention output feature of the first-level decoder, where `gate` represents the gating mechanism, and its value ranges from [value range missing]. , This represents the local branch features of the first-level decoder. This represents a 5×5 convolution. This indicates the axial branch feature of the first-level decoder. This represents a 1×5 convolution. This represents a 5×1 convolution. This indicates grouping by channel dimension. This indicates the axial fusion feature of the first-level decoder.

9. The method according to claim 8, characterized in that, The second-level decoder follows the core logic of the first-level decoder: upsampling → skip fusion → attention extraction → axial fusion → normalization, and it skips the connection features of the first-level encoder. The results were obtained by fusion using frequency and spatial attention LA_FSA. Simultaneously, the axial convolution kernel was adjusted from 5×5 to 3×3 to reduce the axial receptive field and output the axial fusion features of the second-level decoder. The specific formula is as follows: Official 31: , Official 32: , Official 33: , Official 34: , Official 35: , Official 36: , in, This represents the upsampled features of the second-stage decoder. This represents the skip connection feature of the first-level encoder. This indicates that the second-level decoder fuses skip connection features. This indicates the second-level frequency domain enhancement feature. This represents the two-dimensional real-valued fast Fourier transform. This indicates taking the amplitude in the frequency domain. This represents a 1×3 convolution. This represents a 3×1 convolution. This represents the local branch features of the second-level decoder. This indicates the axial branch feature of the second-level decoder. This indicates the axial fusion feature of the second-level decoder.

10. The method according to claim 9, characterized in that, In step S5, the output layer receives the axial fusion features from the second-level decoder. First, the image is processed by depthwise separable convolution consisting of 3×3 grouped convolutions and 1×1 convolutions, and then mapped by the Sigmoid activation function to output the final enhanced underwater image. The specific formula is as follows: Official 37: , in, This indicates the final enhanced underwater image. This indicates the axial fusion feature of the second-level decoder.

Images