Image defogging method based on wavelet domain diffusion model

Abstract

The application discloses an image defogging method based on a wavelet field diffusion model, and is used for solving the problems of long reasoning time, large calculation amount and unstable recovery effect in the prior art. The method comprises the following steps: constructing a training set and a test set; performing wavelet field feature enhancement processing on all foggy images and clear images in the training set respectively, so as to obtain corresponding secondary components and secondary wavelet field enhanced features; loading random Gaussian noise on the secondary low-frequency approximate component in the secondary component of each clear image, and then inputting the secondary low-frequency approximate component of the corresponding foggy image and the secondary low-frequency approximate component of each clear image into a noise prediction network at the same time and outputting predicted noise; denoising the secondary low-frequency approximate component of each foggy image based on the predicted noise, and then performing inverse discrete Haar wavelet transform processing to obtain secondary recovery features of the foggy image; training a wavelet field diffusion model by using a loss function on the noise prediction network; and obtaining a defogged image based on the model.

Description

Technical Field

[0001] This invention relates to image processing methods, specifically to an image dehazing method based on a wavelet domain diffusion model. Background Technology

[0002] The scattering and absorption of light by particles such as water droplets and dust suspended in the atmosphere can easily lead to degradation phenomena such as reduced contrast, color distortion, and blurred details in images captured by imaging systems. Therefore, image dehazing is generally required after image capture. Image dehazing is a key fundamental task in the fields of computer vision and image processing, and it has wide application value in fields such as autonomous driving, video surveillance, and remote sensing monitoring.

[0003] Traditional image dehazing methods often employ prior information-based, convolutional neural network-based, or generative adversarial network-based methods. However, prior information-based image dehazing methods fail in specific scenarios such as snow, sky, and white buildings, while convolutional neural networks and generative adversarial networks are prone to artifacts and unclear restored details.

[0004] In recent years, researchers have discovered that image dehazing methods based on diffusion models exhibit superior performance in image generation tasks. These methods restore image details by defining forward diffusion and backward denoising processes. However, this approach suffers from excessively long inference times and also faces challenges such as high computational cost and unstable restoration results. Summary of the Invention

[0005] The purpose of this invention is to provide an image dehazing method based on a wavelet domain diffusion model, which solves the technical problems of existing image dehazing methods such as long inference time, large computational load, and unstable restoration effect.

[0006] To achieve the above objectives, the technical solution provided by this invention is as follows:

[0007] A wavelet domain diffusion model-based image dehazing method, characterized by the following steps:

[0008] Step 1: Construct a training set and a test set; the training set includes multiple foggy images and corresponding multiple clear images; the test set includes multiple foggy images.

[0009] Step 2: Perform wavelet domain feature enhancement processing on all foggy and clear images in the training set to obtain the first-level components and their first-level wavelet domain enhancement features for each foggy and clear image; the first-level components include a first-level low-frequency approximation component, a first-level horizontal high-frequency component, a first-level vertical high-frequency component, and a first-level diagonal high-frequency component.

[0010] Step 3: Perform wavelet domain feature enhancement processing again on the first-level low-frequency approximation component of each foggy image and each clear image to obtain the second-level component and its second-level wavelet domain enhancement feature of each foggy image and each clear image; the second-level component includes the second-level low-frequency approximation component, the second-level horizontal high-frequency component, the second-level vertical high-frequency component and the second-level diagonal high-frequency component;

[0011] Step 4: Apply random Gaussian noise to the second-order low-frequency approximation component of each clear image, and then input it and the second-order low-frequency approximation component of each foggy image into the noise prediction network, and output the predicted noise.

[0012] Step 5: Based on the predicted noise, the second-order low-frequency approximation component of each foggy image is denoised using the DDIM sampling method;

[0013] Step 6: Combining the first-level wavelet domain enhancement features of each clear image obtained in Step 2 and the second-level wavelet domain enhancement features of each clear image obtained in Step 3, perform inverse discrete Haar wavelet transform processing on the denoising results of the second-level low-frequency approximation components of each foggy image in Step 5 to obtain the second-level restoration features of each foggy image.

[0014] Step 7: Using the loss function, combined with the methods in steps 4 to 6, iteratively train the noise prediction network until the loss function converges, and obtain the wavelet domain diffusion model.

[0015] Step 8: Process all foggy images in the test set according to the methods in steps 2 to 3, then input them into the wavelet domain diffusion model obtained in step 7, and output the corresponding defogging images, thereby realizing image defogging based on the wavelet domain diffusion model.

[0016] Furthermore, step 2 specifically involves:

[0017] 2.1: Construct a first wavelet domain feature enhancement unit, which includes a wavelet transform module and a feature enhancement module connected to the output of the wavelet transform module;

[0018] 2.2: Wavelet domain feature enhancement processing was performed on all foggy images in the training set;

[0019] 2.2.1: Input all foggy images into the wavelet transform module in sequence for discrete Haar wavelet transform processing to obtain the first-level low-frequency approximation component LLh, the first-level horizontal high-frequency component LHh, the first-level vertical high-frequency component HLh, and the first-level diagonal high-frequency component HHh for each foggy image;

[0020] 2.2.2: The first-level horizontal high-frequency component LHh, the first-level vertical high-frequency component HLh, and the first-level diagonal high-frequency component HHh of each foggy image are input into the feature enhancement module for wavelet domain feature enhancement processing to obtain the first-level wavelet domain enhanced feature En1h of each foggy image;

[0021] 2.3: Perform wavelet domain feature enhancement processing on all clear images in the training set;

[0022] 2.3.1: Input all clear images into the wavelet transform module in sequence for discrete Haar wavelet transform processing to obtain the first-level low-frequency approximation component LLC, the first-level horizontal high-frequency component LHc, the first-level vertical high-frequency component HLc, and the first-level diagonal high-frequency component HHc for each clear image.

[0023] 2.3.2: The first-level horizontal high-frequency component LHc, the first-level vertical high-frequency component HLc, and the first-level diagonal high-frequency component HHc of each clear image are respectively input into the feature enhancement module for wavelet domain feature enhancement processing to obtain the first-level wavelet domain enhanced feature En1c of each clear image.

[0024] Further, in step 2.1, the feature enhancement module is a convolutional residual module, including convolution head module one, convolution head module two, convolution head module three, cross attention module one, cross attention module two, strided convolution module one, strided convolution module two, strided convolution module three, and convolution tail module;

[0025] Convolution head module 1, convolution head module 2, and convolution head module 3 are residual modules with the same architecture. Their input terminals are used to receive the first-level horizontal high-frequency component LHh, the first-level vertical high-frequency component HLh, and the first-level diagonal high-frequency component HHh output by the wavelet transform module, respectively. The output terminal of convolution head module 1 is connected to the input terminals of cross-attention module 1 and strided convolution module 1, respectively. The output terminal of convolution head module 2 is connected to the input terminals of cross-attention module 2 and strided convolution module 2, respectively. The output terminal of convolution head module 3 is connected to the input terminals of cross-attention module 1 and cross-attention module 2, respectively.

[0026] The cross-attention module one and cross-attention module two are residual modules with the same architecture, and their outputs are respectively connected to the input of the strided convolution module three.

[0027] The strided convolution module one, strided convolution module two, and strided convolution module three are residual modules with the same architecture, and their outputs are respectively connected to the convolution tail module, which is a residual module.

[0028] Furthermore, step 2.2.2 specifically includes:

[0029] A. Input the first-level horizontal high-frequency component LHh of each foggy image into the first convolution head module for convolution and activation processing, and output the corresponding feature map to the first cross-attention module and the first strided convolution module.

[0030] The first-level vertical high-frequency component HLh of each foggy day image is input into the second convolution head module for convolution and activation processing, and the corresponding feature map is output to the second cross-attention module and the second strided convolution module.

[0031] The first-order diagonal high-frequency component HHh of each foggy day image is input into the third convolution head module for convolution and activation processing, and the corresponding feature map is output to the first and second cross-attention modules.

[0032] B. Cross-attention module one and cross-attention module two process the received feature maps respectively and output the corresponding enhanced feature maps to strided convolution module three;

[0033] C. Stepped convolution module 1, stepped convolution module 2, and stepped convolution module 3 respectively input the corresponding feature maps and enhanced feature maps into the convolution tail module, and output the first-level wavelet domain enhanced feature En1h for each foggy image.

[0034] Furthermore, step 3 specifically involves:

[0035] 3.1: Construct a second wavelet domain feature enhancement unit, which has the same structure as the first wavelet domain feature enhancement unit;

[0036] 3.2: Wavelet domain feature enhancement processing is performed on the first-order components of each foggy image;

[0037] 3.2.1: Input the first-level low-frequency approximation component LLh of each foggy image into the wavelet transform module of the second wavelet domain feature enhancement unit for discrete Haar wavelet transform processing to obtain the second-level low-frequency approximation component LL2h, the second-level horizontal high-frequency component LH2h, the second-level vertical high-frequency component HL2h and the second-level diagonal high-frequency component HH2h of each foggy image;

[0038] 3.2.2: The second-order horizontal high-frequency component LH2h, the second-order vertical high-frequency component HL2h, and the second-order diagonal high-frequency component HH2h of each foggy image are respectively input into the feature enhancement module of the second wavelet domain feature enhancement unit for wavelet domain feature enhancement processing, thereby obtaining the second-order wavelet domain enhanced feature En2h of each foggy image;

[0039] 3.3: Wavelet domain feature enhancement processing is performed on the first-level components of each clear image;

[0040] 3.3.1: The first-level low-frequency approximation component LLc of each clear image is input into the wavelet transform module of the second wavelet domain feature enhancement unit for discrete Haar wavelet transform processing to obtain the second-level low-frequency approximation component LL2c, the second-level horizontal high-frequency component LH2c, the second-level vertical high-frequency component HL2c and the second-level diagonal high-frequency component HH2c of each clear image.

[0041] 3.3.2: The second-order horizontal high-frequency component LH2c, the second-order vertical high-frequency component HL2c, and the second-order diagonal high-frequency component HH2c of each clear image are respectively input into the feature enhancement module of the second wavelet domain feature enhancement unit for wavelet domain feature enhancement processing, thereby obtaining the second-order wavelet domain enhanced feature En2c of each clear image.

[0042] Furthermore, step 4 specifically involves:

[0043] 4.1: Apply random Gaussian noise of the same size to the second-order low-frequency approximation component LL2c of each clear image described in step 3.3.1. This yields the second-order low-frequency approximation components of the noisy sharp image at different time steps for each sharp image;

[0044] 4.2: The second-order low-frequency approximation component LL2h of each foggy image described in step 3.2.1, and the second-order low-frequency approximation components of the noisy clear image at different time steps of each clear image described in step 4.1, are simultaneously input into the noise prediction network to output the predicted noise. .

[0045] Furthermore, in step 4.1, the expressions for the second-order low-frequency approximation components of the noisy sharp image at different time steps for each sharp image are as follows:

[0046] ;

[0047] In the formula, This represents the second-order low-frequency approximation component of the noisy yet sharp image at time step i; Here are the noise scheduling parameters at time step i; This corresponds to the second-order low-frequency approximation component of the clear image;

[0048] In step 4.2, the noise prediction network is a UNet architecture network.

[0049] Furthermore, step 6 specifically includes:

[0050] 6.1: Combining the second-order wavelet domain enhancement feature En2c of each clear image obtained in step 3.3.2, the denoising results of the second-order low-frequency components of each foggy image are processed by inverse discrete Haar wavelet transform to obtain the first-order restoration feature Re1 of each foggy image.

[0051] 6.2: Combining the first-level wavelet domain enhancement feature En1c of each clear image obtained in step 2.3.2, perform inverse discrete Haar wavelet transform on the first-level restoration feature Re1 of each foggy image to obtain the second-level restoration feature Re2 of each foggy image.

[0052] Furthermore, in step 7, the expression for the loss function is as follows:

[0053] ;

[0054] In the formula, Represents the loss function; These represent the preset weighting coefficients;

[0055] The noise loss is represented by the root mean square loss function, and ;

[0056] To represent the loss of high-frequency information, it uses the TV loss function and the feature edge loss function;

[0057] This represents content loss, which uses the L1 loss function, and , where C represents the clear image in the training set;

[0058] The structural loss is represented by a structural similarity function, and .

[0059] Furthermore, in step 7, the expression for the TV loss function is as follows:

[0060] ;

[0061] The expression for the feature edge loss function is as follows:

[0062] .

[0063] Compared with the prior art, the beneficial effects of the present invention are as follows:

[0064] 1. This invention introduces wavelet domain feature enhancement method into diffusion model to obtain wavelet domain diffusion model. After wavelet domain feature enhancement processing, it can reconstruct local details of foggy images with high precision, and the model is less prone to artifacts and has higher robustness.

[0065] 2. In the process of wavelet domain feature enhancement, the image size of each sub-band is halved after discrete wavelet decomposition. This significantly reduces computational complexity while ensuring the overall quality of the restored image without losing detailed information. Compared with the spatial domain diffusion model, the computational complexity and amount of computation are greatly reduced, thereby accelerating the inference time.

[0066] 3. This invention addresses four aspects: noise loss, high-frequency information loss, content loss, and structural loss. It combines the root mean square loss function, TV loss function, feature edge loss function, L1 loss function, and structural similarity function to train the noise prediction network from multiple angles, thereby obtaining a high-precision wavelet domain diffusion model and ensuring the quality of image restoration. Attached Figure Description

[0067] Figure 1 This is a flowchart illustrating steps 2.2 and 3.2 in an embodiment of the present invention;

[0068] Figure 2 This is a flowchart illustrating steps 2.3 and 3.3 in an embodiment of the present invention;

[0069] Figure 3 This is a schematic diagram of the wavelet domain feature enhancement process in an embodiment of the present invention;

[0070] Figure 4 This is a flowchart illustrating step 6 in an embodiment of the present invention. Detailed Implementation

[0071] To make the objectives, advantages, and features of the present invention clearer, the present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments. Those skilled in the art should understand that these embodiments are merely used to explain the technical principles of the present invention and are not intended to limit the scope of protection of the present invention.

[0072] This embodiment provides an image dehazing method based on a wavelet domain diffusion model, including the following steps:

[0073] Step 1: Construct a training set and a test set; the training set includes multiple foggy images H and corresponding multiple clear images C; the test set includes multiple foggy images H1.

[0074] Step 2: Perform wavelet domain feature enhancement processing on all foggy images H and clear images C in the training set to obtain the first-level components and their first-level wavelet domain enhancement features for each foggy image and each clear image.

[0075] 2.1: Constructing the first wavelet domain feature enhancement unit, which includes a wavelet transform module and a feature enhancement module connected to the output of the wavelet transform module. The wavelet transform module is a discrete Haar wavelet transform module, whose core function is to perform multi-scale orthogonal decomposition on the discrete image, separating low-frequency approximations from high-frequency details, thereby achieving efficient compression, denoising, feature extraction, and real-time analysis.

[0076] like Figure 3 As shown, the feature enhancement module in this embodiment is a convolutional residual module, including convolution head module 1, convolution head module 2, convolution head module 3, cross attention module 1, cross attention module 2, strided convolution module 1, strided convolution module 2, strided convolution module 3, and convolution tail module.

[0077] Among them, Convolution Head Module 1, Convolution Head Module 2, and Convolution Head Module 3 are residual modules with the same architecture, each including a regular convolutional layer and an activation function. The input terminals of Convolution Head Module 1, Convolution Head Module 2, and Convolution Head Module 3 are used to receive the first-level horizontal high-frequency component LHh, the first-level vertical high-frequency component HLh, and the first-level diagonal high-frequency component HHh output by the wavelet transform module, respectively. The output terminal of Convolution Head Module 1 is connected to the input terminals of Cross Attention Module 1 and Stepped Convolution Module 1, respectively. The output terminal of Convolution Head Module 2 is connected to the input terminals of Cross Attention Module 2 and Stepped Convolution Module 2, respectively. The output terminal of Convolution Head Module 3 is connected to the input terminals of Cross Attention Module 1 and Cross Attention Module 2, respectively.

[0078] Cross-attention module 1 and cross-attention module 2 are residual modules with the same architecture, each including a strided convolutional layer and an activation function. The outputs of cross-attention module 1 and cross-attention module 2 are connected to the input of strided convolutional module 3.

[0079] Stride convolution module 1, stride convolution module 2, and stride convolution module 3 are residual modules with the same architecture, and their outputs are connected to the convolution tail module respectively.

[0080] The convolution tail module is a residual module, which includes ordinary convolutional layers and activation functions.

[0081] In this embodiment, the core function of the feature enhancement module is to specifically enhance target features and suppress interference information in the multi-scale space of wavelet decomposition, thereby improving the image's recognizability and providing a higher-quality data foundation for subsequent analysis, detection, or recognition tasks. The feature enhancement module utilizes the multi-resolution and time-frequency localization characteristics of wavelet transform to decompose the image into sub-bands of different scales and directions, and then applies differentiated enhancement strategies to different sub-bands. Compared to traditional spatial / temporal enhancement methods, this approach can more accurately protect key features such as edges and details while avoiding the introduction of artifacts.

[0082] 2.2: For example Figure 1 As shown, wavelet domain feature enhancement processing is performed on all foggy images H in the training set.

[0083] 2.2.1: Input all foggy images H into the wavelet transform module in sequence for discrete Haar wavelet transform processing to obtain the first-level low-frequency approximation component LLh, the first-level horizontal high-frequency component LHh, the first-level vertical high-frequency component HLh, and the first-level diagonal high-frequency component HHh for each foggy image.

[0084] In this embodiment, after discrete wavelet decomposition, the image size of each sub-band is halved step by step, resulting in only a fraction of the original image size. While significantly reducing computational complexity, it can ensure the overall quality of the restored image without losing detailed information, and the computational complexity and amount of computation are greatly reduced compared with the spatial domain diffusion model.

[0085] 2.2.2: The first-level horizontal high-frequency component LHh, the first-level vertical high-frequency component HLh, and the first-level diagonal high-frequency component HHh of each foggy image are input into the feature enhancement module for wavelet domain feature enhancement processing to obtain the first-level wavelet domain enhanced feature En1h of each foggy image.

[0086] like Figure 3 As shown, the specific process of wavelet domain feature enhancement is as follows:

[0087] A. Input the first-level horizontal high-frequency component LHh of each foggy image into the first convolution head module for convolution and activation processing, and output the corresponding feature map to the first cross-attention module and the first strided convolution module.

[0088] The first-level vertical high-frequency component HLh of each foggy day image is input into the second convolution head module for convolution and activation processing, and the corresponding feature map is output to the second cross-attention module and the second strided convolution module.

[0089] The first-order diagonal high-frequency component HHh of each foggy day image is input into the third convolution head module for convolution and activation processing, and the corresponding feature map is output to the first and second cross-attention modules.

[0090] B. Cross-attention module one and cross-attention module two process the received feature maps respectively and output the corresponding enhanced feature maps to strided convolution module three.

[0091] In this process, taking the first cross-attention module as an example, the first cross-attention module first converts one of the received feature maps into a query vector Q through depthwise separable convolution, and converts the other received feature map into a key vector K and a value vector V. Then, by calculating the similarity between the query vector Q and the key vector K, the attention weights are obtained.

[0092] The specific calculation method is as follows: Combine the query vector Q and the transposed key vector... Perform matrix multiplication to obtain the attention score, then divide by the scaling factor. Then, the attention scores are normalized using the softmax function (activation function) to obtain the attention weights.

[0093] Finally, the attention weights are used to perform a weighted summation on the value vector V, and the corresponding enhanced feature map is output.

[0094] In this embodiment, the cross-attention module in the feature enhancement module uses horizontal and vertical information to supplement diagonal information details, enabling high-precision reconstruction of local details. The model is less prone to artifacts and has higher robustness.

[0095] C. Stepped convolution module 1, stepped convolution module 2, and stepped convolution module 3 respectively input the corresponding feature maps and enhanced feature maps into the convolution tail module, and output the first-level wavelet domain enhanced feature En1h for each foggy image.

[0096] 2.3: For example Figure 2 As shown, wavelet domain feature enhancement processing is performed on all clear images C in the training set.

[0097] 2.3.1: Input all clear images C into the wavelet transform module in sequence for discrete Haar wavelet transform processing to obtain the first-level low-frequency approximation component LLC, the first-level horizontal high-frequency component LHc, the first-level vertical high-frequency component HLc, and the first-level diagonal high-frequency component HHc for each clear image.

[0098] 2.3.2: The first-level horizontal high-frequency component LHc, the first-level vertical high-frequency component HLc, and the first-level diagonal high-frequency component HHc of each clear image are respectively input into the feature enhancement module for wavelet domain feature enhancement processing to obtain the first-level wavelet domain enhanced feature En1c of each clear image.

[0099] Step 3: Perform wavelet domain feature enhancement processing on the first-level components of each foggy image and each clear image again to obtain the second-level components and their second-level wavelet domain enhancement features of each foggy image and each clear image.

[0100] 3.1: Construct a second wavelet domain feature enhancement unit. The structure of the second wavelet domain feature enhancement unit is the same as that of the first wavelet domain feature enhancement unit.

[0101] 3.2: For example Figure 1 As shown, wavelet domain feature enhancement processing is performed on the first-level components of each foggy image.

[0102] 3.2.1: Input the first-level low-frequency approximation component LLh of each foggy image into the wavelet transform module of the second wavelet domain feature enhancement unit for discrete Haar wavelet transform processing to obtain the second-level low-frequency approximation component LL2h, the second-level horizontal high-frequency component LH2h, the second-level vertical high-frequency component HL2h and the second-level diagonal high-frequency component HH2h of each foggy image;

[0103] 3.2.2: The second-order horizontal high-frequency component LH2h, the second-order vertical high-frequency component HL2h, and the second-order diagonal high-frequency component HH2h of each foggy image are respectively input into the feature enhancement module of the second wavelet domain feature enhancement unit for wavelet domain feature enhancement processing, thereby obtaining the second-order wavelet domain enhanced feature En2h of each foggy image.

[0104] 3.3: For example Figure 2 As shown, wavelet domain feature enhancement processing is performed on the first-level components of each clear image.

[0105] 3.3.1: The first-level low-frequency approximation component LLc of each clear image is input into the wavelet transform module of the second wavelet domain feature enhancement unit for discrete Haar wavelet transform processing to obtain the second-level low-frequency approximation component LL2c, the second-level horizontal high-frequency component LH2c, the second-level vertical high-frequency component HL2c and the second-level diagonal high-frequency component HH2c of each clear image.

[0106] 3.3.2: The second-order horizontal high-frequency component LH2c, the second-order vertical high-frequency component HL2c, and the second-order diagonal high-frequency component HH2c of each clear image are respectively input into the feature enhancement module of the second wavelet domain feature enhancement unit for wavelet domain feature enhancement processing, thereby obtaining the second-order wavelet domain enhanced feature En2c of each clear image.

[0107] Step 4: Apply random Gaussian noise to the second-order component of each clear image, then input it along with the second-order component of each foggy image into the noise prediction network, and output the predicted noise. .

[0108] 4.1: Apply random Gaussian noise of the same size to the second-order low-frequency approximation component LL2c of each clear image described in step 3.3.1. We obtain the second-order low-frequency approximation components of the noisy sharp image at different time steps i for each sharp image.

[0109] In this embodiment, the expressions for the second-order low-frequency approximation components of the noisy sharp image at different time steps i are as follows:

[0110] ;

[0111] In the formula, This represents the second-order low-frequency approximation component of the noisy yet sharp image at time step i; Here are the noise scheduling parameters at time step i; This corresponds to the second-order low-frequency approximation component of the clear image.

[0112] 4.2: Combine the second-order low-frequency approximation component LL2h of each foggy image described in step 3.2.1, and the second-order low-frequency approximation components of the noisy clear image of each clear image described in step 4.1 at different time steps. Simultaneously, it is input into the noise prediction network, thereby outputting the predicted noise. .

[0113] The noise prediction network described in this embodiment is the classic UNet architecture model network.

[0114] Step 5: Based on the predicted noise, the second-order low-frequency approximation component LL2h of each foggy image is denoised using the DDIM (Denoising Diffusion Implicit Models) sampling method to obtain the denoised result DeLL2h of the second-order low-frequency component of the corresponding foggy image.

[0115] Step 6: As Figure 4 As shown, combining the first-level wavelet domain enhancement features of each clear image obtained in step 2 and the second-level wavelet domain enhancement features of each clear image obtained in step 3, the denoising results of the second-level components of each foggy image in step 5 are processed by inverse discrete Haar wavelet transform to obtain the second-level restoration features of each foggy image.

[0116] 6.1: Combining the second-level wavelet domain enhancement feature En2c of each clear image obtained in step 3.3.2, the denoising result DeLL2h of the second-level low-frequency component of each foggy image is processed by inverse discrete Haar wavelet transform to obtain the first-level restoration feature Re1 of each foggy image.

[0117] 6.2: Combining the first-level wavelet domain enhancement feature En1c of each clear image obtained in step 2.3.2, perform inverse discrete Haar wavelet transform on the first-level restoration feature Re1 of each foggy image to obtain the second-level restoration feature Re2 of each foggy image.

[0118] Step 7: Using the loss function and the methods from Steps 4 to 6, iteratively train the noise prediction network to obtain the wavelet domain diffusion model.

[0119] The expression for the loss function described in this embodiment is as follows:

[0120]

[0121] In the formula, Represents the loss function; These represent the preset weighting coefficients;

[0122] The noise loss is represented by the root mean square loss function (RMS). ),and ;in, This represents the predicted noise output from step 4.2; This represents the random Gaussian noise loaded in step 4.1.

[0123] This represents the loss of high-frequency information, which is expressed using the TV loss function ( and feature edge loss function, and Among them, the TV loss function The expression is as follows:

[0124] ;

[0125] In the formula, En1h represents the first-level wavelet domain enhancement feature of the foggy image in the training set; En2h represents the second-level wavelet domain enhancement feature of the foggy image in the training set.

[0126] The Total Variation Loss (TV) function is a regularized loss used to constrain the smoothness of model output. Its core function is to suppress artifacts and unnecessary high-frequency fluctuations in the output, while preserving key structures such as edges as much as possible, making the model output more regular.

[0127] Feature edge loss function The expression is as follows:

[0128] ;

[0129] In the formula, En1c represents the first-level wavelet domain enhancement feature corresponding to the clear image in the training set; En2c represents the second-level wavelet domain enhancement feature corresponding to the clear image in the training set.

[0130] Feature edge loss function is a type of supervised / regularized loss oriented towards the edge structure of an image. Its core function is to force the model to accurately capture and retain high-frequency structural information such as edges and contours of the image in the feature space, thereby solving the problem that ordinary losses (such as MSE) are prone to blurring edges and losing details.

[0131] This represents content loss, which uses the L1 loss function, and , where Re2 represents the secondary restoration feature of the foggy image in the training set; C represents the clear image in the training set.

[0132] The structural loss is represented using the structural similarity function (SSIM), and .

[0133] Step 8: Process all foggy images H1 in the test set according to the methods in steps 2 to 3: First, input the foggy images H1 in the test set into the wavelet transform module to obtain the first low-frequency approximation component LLh1, the first horizontal high-frequency component LHh1, the first vertical high-frequency component HLh1, and the first diagonal high-frequency component HHh1 of the corresponding foggy image H1.

[0134] Next, the first horizontal high-frequency component LHh1, the first vertical high-frequency component HLh1, and the first diagonal high-frequency component HHh1 of the corresponding foggy image H1 are input into the feature enhancement module to obtain the first-order wavelet domain enhancement feature En1h1 of the corresponding foggy image H1.

[0135] Then, the first horizontal high-frequency component LLh1 of the corresponding foggy image H1 is input into the wavelet transform module to obtain the second-order low-frequency approximation component LL2h1, the second-order horizontal high-frequency component LH2h1, the second-order vertical high-frequency component HL2h1, and the second-order diagonal high-frequency component HH2h1 of the corresponding foggy image H1.

[0136] Finally, the second-order horizontal high-frequency component LH2h1, the second-order vertical high-frequency component HL2h1, and the second-order diagonal high-frequency component HH2h1 of the corresponding foggy image H1 are input into the feature enhancement module to obtain the second-order wavelet domain enhancement feature En2h1 of the corresponding foggy image H1.

[0137] After completing the above steps, input the processing results into the wavelet domain diffusion model obtained in step 7, and output the corresponding dehazed image, thereby realizing image dehazing based on the wavelet domain diffusion model.

[0138] The working process of the wavelet domain diffusion model is as follows: The second-order horizontal high-frequency component LL2h1 of the corresponding foggy image H1 is processed using the DDIM sampling method to obtain the second-order low-frequency denoising feature DeLL2h1 of the foggy image H1; then, the second-order low-frequency denoising feature DeLL2h1 and its second-order wavelet domain enhancement feature En2h1 of the foggy image H1 are processed by inverse discrete Haar wavelet transform to obtain the first-order low-frequency denoising feature DeLL1h1 of the corresponding foggy image H1; finally, the first-order low-frequency denoising feature DeLL1h1 and its first-order wavelet domain enhancement feature En1h1 of the foggy image H1 are processed by inverse discrete Haar wavelet transform to obtain the defogging image corresponding to the foggy image H1.

[0139] This invention combines the multi-scale analysis capability of wavelet transform with the generative modeling advantages of diffusion models, achieving efficient and high-quality image dehazing.

[0140] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit it. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some or all of the technical features therein, and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the scope of the technical solutions of the present invention.

Claims

1. An image dehazing method based on a wavelet domain diffusion model, characterized in that, Includes the following steps: Step 1: Construct a training set and a test set; the training set includes multiple foggy images and corresponding multiple clear images; the test set includes multiple foggy images. Step 2: Perform wavelet domain feature enhancement processing on all foggy and clear images in the training set to obtain the first-level components and their first-level wavelet domain enhancement features for each foggy and clear image; the first-level components include a first-level low-frequency approximation component, a first-level horizontal high-frequency component, a first-level vertical high-frequency component, and a first-level diagonal high-frequency component. Step 3: Perform wavelet domain feature enhancement processing again on the first-level low-frequency approximation component of each foggy image and each clear image to obtain the second-level component and its second-level wavelet domain enhancement feature of each foggy image and each clear image; the second-level component includes the second-level low-frequency approximation component, the second-level horizontal high-frequency component, the second-level vertical high-frequency component and the second-level diagonal high-frequency component; Step 4: Apply random Gaussian noise to the second-order low-frequency approximation component of each clear image, and then input it and the second-order low-frequency approximation component of each foggy image into the noise prediction network, and output the predicted noise. Step 5: Based on the predicted noise, the second-order low-frequency approximation component of each foggy image is denoised using the DDIM sampling method; Step 6: Combining the first-level wavelet domain enhancement features of each clear image obtained in Step 2 and the second-level wavelet domain enhancement features of each clear image obtained in Step 3, perform inverse discrete Haar wavelet transform processing on the denoising results of the second-level low-frequency approximation components of each foggy image in Step 5 to obtain the second-level restoration features of each foggy image. Step 7: Using the loss function, combined with the methods in steps 4 to 6, iteratively train the noise prediction network until the loss function converges to obtain the wavelet domain diffusion model. Step 8: Process all foggy images in the test set according to the methods in steps 2 to 3, then input them into the wavelet domain diffusion model obtained in step 7, and output the corresponding defogging images, thereby realizing image defogging based on the wavelet domain diffusion model.

2. The image dehazing method based on the wavelet domain diffusion model according to claim 1, characterized in that, Step 2 is as follows: 2.1: Construct a first wavelet domain feature enhancement unit, which includes a wavelet transform module and a feature enhancement module connected to the output of the wavelet transform module; 2.2: Wavelet domain feature enhancement processing was performed on all foggy images in the training set; 2.2.1: Input all foggy images into the wavelet transform module in sequence for discrete Haar wavelet transform processing to obtain the first-level low-frequency approximation component LLh, the first-level horizontal high-frequency component LHh, the first-level vertical high-frequency component HLh, and the first-level diagonal high-frequency component HHh for each foggy image; 2.2.2: The first-level horizontal high-frequency component LHh, the first-level vertical high-frequency component HLh, and the first-level diagonal high-frequency component HHh of each foggy image are input into the feature enhancement module for wavelet domain feature enhancement processing to obtain the first-level wavelet domain enhanced feature En1h of each foggy image; 2.3: Perform wavelet domain feature enhancement processing on all clear images in the training set; 2.3.1: Input all clear images into the wavelet transform module in sequence for discrete Haar wavelet transform processing to obtain the first-level low-frequency approximation component LLC, the first-level horizontal high-frequency component LHc, the first-level vertical high-frequency component HLc, and the first-level diagonal high-frequency component HHc for each clear image. 2.3.2: The first-level horizontal high-frequency component LHc, the first-level vertical high-frequency component HLc, and the first-level diagonal high-frequency component HHc of each clear image are respectively input into the feature enhancement module for wavelet domain feature enhancement processing to obtain the first-level wavelet domain enhanced feature En1c of each clear image.

3. The image dehazing method based on the wavelet domain diffusion model according to claim 2, characterized in that: In step 2.1, the feature enhancement module is a convolutional residual module, including convolution head module 1, convolution head module 2, convolution head module 3, cross attention module 1, cross attention module 2, strided convolution module 1, strided convolution module 2, strided convolution module 3, and convolution tail module; Convolution head module 1, convolution head module 2, and convolution head module 3 are residual modules with the same architecture. Their input terminals are used to receive the first-level horizontal high-frequency component LHh, the first-level vertical high-frequency component HLh, and the first-level diagonal high-frequency component HHh output by the wavelet transform module, respectively. The output terminal of convolution head module 1 is connected to the input terminals of cross-attention module 1 and strided convolution module 1, respectively. The output terminal of convolution head module 2 is connected to the input terminals of cross-attention module 2 and strided convolution module 2, respectively. The output terminal of convolution head module 3 is connected to the input terminals of cross-attention module 1 and cross-attention module 2, respectively. The cross-attention module one and cross-attention module two are residual modules with the same architecture, and their outputs are respectively connected to the input of the strided convolution module three. The strided convolution module one, strided convolution module two, and strided convolution module three are residual modules with the same architecture, and their outputs are respectively connected to the convolution tail module, which is a residual module.

4. The image dehazing method based on the wavelet domain diffusion model according to claim 3, characterized in that, Step 2.2.2 specifically refers to: A. Input the first-level horizontal high-frequency component LHh of each foggy image into the first convolution head module for convolution and activation processing, and output the corresponding feature map to the first cross-attention module and the first strided convolution module. The first-level vertical high-frequency component HLh of each foggy day image is input into the second convolution head module for convolution and activation processing, and the corresponding feature map is output to the second cross-attention module and the second strided convolution module. The first-order diagonal high-frequency component HHh of each foggy day image is input into the third convolution head module for convolution and activation processing, and the corresponding feature map is output to the first and second cross-attention modules. B. Cross-attention module one and cross-attention module two process the received feature maps respectively and output the corresponding enhanced feature maps to strided convolution module three; C. Stepped convolution module 1, stepped convolution module 2, and stepped convolution module 3 respectively input the corresponding feature maps and enhanced feature maps into the convolution tail module, and output the first-level wavelet domain enhanced feature En1h for each foggy image.

5. The image dehazing method based on the wavelet domain diffusion model according to claim 4, characterized in that, Step 3 specifically involves: 3.1: Construct a second wavelet domain feature enhancement unit, which has the same structure as the first wavelet domain feature enhancement unit; 3.2: Wavelet domain feature enhancement processing is performed on the first-order components of each foggy image; 3.2.1: Input the first-level low-frequency approximation component LLh of each foggy image into the wavelet transform module of the second wavelet domain feature enhancement unit for discrete Haar wavelet transform processing to obtain the second-level low-frequency approximation component LL2h, the second-level horizontal high-frequency component LH2h, the second-level vertical high-frequency component HL2h and the second-level diagonal high-frequency component HH2h of each foggy image; 3.2.2: The second-order horizontal high-frequency component LH2h, the second-order vertical high-frequency component HL2h, and the second-order diagonal high-frequency component HH2h of each foggy image are respectively input into the feature enhancement module of the second wavelet domain feature enhancement unit for wavelet domain feature enhancement processing, thereby obtaining the second-order wavelet domain enhanced feature En2h of each foggy image; 3.3: Wavelet domain feature enhancement processing is performed on the first-level components of each clear image; 3.3.1: The first-level low-frequency approximation component LLc of each clear image is input into the wavelet transform module of the second wavelet domain feature enhancement unit for discrete Haar wavelet transform processing to obtain the second-level low-frequency approximation component LL2c, the second-level horizontal high-frequency component LH2c, the second-level vertical high-frequency component HL2c and the second-level diagonal high-frequency component HH2c of each clear image. 3.3.2: The second-order horizontal high-frequency component LH2c, the second-order vertical high-frequency component HL2c, and the second-order diagonal high-frequency component HH2c of each clear image are respectively input into the feature enhancement module of the second wavelet domain feature enhancement unit for wavelet domain feature enhancement processing, thereby obtaining the second-order wavelet domain enhanced feature En2c of each clear image.

6. The image dehazing method based on the wavelet domain diffusion model according to claim 5, characterized in that, Step 4 is as follows: 4.1: Apply random Gaussian noise of the same size to the second-order low-frequency approximation component LL2c of each clear image described in step 3.3.

1. This yields the second-order low-frequency approximation components of the noisy sharp image at different time steps for each sharp image; 4.2: The second-order low-frequency approximation component LL2h of each foggy image described in step 3.2.1, and the second-order low-frequency approximation components of the noisy clear image at different time steps of each clear image described in step 4.1, are simultaneously input into the noise prediction network to output the predicted noise. .

7. The image dehazing method based on the wavelet domain diffusion model according to claim 6, characterized in that: In step 4.1, the expressions for the second-order low-frequency approximation components of the noisy sharp image at different time steps for each sharp image are as follows: ； In the formula, This represents the second-order low-frequency approximation component of the noisy yet sharp image at time step i; Here are the noise scheduling parameters at time step i; This corresponds to the second-order low-frequency approximation component of the clear image; In step 4.2, the noise prediction network is a UNet architecture network.

8. The image dehazing method based on the wavelet domain diffusion model according to claim 7, characterized in that, Step 6 specifically involves: 6.1: Combining the second-order wavelet domain enhancement feature En2c of each clear image obtained in step 3.3.2, the denoising results of the second-order low-frequency components of each foggy image are processed by inverse discrete Haar wavelet transform to obtain the first-order restoration feature Re1 of each foggy image. 6.2: Combining the first-level wavelet domain enhancement feature En1c of each clear image obtained in step 2.3.2, perform inverse discrete Haar wavelet transform on the first-level restoration feature Re1 of each foggy image to obtain the second-level restoration feature Re2 of each foggy image.

9. The image dehazing method based on the wavelet domain diffusion model according to claim 8, characterized in that: In step 7, the expression for the loss function is as follows: ； In the formula, Represents the loss function; These represent the preset weighting coefficients; The noise loss is represented by the root mean square loss function, and ; To represent the loss of high-frequency information, it uses the TV loss function and the feature edge loss function; This represents content loss, which uses the L1 loss function, and , where C represents the clear image in the training set; The structural loss is represented by a structural similarity function, and .

10. The image dehazing method based on the wavelet domain diffusion model according to claim 9, characterized in that: In step 7, the expression for the TV loss function is as follows: ； The expression for the feature edge loss function is as follows: 。

Images