Underwater image enhancement method based on prior information guidance
By constructing an underwater image enhancement network guided by prior information, the problems of unstable training and computational complexity in existing technologies are solved, achieving more efficient underwater image enhancement results, reducing iteration steps and computational load, and improving the model's inference speed and image enhancement quality.
Patent Information
- Application Number
- CN202511468704.6
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-10-15
- Publication Date
- 2026-02-24
AI Technical Summary
Existing underwater image enhancement algorithms based on diffusion models treat underwater image enhancement tasks as ordinary image generation tasks, ignoring prior knowledge of underwater image enhancement tasks, resulting in unstable training and high computational complexity.
An underwater image enhancement network is constructed, including the Underwater Image Enhancement Compact Prior Representation Extraction Network (UIE-CPREN), the Multi-Color Space Feature Extraction Network (MSFEN), and the Multi-Color Space Feature Fusion Module (MSFFM). Accurate compact prior representations are generated through a diffusion model to guide subsequent underwater image enhancement. The network is divided into two stages: pre-training and joint training, which reduces the number of iteration steps and computational complexity.
It improves the efficiency and accuracy of underwater image enhancement algorithms, reduces computational complexity and inference time, and makes full use of the guidance of compact prior representations, achieving faster inference speed and better enhancement results.
Smart Images
Figure CN121563807A_ABST
Abstract
Description
Technical Field
[0001] This invention belongs to the field of underwater image enhancement technology, specifically relating to an underwater image enhancement method guided by prior information. Background Technology
[0002] With the encouraging results achieved by various deep learning-based underwater image enhancement (UIE) algorithms, generative adversarial networks (GANs)-based underwater image enhancement algorithms can generate clear images. However, due to problems such as model collapse during training, the training process is unstable. Recently, denoising diffusion probabilistic models have attracted considerable attention due to their stable training process and excellent image generation performance, and they have also been applied to underwater image enhancement tasks. However, current diffusion-based underwater image enhancement algorithms treat UIE as a complete image generation task, gradually denoising random Gaussian noise and generating clear underwater images through a large number of reverse denoising steps (100-2000 steps). This blurs the distinction between underwater image enhancement tasks and ordinary image generation tasks. Underwater image enhancement tasks have certain prior knowledge and only need to adjust the detailed information based on the prior representation provided by the corresponding labels of the underwater degraded image. Therefore, the traditional diffusion model generation paradigm is inefficient for underwater image enhancement tasks. Summary of the Invention
[0003] To address the aforementioned problems, this invention provides an underwater image enhancement method guided by prior information.
[0004] The technical solution adopted in this invention is:
[0005] An underwater image enhancement method guided by prior information includes the following steps:
[0006] S1. Prepare the underwater image enhancement dataset. Divide the publicly available underwater image enhancement dataset UIEB into a training set and a test set to provide a data foundation for the subsequent training of underwater image enhancement models.
[0007] S2. Construct an underwater image enhancement network, including an underwater image enhancement compact prior representation extraction network, a multi-color space feature extraction network, a multi-color space feature fusion module, and a diffusion model;
[0008] S3. Perform model training, and divide the model training process into two stages: pre-training and joint training;
[0009] S4. During the model inference phase, a low-dimensional compact vector is randomly sampled from the standard normal distribution, and an accurate compact prior representation is generated through the diffusion model to guide the subsequent underwater image enhancement network.
[0010] Compared with the prior art, the present invention has the following advantages:
[0011] 1. In this invention, MSFEN extracts high-level representations from the RGB and LAB color spaces to obtain richer feature information. Simultaneously, a ResWG module is proposed to reduce the network parameters and computational cost while maintaining MSFEN's feature extraction capabilities and preserving image detail features. The MSFFM module achieves complementarity of different color space features through differential fusion of extracted multi-color space features, improving the enhancement effect of subsequent UIE tasks.
[0012] 2. This invention generates accurate and compact prior representations through a diffusion model to guide subsequent underwater image enhancement networks. Since the compact prior representation is only a low-dimensional compact vector, the diffusion model only requires a small number of denoising steps to generate accurate prior representations, greatly improving the inference speed of the model.
[0013] 3. In the joint training process, this invention jointly optimizes the parameters of the diffusion model and the subsequent enhancement network, making fuller use of the guidance of compact prior representations.
[0014] 4. Compared to generating images, the present invention uses a diffusion model to more accurately generate compact prior representations, which omits a large number of iterative steps of the diffusion model and reduces the computational complexity and inference time of the underwater image enhancement algorithm. Attached Figure Description
[0015] Figure 1 This is a flowchart illustrating the present invention;
[0016] Figure 2 This is a schematic diagram of the multi-color space feature fusion module of the present invention;
[0017] Figure 3 This is a schematic diagram illustrating the enhanced results of this invention on the UIEB dataset;
[0018] Figure 4 This is a schematic diagram of the enhanced result set of the present invention on the U45 dataset. Detailed Implementation
[0019] To better understand the purpose, structure, and function of this invention, the invention will be described in further detail below with reference to the accompanying drawings.
[0020] like Figure 1-2 As shown, the present invention provides an underwater image enhancement method guided by prior information, the method comprising the following steps:
[0021] S1. Prepare the underwater image enhancement dataset. Divide the publicly available underwater image enhancement dataset UIEB into a training set and a test set to provide a data foundation for the subsequent training of underwater image enhancement models.
[0022] The Underwater Image Augmentation Dataset (UIEB) contains multiple real underwater images, some of which have corresponding labels, while others are severely degraded underwater images without labels. A portion of the degraded underwater images with corresponding labels is selected as the training set to train the underwater image augmentation network. The remaining portion is used as the test set to verify the augmentation effect under objective evaluation metrics. The remaining unlabeled and challenging degraded underwater images are used as test sets to verify the augmentation effect under subjective evaluation metrics.
[0023] For example, UIEB contains 950 real underwater images, of which 890 are labeled and the remaining 60 are severely degraded and unlabeled underwater images. 800 labeled degraded underwater images are selected as the training set to train the underwater image augmentation network. The remaining 90 labeled degraded images from the UIEB dataset are used as the test set to verify the augmentation effect under objective evaluation metrics, while the remaining 60 unlabeled and challenging degraded underwater images are used as the test set to verify the augmentation effect under subjective evaluation metrics. U45 contains 45 underwater degraded images from different scenarios, covering various underwater degradation scenarios such as high color cast, blur, and fogging. These are also used as the test set to verify the performance of the underwater augmentation network.
[0024] S2. Construct an underwater image enhancement network, including the Underwater Image Enhancement Compact Prior Representation Extraction Network (UIE-CPREN), the Multi-Color Space Feature Extraction Network (MSFEN), the Multi-Color Space Feature Fusion Module (MSFFM), and a diffusion model;
[0025] MSFEN extracts high-level representations from the RGB and LAB color spaces to obtain richer feature information. A ResWG module is also proposed to reduce the network's parameters and computational cost while maintaining MSFEN's feature extraction capabilities and preserving image details. The MSFFM module achieves complementary features across different color spaces through differential fusion of extracted features, improving the enhancement effect for subsequent underwater image tasks.
[0026] S21. Constructing a Compact Prior Representation Extraction Network for Underwater Image Enhancement (UIE-CPREN): Unlike traditional image generation tasks, underwater image enhancement tasks have certain prior knowledge. They only need to add detailed information (such as color cast removal and deblurring) to the original degraded underwater image to obtain a clear enhanced image. Therefore, it is not necessary to generate enhanced images from random noise through numerous iterations. UIE-CPREN is designed to extract underwater image enhancement prior representations from underwater images and compress them into a low-dimensional space to ensure a reduction in iteration steps during the diffusion model generation process.
[0027] First, the underwater degraded image With label image The channels are concatenated, and then downsampled by a factor of four using a pixel unshuffle operation to obtain the downsampled features:
[0028]
[0029] Next, convolutional layers with residual connections are used to extract underwater image enhancement prior representations, and global average pooling is used to compress these prior representations into a low-dimensional space. Finally, fully connected layers are used to re-integrate the information, resulting in a compact prior representation.
[0030]
[0031] Here, Linear represents a fully connected layer, and ResConv represents a convolutional layer with residual connections.
[0032] The final compact prior representation This is a one-dimensional feature vector. Compared to generating images, the diffusion model is more likely to accurately generate compact prior representations, which eliminates a large number of iterative steps in the diffusion model and reduces the computational complexity and inference time of underwater image enhancement algorithms.
[0033] S22, a Multi-Color Space Feature Extraction Network (MSFEN) is constructed. MSFEN contains two parallel feature extraction branches that independently extract features from the two color spaces. Within these feature extraction branches, a Residual Weighted Module (ReSWG) is proposed. For feature maps... ResWG first uses the Ghost module to extract features. GhostNet pointed out that the feature maps extracted by traditional convolution operations contain a large number of redundant features. Therefore, it simply uses a small number of convolution kernels to obtain useful feature information, while redundant features are obtained through inexpensive operations. Inspired by GhostNet, ResWG uses... Convolutional kernels are used to extract useful features, and depthwise separable convolution (DWConv) is used as a cheap operation in the Ghost module, which can be represented as:
[0034]
[0035] in, Furthermore, an SE (squeeze-and-excitation) module is embedded after the Ghost module, enabling the network to focus on channels with richer features.
[0036] In the SE module, first, the output of the ghost module is processed. Global average pooling is performed along the feature channel dimension, and a fully connected layer is used to predict the weight coefficients of each feature channel. This process can be represented as:
[0037]
[0038] Where H and W represent the width and height of the feature map, represent The features of the k-th channel This represents the weight coefficient for the k-th channel. Each feature extraction branch contains three ResWG modules. No downsampling operation is used during the feature extraction process to avoid loss of image information.
[0039] In the RGB color space feature extraction branch, the compact prior representation extracted by UIE-CPREN is used to guide the extraction of features beneficial to image enhancement. Guided by ResWG and prior information, MSFEN extracts enhancement features in the RGB color space. Enhanced features with Lab color space .
[0040] S23, construct a multi-color space feature fusion module (MSFFM) to enable complementary features from different color spaces. First, analyze the RGB color space features... Features of Lab color space Differential operations are performed to capture the differences between features, and the differential features are input into a gating selection mechanism to select effective complementary features from different color spaces.
[0041] Specifically, global average pooling and max pooling are used at the channel dimension to retain more global feature information, and the global features are input into the fully connected layer to predict feature selection weights. By controlling these weights, complementary cross-color space features are preserved, and these features are concatenated along the channel dimension to achieve effective feature fusion, which can be represented as:
[0042]
[0043] in, This represents the Sigmoid activation function, which is ultimately passed through... Convolutional layer traversal and feature fusion The final enhanced underwater image was obtained. .
[0044] S3. After constructing the underwater image augmentation network, the model is trained in two stages. In the pre-training stage, UIE-CPREN, MSFEN, and MSFFM are trained simultaneously. This enables UIE-CPREN to generate accurate and compact prior representations, and allows MSFEN and MSFFM to fully utilize these compact prior representations for underwater image augmentation. During training, mean absolute error loss (L1 Loss), structural similarity loss (SSIM loss), and perceptual similarity loss are employed. The mean absolute error loss measures the difference between the augmented image and the label, while the structural similarity loss measures the global structural similarity between the augmented image and the label.
[0045]
[0046] in, Indicates an enhanced image. Indicates a label. Represents the mean of the image. Represents the standard deviation of the image. and This is a hyperparameter.
[0047] During the joint training phase, the powerful data generation capabilities of the diffusion model are first leveraged to directly generate compact prior representations from degraded images. The UIE-CPREN pre-trained in the first stage is first weighted and then used to generate accurate prior representation z. The diffusion process can be viewed as a Markov chain, gradually breaking z into a noisy prior representation by continuously adding uncertainty noise. This process can be represented as:
[0048]
[0049] Then, the reverse denoising process is used to... Denoising is performed, progressively generating a compact prior representation identical to z. During denoising, UIE-CPRENv2 is designed to extract only the degraded representation f from the underwater degraded image as a condition for inverse denoising. Its structure is consistent with UIE-CPREN, except that the input dimension of the first convolutional layer is reduced by half. The inverse denoising process can be represented as:
[0050]
[0051] After obtaining the prior representation generated by the diffusion model, consistent with the previous stage, this representation is used to guide subsequent MSFEN and MSFFM underwater image enhancement, resulting in the final enhanced image. We jointly train the diffusion model with the subsequent MSFEN and MSFFM to more fully utilize the guiding role of the compact prior representation. The loss function of the joint training process can be expressed as:
[0052]
[0053] in D represents the number of feature channels in the compact prior representation.
[0054] S4, during the model inference phase, randomly samples a low-dimensional compact vector from the standard normal distribution and generates an accurate compact prior representation through a diffusion model to guide the subsequent underwater image enhancement network.
[0055] During the model inference phase, inverse denoising conditions are first extracted from the underwater degradation model using the pre-trained UIE-CPRENv2. Then, a noise sample is randomly generated. A compact prior representation is obtained by denoising random noise through a reverse denoising process. This compact prior representation is then used to guide MSFEN and MSFFM in enhancing underwater degraded images, resulting in enhanced underwater images. The enhancement effect of this algorithm on the underwater datasets UIEB and U45 is shown in the attached figure. Figure 3 , Figure 4 As shown.
[0056] In summary, this invention proposes an underwater image enhancement algorithm guided by prior information, which consists of a Compact Prior Representation Extraction Network (UIE-CPREN), a Multi-Color Space Feature Extraction Network (MSFEN), a Multi-Color Space Feature Fusion Module (MSFFM), and a diffusion model. MSFEN extracts high-level representations from the RGB and LAB color spaces to obtain richer feature information. A ResWG module is also proposed to reduce the network parameters and computational cost while maintaining the feature extraction capability of MSFEN and preserving image details. The MSFFM module achieves complementarity of different color space features through differential fusion of the extracted multi-color space features, improving the enhancement effect of subsequent UIE tasks.
[0057] This invention trains the proposed underwater image enhancement network in two stages: pre-training and joint training.
[0058] In the pre-training phase, UIE-CPREN is first used to extract low-dimensional compact prior representations from underwater degraded images and their corresponding labels. These representations are then used to guide MSFEN and MSFFM for underwater image enhancement. MSFEN extracts richer feature representations from RGB, LAB, and other color spaces, while MSFFM performs differential fusion of features from multiple color spaces to achieve complementarity and improve the enhancement effect of subsequent UIE tasks. In the joint training phase, conditional vectors for the denoising process are extracted from the underwater degraded images to guide the diffusion model in generating the same compact representations as in the pre-training phase, which then guide subsequent modules for underwater image enhancement. During joint training, the parameters of the diffusion model and subsequent enhancement networks are jointly optimized to make fuller use of the guidance provided by the compact prior representations. In the model inference phase, a low-dimensional compact vector is randomly sampled from a standard normal distribution, and the diffusion model generates accurate compact prior representations to guide the subsequent underwater image enhancement network. Since the compact prior representation is only a low-dimensional compact vector, the diffusion model only requires a few denoising steps to generate accurate prior representations, significantly improving the model's inference speed.
[0059] It is understood that the present invention has been described through some embodiments, and those skilled in the art will recognize that various changes or equivalent substitutions can be made to these features and embodiments without departing from the spirit and scope of the invention. Furthermore, under the teachings of the present invention, these features and embodiments can be modified to adapt to specific situations and materials without departing from the spirit and scope of the invention. Therefore, the present invention is not limited to the specific embodiments disclosed herein, and all embodiments falling within the scope of the claims of this application are within the protection scope of the present invention.
Claims
1. An underwater image enhancement method guided by prior information, characterized in that: Includes the following steps: S1. Prepare the underwater image enhancement dataset. Divide the publicly available underwater image enhancement dataset UIEB into a training set and a test set to provide a data foundation for the subsequent training of underwater image enhancement models. S2. Construct an underwater image enhancement network, including an underwater image enhancement compact prior representation extraction network, a multi-color space feature extraction network, a multi-color space feature fusion module, and a diffusion model; S3. Perform model training, and divide the model training process into two stages: pre-training and joint training; S4. During the model inference phase, a low-dimensional compact vector is randomly sampled from the standard normal distribution, and an accurate compact prior representation is generated through the diffusion model to guide the subsequent underwater image enhancement network.
2. The underwater image enhancement method based on prior information guidance according to claim 1, characterized in that: The underwater image augmentation dataset UIEB in S1 contains multiple real underwater images, some of which have corresponding labels, and others are severely degraded underwater images without labels. A portion of the degraded underwater images with corresponding labels is selected as the training set to train the underwater image augmentation network. The remaining portion is used as the test set to verify the augmentation effect under objective evaluation metrics. The remaining unlabeled and challenging degraded underwater images are used as test sets to verify the augmentation effect under subjective evaluation metrics.
3. The underwater image enhancement method based on prior information guidance according to claim 1, characterized in that: In step S2, constructing the compact prior representation extraction network for underwater image enhancement includes the following steps: Step 1: Convert the underwater degraded image With label image The channels are concatenated, and then downsampled by a factor of four using a channel shuffling operation to obtain the downsampled features: Step 2: Extract underwater image enhancement prior representations using convolutional layers with residual connections, compress these prior representations into a low-dimensional space using global average pooling, and then re-integrate the information using fully connected layers to obtain compact prior representations. Wherein, Linear represents a fully connected layer, and ResConv represents a convolutional layer with residual connections; The final compact prior representation It is a one-dimensional feature vector.
4. The underwater image enhancement method based on prior information guidance according to claim 2, characterized in that: In step S2, constructing a multi-color space feature extraction network includes the following process: it comprises two parallel feature extraction branches that independently extract features from the two color spaces respectively. Within each feature extraction branch, a residual weighting module is proposed. For feature maps The residual weighted module first uses the Ghost module to extract features. use Convolutional kernels extract useful features, and depthwise separable convolution is used as a cheap operation in the Ghost module, which is represented as follows: in, , The SE module is embedded after the Ghost module.
5. The underwater image enhancement method based on prior information guidance according to claim 4, characterized in that: In the SE module, the output of the ghost module is first processed. Global average pooling is performed along the feature channel dimension, and a fully connected layer is used to predict the weight coefficients of each feature channel. This process is represented as: Where H and W represent the width and height of the feature map, represent The features of the k-th channel The weight coefficients for the k-th channel are represented by three ResWG modules in each feature extraction branch.
6. The underwater image enhancement method based on prior information guidance according to claim 5, characterized in that: A multi-color space feature fusion module is used to complement the features of different color spaces, specifically the RGB color space features. Features of Lab color space Differential operations are performed to capture the differences between features, and the differential features are input into a gating selection mechanism to select effective complementary features from different color spaces.
7. The underwater image enhancement method based on prior information guidance according to claim 1, characterized in that: During the model training phase of S3, During the pre-training phase, the underwater image enhancement compact prior representation extraction network is used to extract low-dimensional compact prior representations from the corresponding labels, and these representations are used to guide the multi-color space feature extraction network and the multi-color space fusion module to perform underwater image enhancement. During the joint training phase, the same compact representation as in the pre-training phase is used to estimate the diffusion model to guide subsequent modules for underwater image enhancement, and the diffusion model is jointly optimized with the subsequent enhancement network.
8. The underwater image enhancement method based on prior information guidance according to claim 7, characterized in that: During training, mean absolute error loss, structural similarity loss, and semantic similarity loss are used. Mean absolute error is used to measure the difference between the augmented image and the label, and structural similarity loss is used to measure the global structural similarity between the augmented image and the label. in, Indicates an enhanced image. Indicates a label, Represents the mean of the image. Represents the standard deviation of the image. and This is a hyperparameter.
9. The underwater image enhancement method based on prior information guidance according to claim 8, characterized in that: During the model training phase of S3, In the joint training phase, the weights of the underwater image enhancement compact prior representation extraction network pre-trained in the first stage are first frozen, and an accurate prior representation z is generated using it. Then, the powerful data generation capability of the diffusion model is used to directly generate a similar compact prior representation from the degraded image. After obtaining the prior representation generated by the diffusion model, consistent with the previous stage, this representation is used to guide the subsequent multi-color space feature extraction network and multi-color space feature fusion module to perform underwater image enhancement, resulting in the final enhanced image.
10. The underwater image enhancement method based on prior information guidance according to claim 1, characterized in that: In the inference phase, S4 first extracts inverse denoising conditions from the underwater degradation model using a trained underwater image-enhanced compact prior representation extraction network, and then randomly samples a noise sample. By performing a reverse denoising process to denoise random noise, a compact prior representation is obtained. Finally, this compact prior representation is used to guide the multi-color space feature extraction network and the multi-color space feature fusion module to enhance the underwater degraded image, resulting in an enhanced underwater image.