A polarization recovery imaging method based on a conditional diffusion model

Abstract

The application discloses a polarization restoration imaging method based on a conditional diffusion model, which comprises the following steps: step one, shooting clear polarization images of underwater target objects and corresponding polarization images with different turbidity degrees; step two, dividing the shot polarization images into a training set, a verification set and a test set according to a ratio of 7:2:1; according to the idea of generating images based on a diffusion model, combining a polarization information preprocessing module, a residual convolution module and a self-attention mechanism module, a conditional diffusion network model is constructed, a data set is established, the polarization information of the images is fully utilized, the network is trained so that the network can output clear images according to input turbid polarization images and noise images subject to a standard normal distribution, the turbid polarization images are restored by using the conditional diffusion model, and good enhancement effect can be achieved, and stable performance can be maintained in various environments.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of deep learning and polarization imaging, in particular to a polarization restoration imaging method based on a conditional diffusion model. BACKGROUND

[0002] Computer vision technology is an important branch of artificial intelligence, which enables computers to interpret and understand image and video content like human eyes, and has great research significance. Computer vision technology can automate many tasks such as quality detection, monitoring and navigation, helping to reduce human intervention and improve efficiency and accuracy. For the field of autonomous driving, computer vision is one of the key technologies of autonomous vehicles. It enables vehicles to recognize road signs, pedestrians, other vehicles and obstacles, thereby achieving safe driving. However, in computer vision tasks, the target object often cannot be clearly captured due to the degradation effect of various scattering media such as fog, smoke and dust, so the low-quality blurred image obtained cannot be applied to other advanced vision tasks, and therefore the restoration work of foggy and hazy images is carried out.

[0003] At present, the methods for image dehazing can be roughly divided into two categories: physical-based methods and learning-based methods. The physical-based method is based on the atmospheric scattering model, and by constructing a mathematical model, the concentration of the scattering medium and other information are estimated to restore the clear image. This method has the advantages of high efficiency and fast speed. The learning-based method mainly uses the mapping ability of deep learning neural network, and through the method of training neural network, a network model that can restore hazy images to clear images is obtained. Compared with traditional physical methods, this method requires more data and computing resources, but it performs better in processing complex scene hazy images.

[0004] Compared with traditional intensity imaging technology, polarization imaging has certain advantages due to the additional physical information such as polarization degree and polarization angle of the target. However, due to the limitations of complex scenes and harsh environments, the effect of polarization imaging still has certain limitations, and a polarization restoration imaging method with wider applicability and higher image restoration quality is urgently needed. SUMMARY

[0005] In order to overcome the shortcomings of the prior art, the present application provides a polarization restoration imaging method based on a conditional diffusion model, which establishes a data set in a simulated environment, fully utilizes the polarization information of the image, and trains the network to enable it to output the corresponding clear image according to the input hazy polarization image and Gaussian noise image subject to standard normal distribution. Using this network to restore hazy polarization images can obtain high-quality clear images, and in many scenarios it can maintain good stability.

[0006] In order to achieve the above object, the technical scheme adopted by the present application is as follows: A polarization restoration imaging method based on a conditional diffusion model, comprising the following steps:

[0007] Step one, shoot clear polarization images of underwater target objects and corresponding polarization images of different turbidity levels;

[0008] Step two, divide the shot polarization images into a training set, a validation set and a test set according to a ratio of 7:2:1;

[0009] Step three, construct a conditional diffusion neural network according to the idea of generating images based on a diffusion model, the overall network structure is designed in a U shape, add a certain degree of random Gaussian noise associated with time to the clear polarization images of the target, obtain a feature map vector from the turbid polarization images through a polarization information preprocessing module, and input the feature map and the clear images with added noise into the network; first, pass through four encoders, which contain two residual convolution modules, two self-attention mechanism modules and one down-sampling module, to extract features, and then pass through four decoders for up-sampling, the decoder contains two residual convolution modules, two self-attention mechanism modules and one up-sampling module, a time parameter associated with the degree of added noise is embedded in each encoder and decoder to prompt the network to the degree of noise removal, the network outputs a predicted noise in the training stage, and the mean square error loss between the predicted noise and the standard Gaussian noise is used to optimize the network;

[0010] Step four, train the network, use the data set obtained in step two to train the conditional diffusion model constructed in step three so that it can output the corresponding clear images according to the input turbid images and Gaussian noise images subject to standard normal distribution;

[0011] Step five, restore the turbid polarization images to obtain the restored images with obvious visual effect.

[0012] Further, the clear polarization images of the underwater target objects shot in step one: the light beams emitted by the light source are sequentially irradiated to the target object after passing through the polarizer of the polarization modulation system and the beam expander, and then reach the defocus plane polarization camera after being reflected by the target object, so that the clear polarization images of the target object are obtained.

[0013] Further, the polarization images of the target objects under different turbidity levels shot in step one: gradually add scattering media in the water, and use a 500 million pixel polarization gigabit Ethernet industrial camera to shoot the target objects in different turbidity environments, so that the polarization images of the target objects in different polarization angles and different turbidity concentrations in the turbid underwater environment are obtained.

[0014] Further, a 532nm blue-green laser is used as the light source.

[0015] Further, the residual convolution module in the conditional diffusion neural network in step three is composed of three parts, the first part is a structure composed of a small batch normalization operation for feature map, a Swish activation function and a 3*3 convolution layer; the second part is a module for time embedding, which is composed of a Swish activation function and a 1*1 convolution layer; the third part is composed of a feature batch normalization layer, a Swish activation function, a Dropout regularization layer and a 3*3 convolution layer.

[0016] Further, the self-attention mechanism module in the conditional diffusion neural network in step three is composed of two parts, the first part is group normalization, which standardizes the input to improve the stability and efficiency of training and help reduce internal covariate shift in the network; the second part is a convolution projection layer, which is used to generate query (Q), key (K) and value (V) vectors, the query (Q) determines the focus of attention of the current element in the entire sequence; the key (K) provides the identification of each element, helping the query to find related elements; the value (V) carries the actual information content, and finally generates the output according to the matching degree of the query and the key.

[0017] Further, the target is placed in a glass tank filled with water to simulate an underwater environment.

[0018] Further, the glass tank is made of PMMA material.

[0019] Compared with the prior art, the present application has the following beneficial effects:

[0020] According to the idea of generating images based on the diffusion model, combined with the polarization information preprocessing module, the residual convolution module and the self-attention mechanism module, a conditional diffusion network model is constructed, a data set is established, and the polarization information of the image is fully utilized to train the network so that it can output a clear image according to the input turbid polarization image and noise image subject to standard normal distribution. Using the conditional diffusion model to restore the turbid polarization image can achieve good enhancement effect, and can maintain stable performance in various environments. BRIEF DESCRIPTION OF DRAWINGS

[0021] Figure 1 is the overall flowchart of the present application;

[0022] Figure 2 is the experimental device schematic diagram of the present application;

[0023] Figure 3 is the residual convolution module structure schematic diagram;

[0024] Figure 4 is the self-attention mechanism module structure schematic diagram;

[0025] Figure 5 is a conditional diffusion model recovery experiment effect diagram, (a) is a turbid polarization image, (b) is the image recovered by the present application, and (c) is a clear image;

[0026] Figure 6 is a table of recovery effect evaluation indexes SSIM and PSNR.

[0027] In the figure: light source 1, polarizer 2, beam expander 3, glass jar 4, underwater environment 5, target object 6, defocus plane polarization camera 7. DETAILED DESCRIPTION

[0028] The technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, rather than all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those of ordinary skill in the art without creative labor fall within the scope of the present application.

[0029] As shown in Figures 1-6 , the technical solutions adopted by the present application are as follows: the present embodiment provides a polarization recovery imaging method based on a conditional diffusion model, comprising the following steps:

[0030] Step one, shooting clear polarization images of underwater target objects and corresponding polarization images of different turbidity degrees;

[0031] Step two, dividing the shot polarization images into a training set, a validation set and a test set according to a ratio of 7:2:1;

[0032] Step three, constructing a conditional diffusion neural network according to the idea of generating images based on a diffusion model, and designing the overall network structure in a U shape. Adding a certain degree of time-related random Gaussian noise to the clear polarization images of the target, and obtaining a feature map vector from the turbid polarization images through a polarization information preprocessing module. Inputting the feature map and the clear images with noise into the network. First, passing through four encoders, each encoder being composed of two residual convolution modules, two self-attention mechanism modules and one down-sampling module. The encoder will reduce the size and increase the dimension of the feature map. Then, passing through four decoders, each decoder containing two residual convolution modules, two self-attention mechanism modules and one up-sampling module. The decoder will increase the size and reduce the dimension of the feature map. Each encoder and decoder is embedded with a time parameter associated with the degree of added noise, which is used to prompt the network to the degree of noise removal. The network outputs a predicted noise in the training stage, and the noise and the standard Gaussian noise are subjected to mean square error loss to optimize the network.

[0033] Step four, training the network, using the data set obtained in step two to train the conditional diffusion model constructed in step three to enable it to output the corresponding clear image according to the input turbid image and Gaussian noise image subject to standard normal distribution;

[0034] Step five, restoring the turbid polarized image, and finally obtaining a restored image with obvious visual effect improvement.

[0035] Further, the clear polarized image of the target object in step one: the light beam emitted by the light source 1 passes through the polarizer 2 and the beam expander 3 of the polarization modulation system, and then irradiates the target object 6, and after being reflected by the target object 6, reaches the split focal plane polarized camera 7, thereby obtaining the clear polarized image of the target object.

[0036] Further, the polarized image of the target object under different turbidity in step one: gradually add scattering medium (such as skimmed milk, etc.) in water, and use a 500 million pixel polarized gigabit Ethernet industrial camera to shoot the target object 6 in different turbidity environments, thereby obtaining the polarized image of the target object 6 in different polarization angles and different turbidity concentrations in the turbid underwater environment.

[0037] Further, a 532nm blue-green laser is used as the light source 1.

[0038] Further, the residual convolution module in the conditional diffusion neural network in step three is composed of three parts. The first part is a structure composed of a feature map small batch normalization operation, a Swish activation function and a 3*3 convolution layer; the second part is a module for time embedding, which is composed of a Swish activation function and a 1*1 convolution layer; the third part is composed of a feature batch normalization layer, a Swish activation function, a Dropout layer and a 3*3 convolution layer. This structure can solve the problem of gradient disappearance in the network training process, strengthen the expression ability of the model, and improve the trainability and stability of the network.

[0039] Further, the self-attention mechanism module in the conditional diffusion neural network in step three is composed of two parts. The first part is group normalization. Through group normalization, the input is standardized to improve the stability and efficiency of training, and to help reduce internal covariate shift in the network; the second part is a convolution projection layer, which is used to generate query (Q), key (K) and value (V) vectors. The query (Q) determines the focus of attention of the current element in the entire sequence; the key (K) provides the identification of each element, helping the query to find related elements; the value (V) carries the actual information content, and finally generates the output according to the matching degree of the query and the key.

[0040] Further, the conditional diffusion model described in step three is a model proposed based on the principle of generating images by diffusion model, which describes the change relationship between turbid images and clear images.

[0041] Specifically, step 1, using an underwater active imaging system, using linearly polarized light for active illumination, shooting clear polarization images of underwater target 6 and turbid polarization images of target 6 under different concentrations. In this embodiment, a 532nm blue-green laser is used as the light source 1, and a PMMA (polymethyl methacrylate) glass cylinder 4 is selected.

[0042] The light beam emitted by the light source 1 passes through the polarizer 2 of the polarization modulation system and the beam expander 3 in turn, and then irradiates the underwater target 6. After being reflected by the target 6, it reaches the split focal plane polarization camera 7, and the clear polarization image of the underwater target 6 is shot. Gradually add scattering medium (such as skimmed milk) in water, use 5 million pixel polarization gigabit Ethernet industrial camera to shoot polarization images of target 6 under different turbidity, in this example, 10 groups of images are shot, 20 kinds of different concentrations of milk are added in each group, which is used to simulate 20 kinds of different concentrations of underwater environment 5, and each concentration is shot a polarization image of target 6 (including polarization direction 0 degree, 45 degree, 90 degree, 135 degree) according to the shooting method in step 1;

[0043] Step 2, establish data set; the images obtained in step 1 are expanded by turning over, rotating and cutting, and then divided into training set, validation set and test set according to the ratio of 7:2:1.

[0044] Step 3, according to the idea of generating images by diffusion model, construct conditional diffusion neural network, the overall network structure is designed according to U type, the clear polarization image of the target is added with a certain degree of time related random Gaussian noise, the turbid polarization image is preprocessed by polarization information module to obtain feature map vector, and the feature map and the clear image with noise are jointly input into the network; first, pass through four encoders, each encoder is composed of two residual convolution modules, two self-attention mechanism modules and one down sampling module, the encoder reduces the size of the feature map and increases the dimension. Then pass through four decoders, the decoder includes two residual convolution modules, two self-attention mechanism modules and one up sampling module, the decoder increases the size of the feature map and reduces the dimension. Each encoder and decoder embeds a time parameter associated with the noise degree, which is used to prompt the network to the degree of denoising, the network outputs a predicted noise in the training stage, and the noise and the standard Gaussian noise are used to do mean square error loss to optimize the network.

[0045] Step 4, training the network, using the data set obtained in step two to train the conditional diffusion model constructed in step three to be able to output the corresponding clear image according to the input turbidity image and Gaussian noise image subject to standard normal distribution.

[0046] Step 5, restoring the turbid polarization image, and finally obtaining a restored image with obvious visual effect improvement.

[0047] It can be seen from the experimental results that the polarization image of high-concentration turbid underwater shooting can be effectively restored, and the image restoration effect is remarkable in combination with objective evaluation indexes SSIM (structural similarity) and PSNR (peak signal-to-noise ratio).

[0048] Although the present application has been described in detail with reference to the foregoing embodiments, those skilled in the art can modify the technical solutions recorded in the foregoing embodiments or make equivalent replacements for part of the technical features, and any modification, equivalent replacement, improvement, etc. made within the spirit and principles of the present application shall be included in the protection scope of the present application.

Claims

1. A polarization restoration imaging method based on a conditional diffusion model, characterized in that: The following steps are involved: Step 1: Capture a clear polarization image of the underwater target and polarization images corresponding to different turbidity levels; Step 2: Divide the captured polarization images into training set, validation set, and test set according to the ratio of 7:2:1; Step 3: Based on the idea of ​​generating images using the diffusion model, a conditional diffusion neural network is constructed. The overall network structure follows a U-shaped design. A certain degree of random Gaussian noise associated with time is added to the target clear polarization image. The turbid polarization image is processed through the polarization information preprocessing module to obtain a feature map vector. This feature map and the noisy clear image are jointly input into the network. First, it passes through four encoders, each of which consists of two residual convolution modules, two self-attention mechanism modules, and a downsampling module. Then, it passes through four decoders, each of which contains two residual convolution modules, two self-attention mechanism modules, and an upsampling module. Each encoder and decoder embeds a time parameter associated with the degree of added noise to indicate the degree of denoising required by the network. The network outputs a predicted noise during the training phase, and the predicted noise is compared with the standard Gaussian noise using a mean square error loss to optimize the network. Step 4: Train the network by using the data set obtained in step 2 to train the conditional diffusion model constructed in step 3 so that it can output a corresponding clear image based on the input turbid image and the Gaussian noise image that obeys the standard normal distribution; Step 5: Restore the turbid polarized image, and finally obtain a restored image with significantly improved visual effects.

2. The polarization restoration imaging method based on the conditional diffusion model according to claim 1, characterized in that: The clear polarization image of the underwater target captured in step 1: the light beam emitted by the light source (1) passes through the polarizer (2) and the beam expander (3) of the polarization modulation system in sequence and then irradiates the target (6), and then reaches the focal plane polarization camera (7) after being reflected by the target (6), thereby obtaining a clear polarization image of the target (6).

3. The polarization restoration imaging method based on the conditional diffusion model according to claim 2, characterized in that: In step 1, polarization images of the target object under different turbidity levels are captured: scattering media are gradually added to the water, and a 5-megapixel polarization Gigabit Ethernet industrial camera is used to capture the target object (6) under different turbidity concentration environments, thereby obtaining polarization images of the target object (6) in the turbid underwater environment at different polarization angles and different turbidity concentrations.

4. The polarization restoration imaging method based on the conditional diffusion model according to claim 2 or 3, characterized in that: A 532 nm blue-green laser is used as the light source (1).

5. The polarization restoration imaging method based on the conditional diffusion model according to claim 4, characterized in that: The residual convolution module in the conditional diffusion neural network described in step 3 consists of three parts. The first part is a structure consisting of a small batch normalization operation for the feature map, a Swish activation function and a 3*3 convolution layer; the second part is a module for time embedding, which consists of a Swish activation function and a 1*1 convolution layer; the third part consists of a feature batch normalization layer, a Swish activation function, a Dropout regularization layer and a 3*3 convolution layer.

6. The polarization restoration imaging method based on the conditional diffusion model according to claim 5, characterized in that: The self-attention mechanism module in the conditional diffusion neural network described in step 3 consists of two parts. The first part is group normalization, which standardizes the input to improve the stability and efficiency of training and help reduce internal covariate shift in the network. The second part is the convolutional projection layer, which is used to generate the query (Q), key (K), and value (V) vectors. The query (Q) determines the focus of the current element in the entire sequence; the key (K) provides an identifier for each element, helping the query find related elements. The value (V) carries the actual information content and is ultimately weighted to generate the output based on the degree of match between the query and the key.

7. The polarization restoration imaging method based on the conditional diffusion model according to claim 2 or 3, characterized in that: The target object (6) is placed in a glass tank (4) filled with water to simulate an underwater environment (5).

8. The polarization restoration imaging method based on the conditional diffusion model according to claim 7, characterized in that: The glass cylinder (4) is made of PMMA material.

Images