An underwater image processing method based on a lightweight diffusion model
By using a lightweight diffusion model and neural architecture search algorithm, combined with color distribution priors, a network structure adapted to different terminal devices is automatically designed. This solves the problems of traditional methods in underwater image processing neglecting physical priors and insufficient diversity of generative models, and achieves efficient and lightweight image enhancement effects.
Patent Information
- Application Number
- CN202310914540.X
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2023-07-25
- Publication Date
- 2026-01-16
- Estimated Expiration
- 2043-07-25
AI Technical Summary
Existing underwater image processing methods neglect traditional physical priors, and generative adversarial network-based generative models lack diversity, making it difficult to design optimal network structures that adapt to different edge devices. Furthermore, deep learning models have a large number of parameters, making them difficult to apply to resource-constrained devices.
A lightweight diffusion model is adopted, combined with a neural architecture search algorithm, to automatically search for neural network modules. By utilizing a latent diffusion model and a compressed sensing model, and using color distribution priors to guide the training process, a lightweight network structure adapted to different terminal devices is designed.
It achieves efficient image enhancement on different resource terminal devices, generates high-quality images, improves the model's generalization ability and robustness, and reduces the number of model parameters.
Smart Images

Figure CN117152600B_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The present application relates to the field of image processing, in particular to an underwater image processing method based on a light diffusion model. BACKGROUND
[0002] In recent years, marine scenes have attracted extensive attention in the field of computer vision due to their rich and under-explored resources. Because the image quality captured in underwater scenes is degraded, including but not limited to loss of detail, color distortion, low brightness and low contrast, blurring, etc., underwater image processing technology has become a research hotspot. This technology mainly has the following application scenarios: underwater image enhancement, underwater image generation, underwater image color correction, and underwater image high-resolution reconstruction. The present application mainly aims at the first scenario, in which the propagation of light in water is affected by the absorption and scattering effects of particles in water, resulting in image quality degradation. Traditional methods or deep learning methods are used to improve image quality for subsequent tasks.
[0003] Current underwater image processing algorithms can be roughly divided into traditional methods and deep learning-based methods. Traditional methods are based on the imaging mechanism of underwater images and consider image enhancement as an inverse problem. By estimating the parameters of the underwater image formation model, the image quality is improved. However, this method relies on physical priors and usually requires complex mathematical formulas. At the same time, due to the limitations of the formula, this method usually has limited application scenarios and poor generalization ability. Deep learning-based methods have made great progress in recent years. Peng et al. used Transformer to complete the underwater image enhancement task in “U-shape Transformer for Underwater Image Enhancement” and achieved good enhancement results, but the model structure is complex and has a large number of parameters. Deep learning methods use convolutional neural networks or Transformers to capture image features to improve image quality and complete enhancement tasks. However, this method usually relies on researchers' prior knowledge to design complex network structures, and models with excessive parameters are difficult to apply to resource-constrained devices. Another method based on generative adversarial networks synthesizes high-quality images to fill the gap in the current real underwater data. Although this method generates high-quality images, it has poor diversity, making it difficult for models trained on such synthetic datasets to be extended to real-world scenarios. In addition, this deep learning method ignores the importance of physical priors in traditional methods.
[0004] In 2020, Ho et al. proposed the Denoising Diffusion Probabilistic Model (DDPM) in the paper "Denoising Diffusion Probabilistic Model", which is a probabilistic likelihood-based generative model with strong distribution learning ability. Since the publication of DDPM, the diffusion model field has seen a surge of work exploring and optimizing the performance of diffusion models. In 2022, Rombach et al. proposed the Latent Diffusion Model (LDM) in the paper "High-Resolution Image Synthesis with Latent Diffusion Models", which compresses images from pixel space to latent space and performs diffusion and denoising processes in this space, reducing training costs and speeding up model inference. In 2023, Lu et al. applied DDPM to the underwater image enhancement task in the paper "Underwater Image Enhancement Method Based on Denoising Diffusion Probabilistic Model" and achieved certain results, but the slow sampling process of DDPM limits its potential applications. In addition, many advances in the field of diffusion models have not been applied to underwater image processing tasks.
[0005] In summary, the current problems of underwater image processing include:
[0006] (1) Deep learning-based methods ignore the importance of traditional physical priors, which are essential for describing high-quality image distributions.
[0007] (2) The diversity of underwater image datasets generated by generative adversarial network-based generative models is insufficient, resulting in models trained on such datasets lacking strong robustness and difficulty in extending to real-world scenarios.
[0008] (3) Designing deep learning network models requires researchers' prior knowledge, making it difficult to design the optimal network structure, especially for different edge devices. SUMMARY
[0009] The purpose of the present application is to propose a lightweight diffusion model-based underwater image processing method to solve the above technical problems.
[0010] To achieve the above purpose, the present application provides the following scheme: a lightweight diffusion model-based underwater image processing method, comprising the following steps:
[0011] Step (1), obtaining the color distribution histogram of degraded images and enhanced images in the existing underwater image dataset;
[0012] Taking an underwater image dataset as the research object, the underwater image dataset includes a paired degraded image U d and an enhanced image U e ; the degraded image represents an image directly obtained in an underwater scene, and the enhanced image represents an image obtained after using an image enhancement algorithm on the degraded image; a color histogram of each pair of image combinations in the underwater image dataset is obtained, and the color histogram presents the distribution y of red, green and blue colors in the [0, 255] pixel value range as a color distribution prior;
[0013] Step (2), designing a neural architecture search algorithm for automatically searching neural network modules, including the following steps:
[0014] Step (2.1), defining a search space: the search space is a set of candidate neural network modules;
[0015] The search space S i ={conv(depth,width,kernelsize),attention}, where conv represents a convolution type, attention represents an attention mechanism, depth represents a network depth, width represents a network width, and kernelsize represents a convolution kernel size;
[0016] Step (2.2), training a neural network composed of the search space: the weight of the neural network composed of the search space is W o , the selection mechanism of the neural network module composed of the search space is C, and the loss of the neural network composed of the search space is L val , and the training process takes as the training target;
[0017] Step (2.3), after training the neural network composed of the search space, a plurality of combined neural network modules are obtained;
[0018] Step (3), defining a latent diffusion model;
[0019] The latent diffusion model includes a compressed sensing model, a diffusion model and a conditional encoder; the compressed sensing model includes an encoder ε and a decoder The diffusion model is a U-Net-based neural network, including an encoder and a decoder The conditional encoder only includes an encoder τ θ ;
[0020] The data flow of the latent diffusion model includes a compressed sensing model data flow and a latent diffusion model data flow;
[0021] The data flow of the compressed sensing model is that the encoder of the compressed sensing model compresses the degraded image of the underwater image dataset into a latent variable, and the decoder of the compressed sensing model reconstructs the latent variable into a restored image.
[0022] The data flow of the latent diffusion model is that the encoder of the compressed sensing model compresses the degraded image of the underwater image dataset into a latent variable, the latent variable is input into the encoder of the diffusion model; the encoder of the diffusion model adds noise to the latent variable to generate a noisy latent variable; the noisy latent variable is input into the decoder of the diffusion model; the conditional encoder compresses the color prior into a latent space prior, and the latent space prior is input into the decoder of the diffusion model; the decoder of the diffusion model predicts noise, outputs a denoised latent variable, and inputs the denoised latent variable into the decoder of the compressed sensing model; the decoder of the compressed sensing model reconstructs the denoised latent variable into an enhanced image.
[0023] Step (4), training the latent diffusion model;
[0024] The training of the latent diffusion model includes two processes; first, only the compressed sensing model is trained, and in this process, the diffusion model and the conditional encoder do not participate in the training; second, the diffusion model is trained, and in this process, the trained compressed sensing model is only used to compress the degraded image of the underwater image dataset into a latent variable and to reconstruct the denoised latent variable into an enhanced image, and the model parameters are not updated; the conditional encoder uses a pre-trained open source model and only compresses the color prior into a latent space prior, and the model parameters are not updated;
[0025] Step (5), according to the plurality of combined neural network modules obtained in step (2.3), repeating step (3.2) to obtain a diffusion model based on a plurality of combined U-Net; repeating step (4.2) to carry out an image enhancement task, training and evaluating different diffusion models, and selecting a model for different terminal devices according to the performance.
[0026] The step (3) specifically includes the following steps:
[0027] Step (3.1), defining a compressed sensing model:
[0028] The compressed sensing model includes an encoder ε and a decoder The encoder is used to compress the degraded image x of the underwater image dataset into a latent space variable z; the decoder is used to restore the latent space variable z into a reconstructed image Specifically,
[0029] Given the degraded image x of the underwater image dataset in the RGB space where H represents the height of the degraded image, W represents the width of the degraded image; the encoder ε of the compressed sensing model compresses the degraded image x into a latent space variable z according to the coefficient f wherein, h represents the height of the latent space variable w represents the width of the latent space variable decoder of the compressive sensing model converts the latent space variable into a reconstructed image i.e.
[0030] Step (3.2), defining the diffusion model:
[0031] According to the neural network module obtained in step (2.3) as the skeleton of the diffusion model, i.e. the U-Net network structure; U-Net∈ θ is a symmetric structure of an encoder-decoder neural network, the encoder is composed of multiple convolutional layers, and the decoder is composed of multiple transposed convolutional layers, and the convolutional layers and the transposed convolutional layers use skip connections and attention mechanisms between them;
[0032] Step (3.3), defining the conditional encoder:
[0033] The conditional encoder τ θ is composed of multiple convolutional layers, and is used to project the color distribution prior y to the latent space prior τ θ (y);
[0034] Step (3.4), defining the loss function:
[0035] The loss function of the compressive sensing model is:
[0036]
[0037] wherein, L per represents the perceptual loss, L reg represents the regularization loss, W per represents the perceptual loss weight, W reg represents the regularization loss weight.
[0038] The loss function of the diffusion model is:
[0039]
[0040] wherein, ∈ represents noise subject to a Gaussian distribution, ∈ θ represents noise predicted by the U-Net network model, and t represents a time in a time sequence.
[0041] The step (4) specifically comprises the following steps:
[0042] Step (4.1), training the compressive sensing model:
[0043] Set the number of training epochs epoch, repeat the following steps:
[0044] Step (4.1.1), get the degraded image x0 from the underwater image dataset, input to the encoder ε of the compressed sensing model, get the latent variable
[0045] Step (4.1.2), latent variable Input to the decoder of the compressed sensing model Get the reconstructed image
[0046] Step (4.1.3), according to the loss function L CPM Calculate the loss of this round; update the compressed sensing model parameters using the gradient descent algorithm;
[0047] Step (4.2), train the diffusion model:
[0048] Set the number of training epochs epoch; initialize the noise mechanism α t
[0049] Repeat the following steps at each time t in the time sequence Uniform({1,…,T}):
[0050] Step (4.2.1), forward process:
[0051] Step (4.2.1.1), get the degraded image x0 from the underwater dataset, input to the encoder ε of the compressed sensing model, get the latent variable
[0052] Step (4.2.1.2), according to the noise mechanism α t Get the noise level ∈;
[0053] Step (4.2.1.3), add the noise level ∈ to the latent variable Get the noise latent variable
[0054] After the forward process, get the complete noise latent variable z0 that obeys the Gaussian distribution
[0055] Step (4.2.2), backward process:
[0056] Step (4.2.2.1), get the complete noise latent variable z0
[0057] Step (4.2.2.2), use the conditional encoder τ θ Project the color distribution prior y to the latent space to get the latent space prior τ θ (y);
[0058] Step (4.2.2.3), latent space prior τ θ (y) and noise latent variable Encoder of diffusion model Decoder of diffusion model Output corresponding noise ∈ θ ;
[0059] Step (4.2.2.4), using loss function L LDM Calculate loss;
[0060] Step (4.2.2.5), according to the formula Get latent variable Where σ t And z both represent constants,
[0061] Step (4.2.2.6), latent variable Decoder of compressed sensing model Get enhanced image
[0062] Step (4.2.2.7), update diffusion model parameters using gradient descent;
[0063] The training of the latent diffusion model includes two processes; first, only train the compressed sensing model, in this process, the diffusion model and the conditional encoder do not participate in the training; second, train the diffusion model, in this process, the trained compressed sensing model is only used to compress the degraded image of the underwater image dataset into the latent variable and to reconstruct the denoised latent variable into the enhanced image, and the model parameters are not updated; the conditional encoder uses a pre-trained open source model, which is only used to compress the color prior into the latent space prior, and the model parameters are not updated;
[0064] Specifically includes the following steps:
[0065] Step (4.1), train the compressed sensing model:
[0066] Set the training round epoch, repeat the following steps:
[0067] Step (4.1.1), get degraded image x0 from underwater image dataset, input to encoder ε of compressed sensing model, get latent variable
[0068] Step (4.1.2), latent variable Input to decoder of compressed sensing model Get reconstructed image
[0069] Step (4.1.3), update the compressed sensing model parameters using gradient descent algorithm according to the loss function L CP Calculate the loss of this round; update the compressed sensing model parameters using gradient descent algorithm;
[0070] Step (4.2), train the diffusion model:
[0071] Set the number of training rounds epoch; initialize the noise mechanism a t
[0072] Repeat the following steps at each time t in the time series Uniform({1,…,T}):
[0073] Step (4.2.1), forward process:
[0074] Step (4.2.1.1), get the degraded image x0 from the underwater data set, input it into the encoder ε of the compressed sensing model, and get the latent variable
[0075] Step (4.2.1.2), according to the noise mechanism a t Get the noise level ∈;
[0076] Step (4.2.1.3), add the noise level ∈ to the latent variable Get the noise latent variable
[0077] After the forward process, get the complete noise latent variable
[0078] Step (4.2.2), backward process:
[0079] Step (4.2.2.1), get the complete noise latent variable
[0080] Step (4.2.2.2), use the conditional encoder τ θ Project the color distribution prior y to the latent space to get the latent space prior τ θ (y);
[0081] Step (4.2.2.3), input the latent space prior τ θ (y) and the noise latent variable to the encoder of the diffusion model The decoder of the diffusion model Output the corresponding noise ∈ θ ;
[0082] Step (4.2.2.4), use the loss function L LDM Calculate the loss;
[0083] Step (4.2.2.5), obtaining latent variable z according to formula obtaining latent variable wherein, σ t and z both represent constants,
[0084] Step (4.2.2.6), updating diffusion model parameters using gradient descent input into the decoder of the compressed sensing model obtaining enhanced image
[0085] Step (4.2.2.7), using gradient descent to update diffusion model parameters;
[0086] Step (5), according to the plurality of combined neural network modules obtained in step (2.3), repeating step (3.2) to obtain a diffusion model based on a plurality of combined U-Net; repeating step (4.2) to carry out an image enhancement task, training and evaluating different diffusion models, and selecting a model for different terminal devices according to their performance.
[0087] The present application has the beneficial effects that the present application proposes an underwater image processing method based on a lightweight diffusion model, guides the training process of the diffusion model according to the data set distribution prior, uses a neural architecture search algorithm to find a suitable network module, and uses the diffusion model to complete an image enhancement task. The present application uses the outstanding distribution learning ability of the diffusion model to capture the data distribution, uses a high-quality reference prior to generate a high-quality image, and can find a suitable network structure for different resource terminal devices to achieve lightweight design. BRIEF DESCRIPTION OF DRAWINGS
[0088] Figure 1 is a flowchart of an underwater processing algorithm based on a lightweight diffusion model proposed by the present application.
[0089] Figure 2 is a flowchart of a neural architecture search algorithm in the present application.
[0090] Figure 3 is a model composition diagram of a latent diffusion model in the present application.
[0091] Figure 4 is a training flowchart of a latent diffusion model in the present application.
[0092] In the figure, is an image, is an autoencoder, is a neural architecture search module, is a prior, is a latent variable, and is a skip connection. DETAILED DESCRIPTION
[0093] In order to make the purpose, technical scheme and advantages of the present application more clear, the present application is further described in detail below with examples. It should be understood that the specific examples described herein are only used to explain the present application and not to limit the present application.
[0094] The underwater image processing method based on the lightweight diffusion model is used for training a diffusion model in the field of underwater image processing to complete an underwater image enhancement task.
[0095] The flow of the present application is shown in Figure 1 First, the existing underwater image dataset or underwater images taken by underwater robots, underwater cameras and other equipment are taken as research targets, and the prior distribution of image colors in the underwater dataset is analyzed; second, a neural architecture search algorithm is used to automatically search for network modules of the diffusion model, including defining a search space and training a neural network, and the flow is shown in Figure 2 Third, the latent diffusion model is trained, including compressing images to a latent space, performing diffusion and denoising processes in the latent space, and generating images according to latent variables. The latent diffusion model is shown in Figure 3 The training process is shown in Figure 4 Finally, different neural network modules are obtained according to the neural architecture search algorithm, diffusion models based on the network modules are trained, the performance of the diffusion models is evaluated, and the models for different terminal devices are selected according to the performance.
[0096] The specific steps are as follows:
[0097] The underwater image processing method based on the lightweight diffusion model includes the following steps:
[0098] Step (1), obtaining the color distribution histogram of the degraded image and the enhanced image in the existing underwater image dataset;
[0099] The underwater image dataset is taken as the research target, and the underwater image dataset includes paired image combinations of degraded images U d and enhanced images U e The degraded image represents an image directly obtained in an underwater scene, and the enhanced image represents an image obtained by using an image enhancement algorithm on the degraded image; the color histogram of each pair of image combinations in the underwater image dataset is obtained, and the color histogram presents the distribution y of red, green and blue colors in the [0, 255] pixel value range as a color distribution prior;
[0100] Step (2), designing a neural architecture search algorithm for automatically searching for neural network modules, including the following steps:
[0101] Step (2.1), defining a search space: the search space is a set of candidate neural network modules;
[0102] search space S i = {conv(depth, width, kernelsize), attention}, where conv represents the convolution type, attention represents the attention mechanism, depth represents the network depth, width represents the network width, and kernelsize represents the kernel size;
[0103] Step (2.2), training the neural network composed of the search space: the weight of the neural network composed of the search space is W o , the search space composed of the neural network module selection mechanism is C, and the search space composed of the neural network loss is L val , and the training process is taken as the training target;
[0104] Step (2.3), after training the neural network composed of the search space, a plurality of combined neural network modules are obtained;
[0105] Step (3), defining a latent diffusion model;
[0106] The latent diffusion model includes a compressed sensing model, a diffusion model, and a conditional encoder; the compressed sensing model includes an encoder ε and a decoder The diffusion model is a U-Net-based neural network, including an encoder and a decoder The conditional encoder only includes an encoder τ θ ;
[0107] The data flow of the latent diffusion model includes a compressed sensing model data flow and a latent diffusion model data flow;
[0108] The compressed sensing model data flow is: the encoder of the compressed sensing model compresses the degraded image of the underwater image dataset into a latent variable, and the decoder of the compressed sensing model reconstructs the latent variable into a restored image;
[0109] The data flow of the latent diffusion model is: the encoder of the compressed sensing model compresses the degraded image of the underwater image dataset into a latent variable, and the latent variable is input into the encoder of the diffusion model; the encoder of the diffusion model adds noise to the latent variable to generate a noisy latent variable; the noisy latent variable is input into the decoder of the diffusion model; the conditional encoder compresses the color prior into a latent space prior, and the latent space prior is input into the decoder of the diffusion model; the decoder of the diffusion model predicts noise, outputs a denoised latent variable, and inputs the denoised latent variable into the decoder of the compressed sensing model; the decoder of the compressed sensing model reconstructs the denoised latent variable into an enhanced image;
[0110] Specifically, the following steps are included:
[0111] Step (3.1), define the compressed sensing model:
[0112] The compressed sensing model comprises an encoder ε and a decoder The encoder is configured to compress a degraded image x of an underwater image dataset into a latent space variable z; the decoder is configured to restore the latent space variable z into a reconstructed image Specifically, the degraded image x of the underwater image dataset in the RGB space is defined as:
[0113] The degraded image x of the underwater image dataset in the RGB space is defined as: wherein H represents the height of the degraded image, and W represents the width of the degraded image; the encoder ε of the compressed sensing model compresses the degraded image x into a latent space variable z according to the coefficient f wherein h represents the height of the latent space variable z, and w represents the width of the latent space variable z; the decoder of the compressed sensing model restores the latent space variable z into a reconstructed image
[0114] Step (3.2), define the diffusion model:
[0115] The neural network module obtained in step (2.3) is used as the skeleton of the diffusion model, i.e., the U-Net network structure; the U-Net ∈ θ is a symmetric encoder-decoder neural network, the encoder is composed of multiple convolutional layers, and the decoder is composed of multiple transposed convolutional layers, and the convolutional layers and the transposed convolutional layers use skip connections and attention mechanisms therebetween;
[0116] Step (3.3), define the conditional encoder:
[0117] The conditional encoder τ θ is composed of multiple convolutional layers, and is configured to project the color distribution prior y into the latent space prior τ θ (y);
[0118] Step (3.4), define the loss function:
[0119] The loss function of the compressed sensing model is:
[0120]
[0121] wherein L per represents the perceptual loss, L reg represents the regularization loss, W per represents the perceptual loss weight, and W reg denotes the regular loss weight;
[0122] The loss function of the diffusion model is:
[0123]
[0124] where ∈ denotes the noise obeying Gaussian distribution, ∈ θ denotes the noise predicted by the U-Net network model, and t denotes a time point in the time series;
[0125] Step (4), training the latent diffusion model;
[0126] The training of the latent diffusion model includes two processes; first, only the compressed sensing model is trained, and in this process, the diffusion model and the conditional encoder do not participate in the training; second, the diffusion model is trained, and in this process, the trained compressed sensing model is only used to compress the degraded image of the underwater image dataset into latent variables and reconstruct the denoised latent variables into enhanced images, without updating the model parameters; the conditional encoder adopts a pre-trained open source model, which is only used to compress the color prior into the latent space prior, without updating the model parameters;
[0127] Specifically, the following steps are included:
[0128] Step (4.1), training the compressed sensing model:
[0129] Set the training round epoch, and repeat the following steps:
[0130] Step (4.1.1), obtain the degraded image x0 from the underwater image dataset, input it into the encoder ε of the compressed sensing model, and obtain the latent variable
[0131] Step (4.1.2), the latent variable is input into the decoder of the compressed sensing model to obtain the reconstructed image
[0132] Step (4.1.3), according to the loss function L CPM , calculate the loss of this round; update the parameters of the compressed sensing model using the gradient descent algorithm;
[0133] Step (4.2), training the diffusion model:
[0134] Set the training round epoch; initialize the noise mechanism α t
[0135] Repeat the following steps at each time point t in the time series Uniform({1,…,T}):
[0136] Step (4.2.1), forward process:
[0137] Step (4.2.1.1), get the degraded image x0 from the underwater dataset, input to the encoder ε of the compressive sensing model, get the latent variable
[0138] Step (4.2.1.2), according to the noise mechanism α t Get the noise level ∈;
[0139] Step (4.2.1.3), add the noise level ∈ to the latent variable Get the noise latent variable
[0140] After the forward process, get the complete noise latent variable z0 that obeys the Gaussian distribution
[0141] Step (4.2.2), backward process:
[0142] Step (4.2.2.1), get the complete noise latent variable z0
[0143] Step (4.2.2.2), use the conditional encoder τ θ Project the color distribution prior y to the latent space to get the latent space prior τ θ (y);
[0144] Step (4.2.2.3), input the latent space prior τ θ (y) and the noise latent variable z to the encoder of the diffusion model The decoder of the diffusion model Output the corresponding noise ∈ θ ;
[0145] Step (4.2.2.4), use the loss function L LDM Calculate the loss;
[0146] Step (4.2.2.5), according to the formula Get the latent variable Where, σ t And z both represent constants,
[0147] Step (4.2.2.6), input the latent variable z to the decoder of the compressive sensing model Get the enhanced image
[0148] Step (4.2.2.7), update the diffusion model parameters using gradient descent;
[0149] Step (5), according to the plurality of combined neural network modules obtained in the step (2.3), repeating the step (3.2), obtaining the diffusion model based on the plurality of combined U-Net; repeating the step (4.2), carrying out the image enhancement task, training and evaluating different diffusion models, and selecting the model for different terminal devices according to the performance.
[0150] In summary:
[0151] The application proposes an underwater image processing method based on a lightweight diffusion model, guides the model training process according to high-quality prior, uses a neural architecture search algorithm to find network modules, uses a diffusion model to complete an image enhancement task, flexibly finds a suitable network structure for different resource terminal devices, and achieves lightweight design.
[0152] The basic principles, main features and advantages of the application are shown and described above. It should be understood by those skilled in the art that the application is not limited by the above examples, and the above examples and descriptions in the specification are only to illustrate the principles of the application. Without departing from the spirit and scope of the application, various changes and improvements can be made to the application, and these changes and improvements all fall within the scope of the claimed application. The scope of protection of the application is defined by the appended claims and their equivalents.
Claims
1. A method for underwater image processing based on a lightweight diffusion model, characterized in that, The method comprises the following steps: step (1), obtaining color distribution histograms of degraded images and enhanced images in an existing underwater image dataset; With a underwater image dataset as the research object, the underwater image dataset includes paired degraded images U d and enhanced image U e combination of images; the degraded image represents an image directly obtained in an underwater scene, and the enhanced image represents an image obtained after using an image enhancement algorithm on the degraded image; a color histogram of each pair of image combinations in the underwater image dataset is obtained, and the color histogram presents the distribution y of red, green and blue colors in the [0, 255] pixel value range as a color distribution prior. Step (2), designing a neural architecture search algorithm for automatically searching neural network modules, comprising the following steps: Step (2.1), defining a search space: the search space is a set of candidate neural network modules; Search space S i = {conv(depth, width, kernel size), attention}, where conv denotes a convolution type, attention denotes an attention mechanism, depth denotes a network depth, width denotes a network width, and kernel size denotes a kernel size. Step (2.2), training the neural network composed of the search space: the neural network composed of the search space weights are W o , the neural network composed of the search space module selection mechanism is C, and the neural network composed of the search space loss is L val , the training process is aimed at ; Step (2.3), after training the neural network composed of the search space, a plurality of combined neural network modules are obtained; Step (3), defining a latent diffusion model; The model structure of the latent diffusion model comprises a compression sensing model, a diffusion model and a conditional encoder; the compression sensing model comprises an encoder and a decoder The diffusion model is a neural network based on U-Net, comprising an encoder and a decoder The conditional encoder only comprises an encoder τ θ ; The data flow of the latent diffusion model comprises a compressed sensing model data flow and a latent diffusion model data flow; The compressed sensing model data stream is: an encoder of the compressed sensing model The degraded image of the compressed underwater image dataset is a latent variable, and a decoder of the compressed sensing model The reconstructed latent variable is a restored image; The latent diffusion model data flow is: the encoder of the compressive sensing model The degraded image of the compressed underwater image dataset is a latent variable, and the latent variable is input into the encoder of the diffusion model The encoder of the diffusion model Noise is added to the latent variable to generate a noise latent variable; the noise latent variable is input into the decoder of the diffusion model The conditional encoder compresses the color prior into a latent space prior, and the latent space prior is input into the decoder of the diffusion model The decoder of the diffusion model Predicting noise, outputting a denoised latent variable, and inputting the denoised latent variable into the decoder of the compressive sensing model The decoder of the compressive sensing model The denoised latent variable is reconstructed into an enhanced image; Step (4), training the latent diffusion model; The training of the latent diffusion model comprises two processes; first, only the compressed sensing model is trained, and in this process, the diffusion model and the conditional encoder are not involved in the training; second, the diffusion model is trained, and in this process, the trained compressed sensing model is only used to compress the degraded images of the underwater image dataset into latent variables and to reconstruct the denoised latent variables into enhanced images, and the model parameters are not updated; the conditional encoder adopts a pre-trained open source model, and is only used to compress the color prior into a latent space prior, and the model parameters are not updated; Step (5), according to the plurality of combined neural network modules obtained in step (2.3), the diffusion model is repeatedly defined to obtain diffusion models based on a plurality of combined U-Net; the diffusion model is repeatedly trained to carry out an image enhancement task, different diffusion models are trained and evaluated, and a model for different terminal devices is selected according to the performance.
2. The method according to claim 1, wherein, The step (3) specifically comprises the following steps: Step (3.1), defining a compressed sensing model: The compression sensing model comprises an encoder and a decoder The encoder is used for compressing a degraded image x of an underwater image dataset into latent space variables The decoder is used for restoring the latent space variables into a reconstructed image Specifically: Degraded images of an underwater image dataset given in RGB space where H denotes the height of the degraded image and W denotes the width of the degraded image; encoder of the compressive sensing model compressing the degraded image x into latent space variables where h denotes the height of the latent space variable and w denotes the width of the latent space variable decoder of the compressive sensing model transforming the latent space variables into a reconstructed image i.e. Step (3.2), defining a diffusion model: According to the neural network module obtained in step (2.3) as the skeleton of the diffusion model, that is, the U-Net network structure; U-Net∈ θ is a symmetric structure of an encoder-decoder neural network, the encoder is composed of multiple convolutional layers, and the decoder is composed of multiple transposed convolutional layers, and the convolutional layers and the transposed convolutional layers use skip connections and attention mechanisms; Step (3.3), defining a conditional encoder: Conditional encoder τ θ consisting of multiple convolutional layers for projecting the color distribution prior y to the latent space prior τ θ (y); Step (3.4), defining a loss function: The loss function of the compressed sensing model is: wherein, L per represents a perceptual loss, L reg represents a regularizing loss, W per represents a perceptual loss weight, W reg represents a regularizing loss weight; The loss function of the diffusion model is: wherein ∈ denotes a noise subject to a Gaussian distribution, ∈ θ denotes a noise predicted by the U-Net network model, and t denotes a certain time in a time series.
3. The underwater image processing method based on a lightweight diffusion model according to claim 1 or 2, characterized in that, The step (4) specifically comprises the following steps: Step (4.1), training the compressed sensing model: Set the training epoch, and repeat the following steps: Step (4.1.1), obtaining a degraded image x0from the underwater image dataset, input to the encoder of the compressive sensing model obtaining the latent variable Step (4.1.2), latent variable Input to decoder of compressive sensing model Obtaining a reconstructed image Step (4.1.3), updating the compressed sensing model parameters using a gradient descent algorithm according to the loss function L CPM Compute the loss for this round; update the compressed sensing model parameters using a gradient descent algorithm; Step (4.2), training the diffusion model: Set the number of training rounds epoch; initialize the noise mechanism a t Repeat the following steps at each time t in the time sequence Uniform({1,…,T}): Step (4.2.1), forward process: Step (4.2.1.1), obtaining a degraded image x0from the underwater dataset, input to the encoder of the compressive sensing model Obtaining latent variables Step (4.2.1.2), according to noise mechanism α t Obtain noise level ∈; Step (4.2.1.3), noise level e added to latent variable Obtaining noise latent variable After the forward pass, the complete noise latent variable is obtained, which is subject to a Gaussian distribution Step (4.2.2), backward process: Step (4.2.2.1), obtaining full noise latent variables Step (4.2.2.2), using a conditional encoder τ θ Projecting the color distribution prior y to the latent space gives the latent space prior τ θ (y); Step (4.2.2.3), latent space prior τ θ (y) and noise latent variable Encoder of diffusion model Decoder of diffusion model Output corresponding noise ∈ θ ; Step (4.2.2.4), using the loss function L LDM Computing the loss; Step (4.2.2.5), according to the formula obtain latent variables where σ t and z both represent constants, Step (4.2.2.6), latent variable Input to decoder of compressive sensing model Obtaining an enhanced image Step (4.2.2.7), updating the diffusion model parameters using gradient descent.
Citation Information
Patent Citations
Lightweight self-attention module, neural network model and search method of neural network framework
CN111539524A
Neural network construction method and device
CN111931904A