Underwater image enhancement method based on red dark channel prior constraint and diffusion model

Abstract

This invention discloses an underwater image enhancement method based on a diffusion model with red-dark channel prior constraints. The method includes inputting a single-frame underwater raw image into the model and passing it through a joint estimation module of red-dark channel prior and quadtree, outputting transmittance and background light. These are then encoded by a physical condition encoder to generate a physical condition vector containing spatial attenuation priors and global background light constraints. This vector is used by the diffusion network for encoding and utilizes transmittance distribution and global background light information. In the reverse denoising process, a physical consistency layer and the Jaffe-McGlamery equation are combined to construct a physical consistency residual. The residual gradient calculated based on the physical model is injected as a guiding term into the noise prediction process. The enhanced image is iteratively generated through a denoising diffusion implicit model sampling process. The enhanced image is then forward-degraded and re-rendered through a cyclic consistency reconstruction branch. The cyclic consistency loss is calculated and used for reverse optimization of network parameters. This invention can reduce the color shift problems caused by background light misjudgment and transmittance overestimation.

Description

Technical Field

[0001] This invention relates to an underwater image enhancement method based on a diffusion model constrained by the Red Dark Channel Prior (RDCP). Background Technology

[0002] The main challenges in underwater image enhancement stem from the complexity of underwater light propagation mechanisms and the insufficient physical constraints of existing data-driven methods. Because water selectively absorbs different wavelengths of light, red light attenuates the fastest during propagation, and the degree of attenuation varies with scene depth, water turbidity, and suspended particle concentration. Simultaneously, forward and backscattering caused by suspended particles in the water can lead to decreased image contrast, blurred edges, localized fogging, and bright scattered points. This can easily cause artificial targets or localized bright spots to be misidentified as infinity-based background areas, resulting in incorrect background light selection and affecting subsequent defogging and enhancement results. Furthermore, transmittance t and background light A are usually difficult to measure directly in real underwater scenes. Existing methods often rely on empirical assumptions or indirect estimations, which can easily lead to the accumulation of errors introduced by prior biases.

[0003] Diffusion models possess strong capabilities in high-dimensional distribution modeling and image restoration, gradually reconstructing target images through forward noise addition and backward denoising, thus demonstrating significant application potential in image enhancement. However, existing diffusion models still suffer from the following problems when applied to underwater image enhancement: First, diffusion models typically rely on multi-step iterative sampling, resulting in lengthy inference processes that struggle to balance enhancement effectiveness with processing efficiency. Artificially reducing the number of sampling steps can easily lead to insufficient color restoration, color cast reversal, or detail distortion. Second, existing methods often learn the mapping relationship between degraded and enhanced images using a purely data-driven approach, lacking physical constraints consistent with underwater imaging mechanisms. This results in visually improved images that may exhibit color distortion, brightness drift, or structural artifacts that violate the Jaffe-McGlamery underwater imaging principles, thus limiting their application in real-world, complex underwater scenarios.

[0004] Dark Channel Prior (DCP), a classic image dehazing method, estimates transmittance and background light based on the statistical law that at least one color channel intensity approaches zero in a local neighborhood of a naturally clear image. While this method performs well in atmospheric dehazing scenarios, it has significant limitations when directly applied to underwater environments. Because red light attenuates significantly in water, traditional DCP often misjudges the medium attenuation of the red channel as a low-reflectivity feature of the scene itself, leading to an overestimation of transmittance. Furthermore, traditional methods typically use the brightest pixels in the image to statistically determine background light, which is easily affected by white suspended particles, localized strong reflections, or artificially bright targets, causing misjudgments of background light and further exacerbating problems such as overly bright distant scenes, saturated foreground scenes, and overall graying. To improve the applicability of traditional DCP in underwater scenes, research has proposed Red Dark Channel Prior (RDCP), which first inverts the red channel and then constructs the dark channel response together with the green and blue channels. This method utilizes the characteristic that the stronger the red light attenuation, the greater the response after reversal, making the transmittance estimation more consistent with the underwater medium propagation law, and can alleviate the color shift problem of traditional dark channel priors to some extent. Nevertheless, existing methods based on dark channel priors are still easily affected by complex aquatic environments and the distribution of suspended particles, resulting in unstable transmittance estimation and large errors in background light selection.

[0005] On the other hand, existing underwater conditional diffusion models typically only stitch transmittance into the input channel or completely discard key physical parameters such as background light, using only random latent variables or monocular depth maps as conditional inputs. These methods lack explicit physical consistency constraints during the backsampling process, causing the noise prediction process to be disconnected from the underwater imaging mechanism. This easily leads to insufficient physical consistency, unstable enhancement results, and high model parameter count and inference costs. Therefore, how to effectively integrate physical priors such as transmittance and background light into the denoising process while maintaining the diffusion model's generative capabilities, and how to introduce differentiable physical constraints in each backsampling step, remains a technical problem that needs to be solved in this field. Summary of the Invention

[0006] To address the problems of large transmittance estimation errors, easy misjudgment of background light, and lack of physical consistency in diffusion enhancement results in existing technologies, this invention proposes an underwater image enhancement method based on a diffusion model under prior constraints of the red and dark channels. This method estimates transmittance and background light using a single-frame underwater image and explicitly introduces these as physical conditions into the inverse denoising process of the diffusion model. This improves the consistency between the enhancement results and the underwater imaging mechanism while addressing issues such as color cast, contrast reduction, and detail blurring.

[0007] This invention can be achieved through the following technical solutions: The underwater image enhancement method based on the diffusion model under red and dark channel prior constraints includes the following steps: 1) Input a single frame of raw underwater image to the model, wherein the raw image is a three-channel RGB image; 2) Input the original image into the red and dark channel prior-quadtree joint estimation module, and output the transmittance and background light; 3) The physical condition encoder includes a spatial branch, a global branch, and a fusion unit. The spatial branch takes the prior transmittance as input and extracts spatial attenuation features that reflect the transmittance distribution per pixel. The global branch takes the background light parameter A as input and extracts global constraint features that reflect the overall lighting conditions of the scene. The fusion unit maps the spatial attenuation features and global constraint features to the same channel dimension through a projection layer, then splices and fuses them to generate a physical condition vector, so that the subsequent diffusion network can simultaneously read pixel transmittance information and background light statistics of the entire image. 4) The latent noise variables and time steps are embedded and the above physical condition vectors are input into the conditional diffusion denoising network. The predicted noise is output during the reverse denoising process. The imaging residual is constructed according to the underwater imaging physical model through the physical consistency layer. The gradient correction term corresponding to the residual is injected into the noise prediction value to form a physical guide for the denoising trajectory. 5) After sampling by the N-step denoising diffusion implicit model, the enhanced image is obtained. At the same time, the enhanced image, transmittance and background light parameters are input into the cyclic consistency reconstruction branch. The re-rendered degraded image is generated according to the forward underwater imaging model, and the cyclic loss between the original image and the re-rendered degraded image is calculated for backpropagation to optimize the network parameters. 6) Output imaging results.

[0008] Furthermore, in step 2), the red-dark channel prior is applied to each pixel of the input image's three RGB channels. Perform red channel inversion and minimum filtering within a local neighborhood to construct the red-dark channel response: The corresponding red and dark channel response estimates the transmittance. The more red light is absorbed by water, the larger 1-R becomes, the smaller the RDCP value becomes, and t decreases accordingly, thus automatically compensating for red light attenuation. A quadtree hierarchical search strategy is used to recursively divide the input image into sub-blocks, calculate the joint score of the brightness mean, color saturation distribution, and gradient sparsity of each sub-block, select the sub-block with the highest score to continue dividing, until the sub-block size is no larger than the preset threshold of 64×64, and select the top 0.1% of pixels with the largest red and dark channels in the final candidate block, and take its RGB mean in the original image as the background light parameter A, so that the background light parameter preferentially comes from the water body at infinity rather than white suspended objects, strong reflective areas, or artificial target areas.

[0009] Furthermore, in the diffusion model inference stage of step 4), at the first... In the backsampling of the Denoising Diffusion Implicit Model (DDIM), the conditional diffusion denoising network adjusts the current latent variable... Time step T and physical condition vector Calculate the direction of natural denoising The physical consistency layer will handle the current latent variables. Restore to predicted brightness map And construct the physical consistency residuals based on the underwater imaging equation. Construct the physical consistency loss based on the physical consistency residuals, and calculate the loss for the current latent variables. gradient The gradient is adjusted according to the physical guidance coefficient. The natural denoising direction is corrected to obtain the physically corrected noise prediction direction. ;Use the corrected noise prediction direction to perform the DDIM update for the current step, and obtain the latent variables for the next step. This results in a reverse denoising process that combines conditional guidance and physical guidance constraints.

[0010] Furthermore, the physical condition encoder in step 3) includes a spatial branch, a global branch, and a fusion unit; the spatial branch uses single-channel transmittance. As input, local spatial decay priors are extracted through convolutional layers, normalization layers, and nonlinear activation layers, and the output is a spatial prior feature map. Global branch based on background light value As input, the global condition vector is obtained through at least two layers of linear mapping. Spatial feature map With global condition vector After being mapped to the same channel dimension by a projection layer, the vectors are concatenated to generate a physical condition vector to guide the diffusion denoising process. .

[0011] Furthermore, the spatial branch uses a convolution with a stride of 4 to downsample and encode the transmittance, outputting an 8-channel spatial prior feature map; the global branch maps the 3D background light vector to a 128-dimensional global conditional vector; the input to the conditional diffusion denoising network is the noisy image. The conditional image and spatial prior feature map are concatenated in the channel dimension; the global conditional vector is applied to multiple residual blocks of the conditional diffusion denoising network in FiLM modulation to generate scaling and bias parameters for modulating network features; an attention module is introduced into the intermediate layer of the conditional diffusion denoising network to enhance the long-range dependency modeling capability of mid-to-high-level features.

[0012] Further, the cyclic consistency reconstruction branch in step 5) includes: generating a re-rendered degraded image according to the estimated enhanced image, transmittance, and background light parameter A, based on the forward underwater imaging model; and calculating the original input image. With re-rendering degraded images Cyclic consistency loss The transmittance parameter, background light parameter, and at least some network weights in the conditional diffusion denoising network are jointly fine-tuned based on the cycle consistency loss to achieve self-supervised optimization in scenarios where a reference ground truth image is lacking.

[0013] Furthermore, the model training employs a multi-objective joint loss function, with the total loss consisting of a weighted combination of diffusion loss, physical consistency loss, and cyclic consistency loss, to simultaneously constrain noise prediction accuracy, imaging mechanism consistency, and enhancement-re-rendering closed-loop consistency.

[0014] Furthermore, the physical guiding coefficient The value range is 0.05 to 0.2, and the number of DDIM sampling steps is a preset number of steps N, in order to balance the enhancement effect and inference speed.

[0015] The present invention has the following beneficial effects: 1) Introducing transmittance and background light as interpretable physical priors. Red light attenuates the fastest underwater, causing traditional dark channel priors to misjudge the dark red channel as a dark scene, resulting in overestimation of transmittance, overly bright distant scenes, and saturated near scenes. To address this, a method is proposed to first invert the 1-R channel before taking the dark channel. The more severe the red light attenuation, the brighter the scene after inversion, while the dark channel value decreases. This breaks the causal chain of the dark channel prior misjudging the dark red channel as a dark scene, allowing transmittance to return from systematic overestimation to the true attenuation amount. This eliminates the root cause of the color bias in the dark channel prior, such as bright distant scenes and saturated near scenes, and global purple-gray color. Furthermore, automatic weighting through quadtree joint scoring avoids interference from white suspended objects, artificial targets, etc., which could lead to incorrect selection of background light, thereby improving the accuracy of background light positioning in the infinite water area. 2) This invention does not simply concatenate conditional information as input, but encodes and fuses spatial priors and global priors separately and injects them into the conditional diffusion denoising network. At the same time, physical consistency constraints based on imaging equation residuals are introduced in each step of the reverse denoising process, so that the denoising trajectory of the diffusion model further satisfies the underwater imaging mechanism in addition to data-driven, thereby improving the physical rationality of the enhancement results, color recovery stability and structure preservation ability.

[0016] 3) To address the issue that real underwater images often lack strict reference values, this invention sets up a cyclic consistency reconstruction branch, constructs closed-loop constraints by enhancing the image and re-rendering the degraded image, and uses the error between the original input image and the re-rendered image for online backpropagation optimization, thereby improving the model's adaptability and cross-scene robustness in no-reference or weak-reference scenarios. Attached Figure Description

[0017] Figure 1 This is a flowchart of the present invention; Figure 2 This is a schematic diagram of the red-dark channel prior-quadtree joint estimation module of the present invention; Figure 3 This is a schematic diagram of the condition encoder of the present invention. Detailed Implementation

[0018] The following specific embodiments illustrate the implementation of the present invention. Those skilled in the art can easily understand other advantages and effects of the present invention from the content disclosed in this specification.

[0019] like Figure 1 As shown, an underwater image enhancement method based on a diffusion model under prior constraints of red and dark channels according to the present invention includes the following steps: Input a single frame of raw underwater image to the model; Red-dark channel prior-quadritree joint estimation module (e.g.) Figure 2 (As shown) Output transmittance and background light. Red and dark channel priors for each pixel of the input image's three RGB channels. Perform red channel inversion and minimum filtering within a local neighborhood to construct the red-dark channel response: The corresponding red and dark channel response estimates the transmittance. The more red light is absorbed by water, the larger 1-R becomes and the smaller the RDCP value becomes, thus automatically compensating for red light attenuation. A quadtree hierarchical search strategy is used to recursively divide the input image into sub-blocks. The joint score of the brightness mean, color saturation distribution, and gradient sparsity of each sub-block is calculated. The sub-block with the highest score is selected and further divided until the sub-block size is no larger than the preset threshold of 64×64. The top 0.1% of pixels with the largest red and dark channels in the final candidate block are selected, and their RGB mean in the original image is taken as the background light parameter A. This ensures that the background light parameter comes from the water body at infinity rather than white suspended objects, strongly reflective areas, or artificial target areas. The physical condition encoder includes a spatial branch, a global branch, and a fusion unit. The spatial branch takes prior transmittance as input and extracts spatial attenuation features reflecting the transmittance distribution per pixel. The global branch takes background light parameter A as input and extracts global constraint features reflecting the overall lighting conditions of the scene. The fusion unit maps the spatial attenuation features and global constraint features to the same channel dimension through a projection layer, then concatenates and fuses them to generate a physical condition vector, enabling the subsequent diffusion network to simultaneously read pixel transmittance information and background light statistics of the entire image. Figure 3 The diagram shows a schematic of the conditional encoder of the present invention. The physical conditional encoder includes a spatial branch, a global branch, and a fusion unit. The spatial branch uses a convolution with a stride of 4 to downsample and encode the transmittance to extract local spatial attenuation priors reflecting the pixel-by-pixel transmittance distribution, outputting an 8-channel spatial prior feature map. The global branch takes a 3D background light vector composed of background light parameters A as input and encodes it into a 128-dimensional global conditional vector through a linear mapping method to extract global prior constraint information reflecting the overall lighting conditions of the scene. The fusion unit is used to map the spatial prior feature map and the global conditional vector to the same channel dimension through a projection layer and then concatenate them to generate a joint physical conditional embedding for use by the conditional diffusion denoising network.

[0020] In each iteration, the conditional diffusion denoising network receives a conditional vector and outputs a noise prediction value. The physical consistency layer (physical constraint gradient injection layer) adds the imaging law residual back to the noise prediction value in the form of a gradient. In the T-th step of the denoising diffusion implicit model (DDIM) backsampling, the conditional diffusion denoising network adjusts the noise prediction value according to the current latent variables. Time step T and physical condition vector Calculate the direction of natural denoising The physical consistency layer will handle the current latent variables. Restore to predicted brightness map And construct the physical consistency residuals based on the underwater imaging equation. Construct the physical consistency loss based on the physical consistency residuals, and calculate the loss for the current latent variables. gradient The gradient is adjusted according to the physical guidance coefficient. The natural denoising direction is corrected to obtain the physically corrected noise prediction direction. ;Use the corrected noise prediction direction to perform the DDIM update for the current step, and obtain the latent variables for the next step. This forms a reverse denoising process that combines conditional guidance and physical guidance. After N sampling steps, an enhanced image is obtained. Simultaneously, the cycle consistency reconstruction branch uses the estimated... and A degraded image is resynthesized using a forward imaging model. and calculate With the original input The losses between them are used to form a physical closed-loop constraint; Output the imaging results.

[0021] Experimental results: On the UIEB benchmark dataset, the proposed method is compared with two representative diffusion enhancement methods: UM-DDPM and DiffWater. The results are shown in the table above. On the U90 subset with a reference image, the proposed method achieves a PSNR of 21.05 dB and an SSIM of 0.898, which are improvements of +0.26 dB and +0.098 respectively compared to UM-DDPM, and also outperform DiffWater (PSNR +0.08 dB, SSIM +0.003), indicating that it further reduces pixel-level reconstruction errors while maintaining structural consistency. It is worth noting that DiffWater already shows strong structural fidelity on U90 (SSIM=0.895), while the proposed method still achieves additional gains, indicating that physical prior constraints can effectively suppress oversmoothing and local texture drift during diffusion inversion, making the prediction closer to the distribution of the reference image. On the U45 subset without reference evaluation, our method achieves UIQM=4.92 and UCIQE=0.474, representing improvements of +0.02 and +0.029 respectively compared to UM-DDPM (UIQM=4.90, UCIQE=0.445), and also surpasses DiffWater (UIQM=4.73, UCIQE=0.462), with UIQM improvement of +0.19 and UCIQE improvement of +0.012. This indicates that our method not only improves color and contrast but also achieves more stable gains in sharpness, acuity, and overall visual quality.

[0022] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the protection scope of the present invention.

Claims

1. An underwater image enhancement method based on a diffusion model under prior constraints of red and dark channels, characterized in that, Includes the following steps: 1) Acquire a single frame of raw underwater image, wherein the raw image is a three-channel RGB image; 2) Input the original image into the red-dark channel prior-quadtree joint estimation module, and output the transmittance and background light; 3) The physical condition encoder includes a spatial branch, a global branch, and a fusion unit. The spatial branch takes the prior transmittance as input and extracts spatial attenuation features that reflect the transmittance distribution per pixel. The global branch takes the background light parameter A as input and extracts global constraint features that reflect the overall lighting conditions of the scene. The fusion unit maps the spatial attenuation features and global constraint features to the same channel dimension through a projection layer, then splices and fuses them to generate a physical condition vector, so that the subsequent diffusion network can simultaneously read pixel transmittance information and background light statistics of the entire image. 4) The latent noise variables and time steps are embedded and the above physical condition vectors are input into the conditional diffusion denoising network. The predicted noise is output during the reverse denoising process. The imaging residual is constructed according to the underwater imaging physical model through the physical consistency layer. The gradient correction term corresponding to the residual is injected into the noise prediction value to form a physical guide for the denoising trajectory. 5) After sampling by the N-step denoising diffusion implicit model, the enhanced image is obtained. At the same time, the enhanced image, transmittance and background light parameters are input into the cyclic consistency reconstruction branch. The re-rendered degraded image is generated according to the forward underwater imaging model, and the cyclic loss between the original image and the re-rendered degraded image is calculated for backpropagation to optimize the network parameters. 6) Output the enhanced underwater image.

2. The underwater image enhancement method based on a diffusion model under prior constraints of red and dark channels according to claim 1, characterized in that, In step 2), the red and dark channels prior are applied to each pixel of the input image's three color channels (RGB). Perform red channel inversion and minimum filtering within a local neighborhood to construct the red-dark channel response: The corresponding red and dark channel response estimates the transmittance. The more red light is absorbed by water, the larger 1-R becomes and the smaller the RDCP value becomes, and t decreases accordingly, thus automatically compensating for red light attenuation. A quadtree hierarchical search strategy is used to recursively divide the input image into sub-blocks, calculate the joint score of the brightness mean, color saturation distribution, and gradient sparsity of each sub-block, select the sub-block with the highest score to continue dividing until the sub-block size is no larger than the preset threshold of 64×64, and select the top 0.1% of pixels with the largest red and dark channels in the final candidate block, and take its RGB mean in the original image as the background light parameter A, so that the background light parameter preferentially comes from the water body at infinity rather than white suspended objects, strong reflective areas, or artificial target areas.

3. The underwater image enhancement method based on a diffusion model under prior constraints of red and dark channels according to claim 1, characterized in that, In step 4), during the diffusion model inference stage, in the... In the backsampling of the Denoising Diffusion Implicit Model (DDIM), the conditional diffusion denoising network adjusts the current latent variable... Time step T and physical condition vector Calculate the direction of natural denoising The physical consistency layer will handle the current latent variables. Restore to predicted brightness map And construct the physical consistency residuals based on the underwater imaging equation. Construct the physical consistency loss based on the physical consistency residuals, and calculate the loss for the current latent variables. gradient ; The gradient is adjusted according to the physical guidance coefficient. The natural denoising direction is corrected to obtain the physically corrected noise prediction direction. ;Use the corrected noise prediction direction to perform the DDIM update for the current step, and obtain the latent variables for the next step. This results in a reverse denoising process that combines conditional guidance and physical guidance constraints.

4. The underwater image enhancement method based on a diffusion model under prior constraints of red and dark channels according to claim 1, characterized in that, The physical condition encoder in step 3) includes a spatial branch, a global branch, and a fusion unit; the spatial branch uses single-channel transmittance. As input, local spatial decay priors are extracted through convolutional layers, normalization layers, and nonlinear activation layers, and the output is a spatial prior feature map. Global branches with background light As input, the global condition vector is obtained through at least two layers of linear mapping. Spatial feature map With global condition vector After being mapped to the same channel dimension by a projection layer, the vectors are concatenated to generate a physical condition vector to guide the diffusion denoising process. .

5. The underwater image enhancement method based on a diffusion model under prior constraints of red and dark channels according to claim 4, characterized in that, The spatial branch uses a convolution with a stride of 4 to downsample and encode the transmittance, outputting an 8-channel spatial prior feature map; the global branch maps the 3D background light vector to a 128-dimensional global conditional vector; the input to the conditional diffusion denoising network is the noisy image. The conditional image and spatial prior feature map are concatenated in the channel dimension; the global conditional vector is applied to multiple residual blocks of the conditional diffusion denoising network using Feature-wise Linear Modulation (FiLM) to generate scaling and bias parameters for modulating network features; an attention module is introduced into the intermediate layer of the conditional diffusion denoising network to enhance the long-range dependency modeling capability of mid-to-high-level features.

6. The underwater image enhancement method based on a diffusion model under prior constraints of red and dark channels according to claim 1, characterized in that, The cyclic consistency reconstruction branch in step 5) includes: generating a re-rendered degraded image according to the estimated enhanced image, transmittance, and background light parameter A, based on the forward underwater imaging model; and calculating the original input image. With re-rendering degraded images Cyclic consistency loss The transmittance parameter, background light parameter, and at least some network weights in the conditional diffusion denoising network are jointly fine-tuned based on the cycle consistency loss to achieve self-supervised optimization in scenarios where a reference ground truth image is lacking.

7. The underwater image enhancement method based on a diffusion model under prior constraints of red and dark channels according to claim 3 or 6, characterized in that, The model training employs a multi-objective joint loss function, with the total loss consisting of a weighted combination of diffusion loss, physical consistency loss, and cyclic consistency loss, to simultaneously constrain noise prediction accuracy, imaging mechanism consistency, and enhanced re-rendering closed-loop consistency.

8. The underwater image enhancement method based on a diffusion model under prior constraints of red and dark channels according to claim 1, characterized in that, The physical guiding coefficient The value range is 0.05 to 0.2, and the number of DDIM sampling steps is a preset number of steps N, in order to balance the enhancement effect and inference speed.

Images