Underwater image enhancement method based on UFMFormer model

Abstract

The invention relates to the field of image enhancement, in particular to an underwater image enhancement method based on a UFMFormer model. The method comprises the steps that a UFMFormer model is constructed, the UFMFormer model is of an encoder-decoder structure, local feature extraction and detail recovery are carried out on a high-resolution layer by using a Transform module, and global modeling and decoding are carried out on a low-resolution layer by using a Mama module; using the underwater image data set to train a UFMFormer model, and obtaining a trained UFMFormer model; and enhancing the underwater image by using the trained UFMFormer model. According to the method, the advantages of Transform in local enhancement and the capability of Mama in global modeling are fully played, so that the problem of image distortion of underwater images caused by light scattering, color attenuation and low visibility is solved.

Description

TECHNICAL FIELD

[0001] The present application relates to the field of image enhancement, in particular to an underwater image enhancement method based on a UFMFormer model. BACKGROUND

[0002] Underwater image enhancement (UIE) technology is an indispensable technology for obtaining underwater images and investigating underwater environments, and has a wide range of applications in fields such as marine exploration, biology, archaeology, underwater robots, etc. However, due to the combined effects of various factors, especially light scattering caused by particle interaction, color fading caused by light wavelength absorption, turbidity of suspended particles, water depth and temperature effects, light penetration changes, and blurring and surface interference caused by motion, underwater images are often of poor quality. Therefore, innovations related to underwater image enhancement are of great significance for improving the visual quality and value of images and accurately understanding the underwater world.

[0003] Deep learning-based methods have revolutionized the field of image enhancement. Unlike traditional methods, deep learning can learn directly from data, overcoming limitations and adapting to new scenarios. These methods excel in extracting complex image features and generalizing the relationship between degraded and clean images, such as the impressive performance of convolutional neural networks (CNNs) in underwater image enhancement in recent years. However, CNNs have limitations, such as limited receptive fields that prevent modeling long-range pixel dependencies, and static weight convolution filters that cannot adapt to input images during inference. To address these issues, some researchers have proposed the Transformer based on the self-attention mechanism, which has made significant progress in natural language processing (NLP) and has been extended to computer vision tasks. The emergence of Vision Transformer allows for the capture of local and global information. This ability to interact across positions gives it an advantage when dealing with tasks that require a broader context and complex relationships within images. However, its quadratic computational complexity relative to sequence length poses challenges when dealing with high-resolution images.

[0004] Therefore, there is an urgent need to solve the above technical problems. SUMMARY

[0005] The technical problem to be solved by the present application is to overcome the defects of the prior art and provide an underwater image enhancement method based on a UFMFormer model, which fully utilizes the advantages of Transformer in local enhancement and the capabilities of Mamba in global modeling to solve the image distortion problems caused by light scattering, color fading, and low visibility in underwater images.

[0006] The technical solution adopted by the present application to solve the technical problem is: an underwater image enhancement method based on a UFMFormer model, comprising: Construct the UFMFormer model. The UFMFormer model has an encoder-decoder structure. It uses the Transformer module for local feature extraction and detail recovery in the high-resolution layer, and the Mamba module for global modeling and decoding in the low-resolution layer. The UFMFormer model is trained using the underwater image dataset to obtain a trained UFMFormer model; Underwater images are enhanced using the trained UFMFormer model.

[0007] Furthermore, the encoder-decoder structure has three layers, using Transformer modules in the first and second layers respectively, and a Mamba module in the third layer; the decoder part introduces the FPN module in the second layer for multi-scale feature fusion, and also introduces the CBAM mechanism to enhance the key feature extraction capability.

[0008] Furthermore, the specific formula of the second layer of the decoder is:

[0009] Where, represents feature map upsampling, represents the output of the third layer of the decoder part, Represents the initial feature map of the second layer of the decoder part; , Where, Represents feature concatenation operation; Represents the output of the second layer of the encoder; Represents 1×1 convolution; Represents the feature map after the second layer fusion of the decoder part; , Representation layer normalization; Indicates that the Mamba module is used for processing; Indicates that it is processed by the feedforward network; , Where, Indicates that it has been processed by the channel plus spatial attention mechanism module;

[0010] Where, Indicates fusion The weight for 1×1 mapping;

[0011] Where, Indicates that the i-th Transformer module is used for operation; a parameter set representing the second layer of the Transformer module; a number of the second layer of the decoder part Transformer module; an output of the second layer of the decoder part.

[0012] Further, the encoder part and the decoder part respectively introduce an underwater color cast correction module in the Transformer module; wherein, The working process of the underwater color cast correction module of the encoder part is: First, separate the input features according to the RGB channel: Then, input the features of each channel into a DFFN module for processing respectively:

[0013] In the formula, denotes a layer normalization operation; denotes a frequency domain enhanced feedforward network processing; denotes the features input into the DFFN module; denotes the output features of the DFFN module; Finally, the processed RGB channels are combined; The underwater color cast correction module of the decoder part is compared with the underwater color cast correction module of the encoder part. Before inputting into the DFFN module, the features of each channel are first processed by the GFLSA module:

[0014] In the formula, denotes a fusion operation of global frequency domain attention and local spatial attention; denotes the features input into the GFLSA module; denotes the output features of the GFLSA module.

[0015] Further, the frequency domain enhanced feedforward network adopts a learnable frequency domain filter, and the specific formula is:

[0016]

[0017]

[0018]

[0019]

[0020] In the formula, denotes the input feature map; denotes a 1x1 convolution operation; denotes a weighted Fourier transform process; denotes a GELU activation function process; denotes a weight, respectively; denotes an output feature map.

[0021] Further, the specific formula of the fusion operation of the global frequency domain attention and the local spatial attention is:

[0022]

[0023]

[0024]

[0025] wherein, denotes an input image feature, W t and W a denote corresponding weight parameters, respectively, denotes a frequency domain representation; denotes an inverse Fourier transform operation, and is a learnable dynamic weight; denotes an output; denotes a GELU activation function process; denotes a 1x1 convolution operation; denotes a 3x3 convolution operation.

[0026] Further, the loss function adopted in the training process is: Pixel-level loss: wherein, is the i-th pixel value of the enhanced image predicted by the model, is the i-th pixel value of the target real image, and N is the total sum of image pixels, is a pixel loss weight; Frequency domain loss: wherein, and denote the two-dimensional fast Fourier transform of the predicted image and the target image; is a frequency domain loss weight; Perception loss: wherein, is a VGG16 network i-th layer feature extraction function, and are the i-th value of the predicted image and the target image in the i-th layer feature map; is a perception loss weight; Color loss: wherein is the predicted Lab color space (L, a, b) of the i-th pixel of the image, is the Lab classification of the i-th pixel of the target image; total loss function: .

[0027] After the above technical scheme, the UFMFormer model constructed by the application fully gives play to the advantages of the Transformer in local enhancement and the ability of the Mamba in global modeling, adopts the Transformer module for local feature extraction on the high-resolution layer, utilizes the self-attention mechanism to efficiently model the spatial dependency, thereby enhancing the detail restoration capability, while the Mamba module is adopted for global modeling on the low-resolution layer, which utilizes the state space model to process long sequence data, ensuring the global color consistency. The application effectively balances the detail restoration capability and the global context consistency, can enhance the underwater image, and significantly improves the image color and clarity. BRIEF DESCRIPTION OF DRAWINGS

[0028] Figure 1 is the framework diagram of the UFMFormer model of the application; Figure 2 is the schematic diagram of the underwater color cast correction module of the encoder part of the application; Figure 3 is the schematic diagram of the underwater color cast correction module of the decoder part of the application; Figure 4 is the schematic diagram of the DFFN module of the application; Figure 5 is the schematic diagram of the GFLSA module of the application; Figure 6 is the schematic diagram of the Mamba module of the application; Figure 7 is the visual comparison result of the model in the application and the traditional model on the UIEB dataset; Figure 8 is the visual comparison result of the model in the application and the traditional model on the Test-L400 dataset; Figure 9 is the visual comparison result of the model in the application and the traditional model on the U45 and C60 two datasets; Figure 7 and Figure 8In the figure, (a) is the Raws model, (b) is the UWCNN model, (c) is the Water-Net model, (d) is the U-shape model, (e) is the PWRNet model, (f) is the DM-water model, (g) is the UDAformer model, (h) is the WF-Diff model, (i) is the model in the application, and (j) is the real image GT. Figure 9 In the figure, (a) is the original image, (b) is the UWCNN model, (c) is the Water-Net model, (d) is the U-shape model, (e) is the PWRNet model, (f) is the DM-water model, (g) is the UDAformer model, (h) is the WF-Diff model, and (i) is the model in the application. DETAILED DESCRIPTION

[0029] As shown in the figure, a water image enhancement method based on a UFMFormer model comprises the following steps: Figure 1 The UFMFormer model is constructed, and the UFMFormer model is an encoder-decoder structure, uses a Transformer module for local feature extraction and detail restoration at a high-resolution layer, and uses a Mamba module for global modeling and decoding at a low-resolution layer. This hybrid architecture combines the local perception ability of the Transformer and the long sequence modeling ability of the Mamba, and can better handle the complex texture and global information of underwater images. The UFMFormer model is trained using an underwater image dataset to obtain a trained UFMFormer model. The trained UFMFormer model is used to enhance underwater images.

[0030] More specifically, as shown in the figure, the encoder-decoder structure has three levels, uses a Transformer module at the first and second levels, and uses a Mamba module at the third level; the decoder part introduces an FPN module at the second level for multi-scale feature fusion, and also introduces a CBAM mechanism to enhance the key feature extraction capability. The specific process and formula are as follows: Figure 1 The input image I is converted to a feature space through overlapping embedding blocks and is used as the input of the first layer of the encoder part:

[0031] In the formula, x is the feature map with a dimension D=48, is the feature map with a dimension D=48, is the convolution kernel weight, represents a 3x3 convolution.​​

[0032] For the encoder part:

[0033] As the local details (e.g. edges and textures) of underwater images are vulnerable to scattering interference. First, capture spatial dependencies at high resolution through self-attention mechanism, strengthen detailed features, and pass reliable information to subsequent levels. Then, the features are down-sampled into the middle resolution layer.

[0034]

[0035]

[0036] This stage focuses on medium-scale features (e.g. object contours), enhances the modeling ability of complex underwater scenes through Transformer, and finally down-samples to the low-resolution layer.

[0037]

[0038]

[0039] wherein, , , respectively represent the parameter set of the first layer of the Transformer module, the second layer of the Transformer module and the third layer of the Mamba module; , respectively represent the weight parameters of the convolution kernel of the first down-sampling and the second down-sampling; , , respectively represent the input of the first layer, the second layer and the third layer of the encoder part; , , respectively represent the output of the first layer, the second layer and the third layer of the encoder part; , , respectively represent the number of the first layer of the Transformer module, the number of the second layer of the Transformer module and the number of the third layer of the Mamba module; represents processing by the Transformer module; represents up-sampling for feature size adjustment; represents down-sampling operation by 3x3 convolution; wherein, As shown in detail, Figure 6 the Mamba module is as shown in the part in the dashed box in Figure 6 , Layer normalization is used to normalize the input feature map X along the channel dimension, stabilizing the training process. Mamba module is used to model long sequences and capture global dependencies. FFN is used to enhance the expressive power of features. Because the global color attenuation and uneven illumination of underwater images require long-range modeling, the state space mechanism of Mamba can efficiently process low-resolution features and combine FFN to ensure global consistency.

[0040] For the decoder part, low-resolution layers are connected through a skip connection to better transfer features and ensure efficient use of global information in low-resolution layers.

[0041]

[0042]

[0043] FPN modules are introduced in the middle resolution layer, using By fusing multi-scale features, the model's ability to capture different scale features is enhanced, especially for the multi-scale details commonly found in underwater images. CBAM mechanism is used in the decoder to further enhance the extraction ability of key features.

[0044]

[0045]

[0046]

[0047]

[0048] Fuse and Transformer are used to restore fine textures and ensure clear details.

[0049]

[0050]

[0051]

[0052] Finally, the features are further optimized through a refining stage and the output is generated through a residual connection:

[0053]

[0054] wherein ,​​ are weight parameters, respectively; , , , are the number of third-layer Mamba modules, the number of second-layer Transformer modules, the number of first-layer Transformer modules, and the number of refining-stage Transformer modules, respectively; indicates that the fusion features are further modeled to improve the global and dynamic information modeling capability; indicates that a concatenation operation is adopted; indicates that the channel-plus-space attention mechanism module is processed, which is used to enhance the key features and suppress the background interference; indicates a 1x1 convolution; indicates the finally generated enhanced image; ; indicates up-sampling, which can adopt bilinear interpolation, etc.

[0055] In the embodiment, considering that different color channels have different degrees of attenuation in the underwater environment, especially the red channel has the most serious attenuation due to the longer wavelength, while the attenuation of the blue and green channels is lighter. Therefore, the embodiment preferably separates the input features by RGB channels, decomposes the image into three independent sub-feature maps of R, G, and B, and each channel dimension is 1 / 3. Through independent optimization strategy, the model can learn color compensation of the red channel and detail recovery of the blue-green channel respectively, so as to more effectively solve the color distortion problem of underwater images, as shown in Figure 2 and Figure 3 , that is, the underwater color cast correction module is introduced in the Transformer module of the encoder part and the decoder part, respectively.

[0056] As shown in Figure 2 , the working process of the underwater color cast correction module of the encoder part is as follows: First, the input features are separated according to the RGB channels: Then, the features of each channel are independently input into a DFFN module for processing:

[0057] In the formula, indicates a layer normalization operation; indicates a frequency domain enhancement feedforward network processing; indicates the features input into the DFFN module; indicates the output features of the DFFN module; Finally, the processed RGB channels are combined.

[0058] AsFigure 3 As shown in the figure, the underwater color correction module in the decoder part is different from the underwater color correction module in the encoder part. The features of each channel are first processed by the GFLSA module before being input into the DFFN module. The specific process of the underwater color correction module in the decoder part is as follows: This module first inputs the features Separate by RGB channels: ,During the processing, the features of each channel are independently input into the GFLSA module and DFFN module for ,processing.

[0059]

[0060] Where, Representation layer normalization operation; It represents the processing of frequency domain enhanced feedforward network; Represents the fusion operation of global frequency domain attention and local spatial attention.

[0061] This module can capture global illumination information and enhance local details. Since the three channels are processed separately, this method can avoid the interference of different color channels, allowing the model to more accurately adjust the color balance of underwater images while reducing color drift. and This operation can make the enhanced underwater image more consistent with the real scene. The red channel can better restore the original color information after processing, while the blue and green channels can retain more details.

[0062] Finally, merge the processed RGB channels: .

[0063] This ensures effective complementarity between information from each channel, thereby generating underwater images with balanced colors and natural visuals.

[0064] The DFFN module is introduced in detail below.

[0065] In the feedforward enhancement network part, this embodiment adopts a frequency domain enhancement feedforward network, breaking through the limitation of traditional feedforward networks that only rely on time domain calculations. Because not all low-frequency and high-frequency information are helpful for underwater image recovery, the frequency domain enhancement feedforward network in this embodiment can use a learnable frequency domain filter to perform feature enhancement operations in the Fourier transform domain, which can adaptively retain important frequency information in underwater images, improve image clarity, and make the enhanced image sharper in edges and textures. Figure 4 As shown, the specific formula is:

[0066]

[0067]

[0068]

[0069]

[0070] wherein, denotes an input feature map; denotes a 1x1 convolution operation; denotes a weighted Fourier transform; denotes a GELU activation function processing; denotes a weight, respectively; denotes an output feature map.

[0071] The GFLSA module is described in detail below.

[0072] In an underwater environment, due to the attenuation and scattering of light, distant targets are often severely color-shifted, and traditional convolutional neural networks or Transformers are limited by the receptive field and cannot fully capture long-distance dependencies. Therefore, the GFLSA module is designed to model long-distance dependencies and local texture features by fusing global frequency domain attention and local spatial attention, as shown in Figure 5 The specific process is as follows: First, calculate the query Q and the key K;

[0073] Then, perform correlation modeling in the frequency domain, so that the network can adaptively correct the global color distribution and compensate for color attenuation caused by the underwater environment. The features output by the frequency domain attention are mapped back to the spatial domain through inverse Fourier transform (iFFT):

[0074] Meanwhile, in order to prevent the problem of high-frequency information loss, local spatial attention is added to the structure, which extracts detailed information through local window convolution, so that the enhanced image can retain clear edges and texture features

[0075] Finally, fuse the global frequency domain attention and the local spatial attention:

[0076] wherein, denotes an input image feature, W t and Wa respectively represent the corresponding weight parameters, represents a frequency domain representation; represents an inverse Fourier transform operation, represents an activation function, enhancing non-linear expression capability, and are learnable dynamic weights; represents an output, which is normalized and projected.

[0077] The entire GFLSA module can preserve details while correcting colors, thereby improving the overall visual quality of underwater images

[0078] It should be noted that in the task of underwater image enhancement, uneven illumination, color distortion and blurring are major challenges. The embodiment proposes a frequency-space hybrid attention mechanism, and combines a frequency domain enhancement feedforward network to realize adaptive enhancement of underwater images, which can effectively improve the global illumination correction capability and local detail recovery capability of the image.

[0079] The loss function is described in detail below.

[0080] In the task of underwater image enhancement, in order to improve image quality and restore true colors and details, we designed a variety of loss functions to constrain the training process of the model. Including pixel loss (Pixel Loss), frequency domain loss (FFTLoss), perceptual loss (Perceptual Loss) and color loss (Color Loss).

[0081] Pixel loss calculates the pixel-by-pixel difference between the predicted image and the target image to ensure that the enhanced image is consistent with the real image at the pixel level, i.e. L1 loss, mathematically represented as follows: where is the i-th pixel value of the model-predicted enhanced image, is the i-th pixel value of the target real image, and N is the total number of image pixels, is the pixel loss weight; In underwater image enhancement, L1 loss helps to preserve the basic structure and details of the image, avoiding blurring or distortion caused by underwater scattering and absorption.

[0082] Frequency domain loss compares the difference between the predicted image and the target image in the Fourier transform domain to constrain the high-frequency information (such as edges and textures) generated by the model. It is defined as follows: where and represent the two-dimensional fast Fourier transform of the predicted image and the target image; The frequency domain loss weight.

[0083] Underwater images often lose high-frequency details due to light scattering. The FFT loss enhances the model's ability to recover edges and textures by constraining the frequency domain.

[0084] The perceptual loss uses a pre-trained VGG16 network to extract image features and measures the similarity between the predicted image and the target image in high-level semantics. The mathematical representation is as follows: where is the VGG16 network's l-th layer feature extraction function, and is the i-th value of the predicted image and the target image in the l-th layer feature map. The perceptual loss weight.

[0085] The perceptual function focuses on the semantics and perceptual quality of the image, helping to restore a more natural visual effect in underwater image enhancement and making up for the neglect of global structure by pixel-level loss.

[0086] The color loss is based on the Lab color space and measures the difference in color consistency between the predicted image and the target image. It is defined as follows: where is the i-th pixel of the predicted image in the Lab color space, is the Lab classification of the i-th pixel of the target image.

[0087] Because underwater images often exhibit color cast due to water absorption, the color loss enhances the model's ability to recover true colors by constraining the difference in the Lab space.

[0088] The total loss function of the model is defined as

[0089] The weights of each loss can be: The weight coefficients of each loss function are based on the actual contribution of different loss functions in the numerical space, training convergence, and performance improvement, and are set through experimental optimization. This setting method is conducive to more effectively balancing structure recovery, detail enhancement, semantic quality, and color restoration in underwater image enhancement tasks.

[0090] In summary, the UFMFormer model constructed in this embodiment fully utilizes the advantages of the Transformer in local enhancement and the capabilities of the Mamba in global modeling. The Transformer module is used for local feature extraction on high-resolution layers, and the self-attention mechanism is used to efficiently model spatial dependencies, thereby enhancing the detail restoration capability. The Mamba module is used for global modeling on low-resolution layers, and its state space model is used to process long sequence data to ensure global color consistency. To further highlight key features and balance multi-level information transmission, ensuring the reconstruction quality from low-resolution global information to high-resolution details, the model uses a multi-scale feature fusion mechanism to integrate the encoder and decoder features in depth by combining the Mamba and convolution operations, supplemented by CBAM (Channel-Spatial Attention) and FPN (Feature Pyramid Network) designs. The RGB channel separation strategy designed in this embodiment effectively solves the problems of color distortion and detail blur. By decomposing the input features into independent red, green, and blue channels, the uneven attenuation characteristics of each channel in the underwater environment are utilized to specifically strengthen the severely damaged color information. The decomposed features are then passed through the GLFSA module to capture global lighting characteristics using frequency domain analysis and highlight local details using spatial domain attention, achieving a dynamic balance between global and local features. Finally, the DFFN linear transformation and frequency domain modeling improve the detail perception ability. This design not only optimizes the accuracy of color restoration but also significantly improves the clarity and realism of underwater images.

[0091] The above embodiments are described in detail below in conjunction with specific experiments.

[0092] Implementation details: The UFMformer model proposed in this embodiment is implemented using the Pytorch framework on an NVIDIA RTX3090 and trained for 3000 epochs, with a training time of approximately 48h. The training uses the ADAW optimizer (β1=0.9, β2=0.9, weight decay of 1e-3) and L2 loss, with a patch size of 256x256 and a batch size of 2. The initial learning rate is 1e-3.

[0093] Dataset: The UIEB dataset and the LSUI dataset are used to train and evaluate the model in this embodiment. The UIEB dataset contains 890 real underwater images and corresponding reference images, with 800 images randomly selected for training and 90 images for testing. The LSUI dataset is randomly divided into 3879 training images and 400 test images. During network training, all images are adjusted to a fixed size (256x256).

[0094] Evaluation metrics: For the test dataset with reference images, we adopted the peak signal-to-noise ratio (PSNR) and structural similarity index (SSIM) metrics for full-reference evaluation. These two metrics reflect the closeness to the reference, where a higher PSNR value indicates a closer image content, and a higher SSIM value indicates a more similar structure and texture. In addition, we included the LPIPS and FID metrics in the full-reference image evaluation. LPIPS is a deep neural network-based image quality metric that assesses the perceptual similarity between an image and a reference image. FID measures the distance between the distribution of real images and generated images. Lower LPIPS and FID scores indicate that the UIE method is more effective. For images in the non-reference test dataset, we adopted the non-reference evaluation metrics UCIQE and UIQM, where higher UCIQE or UIQM scores indicate better human visual perception. In addition, we also adopted NIQE as an evaluation metric, where a lower NIQE value indicates higher visual quality.

[0095] Comparison methods: We compared the UFMFormer with eight state-of-the-art (SOTA) underwater image enhancement (UIE) methods, including UWCNN, WaterNet, U-color, U-shape, PRWNet, DM-water, UDAformer, and WF-Diff. The results of all the above methods were based on the publicly released code by their authors and strictly followed the same experimental settings in all evaluations.

[0096] Full-reference evaluation: We used the Test-U90 and Test-L400 subsets for full-reference evaluation, and the quantitative comparison results and visual comparison graphs are listed in Table 1, Figure 7 and Figure 8 , respectively. As shown in Table 1, the proposed method achieved the best PSNR performance on both datasets and obtained the second-best performance in the SSIM metric.

[0097] Table 1

[0098] In the UIEB dataset, our method achieved a PSNR of 24.03 dB, which is 0.17 dB higher than the WF-Diff method; the SSIM metric was 0.8759, slightly lower than the UDAformer's 0.8780. In addition, it achieved the best score in FID (27.75) and ranked second in LPIPS (0.1497), only behind DM-water (0.1436).

[0099] Figure 7Visualization comparison results on the UIEB dataset are shown, covering multiple challenging underwater scenes. Compared with traditional methods such as UWCNN and Water-Net, our method effectively improves color fidelity and image contrast. Compared with U-shape and PRWNet, our method avoids color shift or oversaturation, and can preserve real colors and clear texture details. As shown in Figure 7 (i), our enhanced images are closest to the real images (GT, Figure 7 (j)).

[0100] In the LSUI dataset, our method achieves a PSNR of 29.00 dB and an SSIM of 0.9232, which are 0.52 dB and 0.011 higher than the previous best method, respectively. In terms of FID (26.43) and LPIPS (0.1081) indicators, it also performs well, ranking second in both. Figure 8 Visualization results in Table 2 further demonstrate the advantages of our method in maintaining image clarity, texture details, and natural color tones, significantly outperforming other deep learning-based methods that often suffer from color drift or over-smoothing.

[0101] These results verify the superior performance and good generalization ability of the proposed method on two full-reference datasets, fully demonstrating its application potential in underwater image enhancement tasks.

[0102] No-reference evaluation: In the no-reference evaluation experiment, we selected the U45 and C60 datasets for testing. The quantitative comparison results are shown in Table 1, and the visualization results are shown in Figure 9 (the first three rows are for the U45 dataset, and the last three rows are for the C60 dataset).

[0103] On the U45 dataset, our method achieved a score of 3.189 and 0.627 in UIQM and UCIQE, respectively, ranking second in both. Although DM-water and UDAformer achieved the highest scores in UIQM (3.212) and UCIQE (0.634), respectively, our method performed more evenly in contrast, clarity, and color restoration. As shown in Figure 9 (i), our method effectively corrected color deviation and removed water mist from the image, generating clear and visually pleasing images. In contrast, traditional methods such as UWCNN and Water-Net often produce images with blurring or color shift, while deep learning models such as U-shape and PRWNet may introduce artifacts or over-enhancement problems.

[0104] On the C60 dataset, our method also achieves sub-optimal results on the UIQM (3.105) and UCIQE (0.605) metrics, very close to the optimal results of DM-water and WF-Diff. Figure 9 The visualization of the lower part of the results shows that our method can restore a more natural underwater scene, significantly improve image visibility, and maintain realistic color performance. The comparison methods often have the problem of poor defogging effect (such as UWCNN) or image oversaturation (such as PRWNet, UDAformer).

[0105] In summary, the proposed method achieves superior visual enhancement effect and robust quantitative indicators on multiple datasets, verifying its effectiveness and robustness in the underwater image enhancement task.

[0106] With the above ideal embodiments according to the present application as the inspiration, through the above description, relevant staff can make various changes and modifications without deviating from the scope of the technical idea of the present application. The technical scope of the present application is not limited to the contents in the specification, and must be determined according to the scope of the claims.

Claims

1. An underwater image enhancement method based on the UFMFormer model, characterized in that: include: Construct the UFMFormer model. The UFMFormer model has an encoder-decoder structure. It uses the Transformer module for local feature extraction and detail recovery in the high-resolution layer, and the Mamba module for global modeling and decoding in the low-resolution layer. The UFMFormer model is trained using the underwater image dataset to obtain a trained UFMFormer model; Underwater images are enhanced using the trained UFMFormer model.

2. The underwater image enhancement method based on the UFMFormer model according to claim 1, characterized in that: The encoder-decoder structure has three layers, using Transformer modules in the first and second layers respectively, and a Mamba module in the third layer; the decoder part introduces an FPN module in the second layer for multi-scale feature fusion, and also introduces the CBAM mechanism to enhance the ability to extract key features.

3. The underwater image enhancement method based on the UFMFormer model according to claim 2, characterized in that: The specific formula of the second layer of the decoder is: Where, represents feature map upsampling, represents the output of the third layer of the decoder part, Represents the initial feature map of the second layer of the decoder part; , Where, Represents feature concatenation operation; Represents the output of the second layer of the encoder; Represents 1×1 convolution; Represents the feature map after the second layer fusion of the decoder; , Representation layer normalization; Indicates that the Mamba module is used for processing; Indicates that it is processed by the feedforward network; , Where, Indicates that it has been processed by the channel plus spatial attention mechanism module; Where, Indicates fusion The weight for 1×1 mapping; Where, Indicates that the i-th Transformer module is used for operation; Represents the parameter set of the Transformer module of the second layer; Indicates the number of Transformer modules in the second layer of the decoder; Represents the output of the second layer of the decoder part.

4. The underwater image enhancement method based on the UFMFormer model according to claim 1, characterized in that: The encoder part and the decoder part respectively introduce underwater color correction modules into the Transformer module; The working process of the underwater color deviation correction module in the encoder part is as follows: First, separate the input features according to RGB channels: Then, the features of each channel are independently input into a DFFN module for processing: Where, Representation layer normalization operation; It represents the processing of frequency domain enhanced feedforward network; Represents the features of the input DFFN module; Represents the output features of the DFFN module; Finally, the processed RGB channels are merged; Compared to the underwater color skew correction module in the encoder, the features of each channel in the decoder are processed by the GFLSA module before being input into the DFFN module: Where, Represents the fusion operation of global frequency domain attention and local spatial attention; Indicates the characteristics of the input GFLSA module; Represents the output characteristics of the GFLSA module.

5. The underwater image enhancement method based on the UFMFormer model according to claim 4, characterized in that: The frequency domain enhanced feedforward network uses a learnable frequency domain filter. The specific formula is: Where, Represents the input feature map; Represents a 1×1 convolution operation; represents weighted Fourier transform processing; Indicates processing by GELU activation function; Respectively represent weights; Represents the output feature map.

6. The underwater image enhancement method based on the UFMFormer model according to claim 4, characterized in that: The specific formula for the fusion operation of global frequency domain attention and local spatial attention is: Where, represents the input image features, W t and W a Represent the corresponding weight parameters, represents the frequency domain representation; represents the inverse Fourier transform operation, and is a learnable dynamic weight; Indicates output; Indicates processing by GELU activation function; Represents a 1×1 convolution operation; Represents a 3×3 convolution operation.

7. The underwater image enhancement method based on the UFMFormer model according to claim 1, characterized in that: The loss function used in the training process is: Pixel-level loss: ,in is the i-th pixel value of the enhanced image predicted by the model, is the i-th pixel value of the target real image, N is the total number of image pixels, is the pixel loss weight; Frequency domain loss: ,in and Represents the two-dimensional fast Fourier transform of the prediction image and the target image; is the frequency domain loss weight; Perceptual Loss: ,in is the feature extraction function of the lth layer of the VGG16 network, and is the i-th value of the feature map of the predicted image and the target image in the l-th layer; is the perceptual loss weight; Color loss: ,in is the brightness (L) and chrominance (a, b) components of the i-th pixel in the predicted image in Lab color space, is the Lab classification of the i-th pixel of the target image; Total loss function: .

Images