High dynamic range image tone mapping method based on diffusion model and zero-shot training

Abstract

This invention belongs to the field of image processing technology, specifically relating to a high dynamic range (HDR) image tone mapping method based on a diffusion model and zero-shot training. The invention includes constructing a tone mapping network model and integrating it into the framework of a diffusion model. Structural information and a rough brightness distribution map are used as conditional inputs to the diffusion model. The structural information of the original image is used to correct the intermediate results of the diffusion model's backsampling. The network model includes a generator model and its control module. The generator model includes cascaded encoding modules at four scales, an intermediate module, and cascaded decoding modules at four scales. The control module includes the same four-scale encoding modules as the SD model, with zero-convolutional layers added at the end of each scale. The output of the control module at each scale is added to the feature map output of the generator model at the corresponding scale, thereby controlling the final output of the generator model. This invention can effectively preserve the original image information and improve the tone mapping effect of HDR images.

Description

Technical Field

[0001] This invention belongs to the field of image processing technology, specifically relating to a high dynamic range image tone mapping method based on diffusion models and zero-shot training. Background Technology

[0002] High Dynamic Range (HDR) images have a wider bit width than ordinary images, thus recording more information, especially advantageous under extreme lighting conditions. However, HDR images cannot be directly displayed on ordinary display devices, so tone mapping is required. This maps the pixel values ​​of the HDR image to the displayable Low Dynamic Range (LDR) range, preserving as much of the rich information recorded in the original image as possible while generating normal tones.

[0003] However, it is theoretically impossible to obtain paired HDR-high-quality LDR datasets. Existing state-of-the-art methods [1][2] mostly utilize generative adversarial networks to address this problem, and the loss function used includes a structure preservation term and a style naturalness term. However, the optimization directions of these two objectives may be different, resulting in results that neither preserve the original image information well nor have natural tones. Some other works [3] train models on datasets for low-light image enhancement, but they do not perform well for HDR image tone mapping tasks. Summary of the Invention

[0004] The purpose of this invention is to provide a high dynamic range image tone mapping method based on a diffusion model and zero-shot training, so as to preserve the original image information well and improve the tone mapping effect of HDR images.

[0005] The high dynamic range image tone mapping method based on a diffusion model and zero-shot training provided by this invention includes constructing a tone mapping network model and integrating it into the framework of a diffusion model, using structural information and a rough brightness distribution map as conditional inputs to the diffusion model, and using the structural information of the original image to correct the intermediate results of backsampling in the diffusion model. The flowchart of this invention is described below. Figure 1 The specific steps are as follows:

[0006] (1) Construct a tone mapping network model for each step of the diffusion model. This network model structure includes a pre-trained generative model and a control module for that model, such as... Figure 2As shown. The pre-trained generative model adopts the network structure of the Stable Diffusion (SD model) [4], which mainly includes: a cascaded four-scale encoding module, an intermediate module, and a cascaded four-scale decoding module; the structure of the control module includes the same four-scale encoding module as the SD model, but a zero convolutional layer is added at the end of the neural network at each scale. The output of the control module at each scale is added to the feature map of the corresponding scale output of the generative model, thereby controlling the final output result of the generative model. The input of the control module includes the structural information and the approximate brightness distribution map at each pixel of the image. Wherein:

[0007] The formula for calculating the structural information at each pixel is:

[0008]

[0009] In the formula, I ori Represents the original image, μ ori and σ ori Let represent the local Gaussian mean and local standard deviation of the image, respectively; ε is a small value introduced to prevent division by zero; and (i,j) are the coordinates of the current pixel in the image; μ ori and σ ori The specific calculation formula is as follows:

[0010]

[0011] Where K represents the window size of the Gaussian filter, and ω represents the Gaussian filter kernel.

[0012] The calculation method for the approximate brightness distribution map differs between the training and testing phases. During the training phase, the calculation method involves mean filtering of the brightness channel of the image; during the testing phase, the image is preprocessed first, specifically using global tone mapping [2], and then mean filtering is performed. The calculation formula for the preprocessing method is as follows:

[0013]

[0014] Where log represents the logarithmic function, max is the operation of taking the maximum value, Y represents the brightness channel of the input image, ε is a small value set to prevent division by zero, and the parameter λ is obtained by solving the optimization problem of equation (5):

[0015]

[0016] Among them, H l (Y c ) represents Y c Histogram, H l(LDR) is the average histogram of the brightness channels of a large number of high-quality natural LDR images.

[0017] During the training phase, the network structure described above is integrated into the Denoising Diffusion Probabilistic Models (DDPM)[5]; during the testing phase, the network structure described above is integrated into the Denoising Diffusion Implicit Models (DDIM)[6]. In both diffusion models, the network designed in this invention is used to predict the noise added at each step of the diffusion model.

[0018] (2) Based on the network model described in step (1), the structural information is input into the first scale encoding module of the control module, and the brightness approximate distribution map is input into the second scale encoding module of the control module, which can further enhance the detail information of the output image.

[0019] (3) Based on DDIM mentioned in step (1), the intermediate results generated in each iteration or part of the iterations of its backsampling iteration process are corrected, that is: the hue information of the generated result is taken and combined with the structural information of the original image to further protect the structural information of the original image. The corrected result is then used for subsequent calculations of the backsampling process according to the formula of the diffusion model.

[0020] Specifically, for any iteration of DDIM backsampling, the initial estimate of the noiseless data z0 is first calculated using the formula of the original DDIM algorithm, which is:

[0021]

[0022] in, This is the initial estimate of the latent space embedding z0 of the noiseless data obtained in this iteration of DDIM. t This is the result of the previous iteration. For DDIM predefined weights, ∈ θ (z t ,t,c) is the network model described in step (1), where t is the current iteration time step and c is the conditional input of the network model.

[0023] Then, using the structural information of the original image... Make corrections. The correction method is as follows:

[0024]

[0025] in, It is The operation of transferring from the latent space of the SD model to the image domain, TSD represents the tonal-structural decomposition of the image, that is, calculating the local Gaussian mean, local variance and structural information of the image according to formulas (1) to (3) in step (1). γ is a hyperparameter that controls the degree of image detail enhancement. The structural information of the original image is given, and ε is the operation that transforms the image into the latent space of the SD model; thus, we obtain... Then, this is used as the final estimate of z0 for this iteration. Following the principles of the original DDIM algorithm, subsequent calculations for this iteration are performed to obtain the result of this iteration. The calculation formula for the subsequent calculations in this iteration is:

[0026]

[0027] In the formula These are also weights predefined by DDIM.

[0028] (4) Network Training and Testing. The network was trained on a high-quality LDR image dataset. During training, the parameters of the generative model were frozen, and only the control module was trained. The loss function used was the same as that used in [4], [5] and other literature, namely: the L2 loss of the noise predicted by the network and the noise actually added to the noise-free image. After training, the model can be directly used for tone mapping of HDR images or other similar tasks.

[0029] During testing, only the luminance component of the image generated by the network was taken, and then converted into an RGB image using the following formula:

[0030]

[0031] Among them, Y pred Y and are the brightness components of the network-generated image and the original image, respectively, i∈{R,G,B}. and These represent one of the three RGB channels of the input image and the final output result, respectively, and s is a hyperparameter.

[0032] The main technical features and performance advantages of this invention are:

[0033] The network can be trained using only LDR images, and the resulting network model can be directly used for tone mapping of HDR images or other similar tasks without needing to be retrained on the image category to be processed. This means that the original image information can be well preserved and the tone mapping effect of HDR images can be improved. Attached Figure Description

[0034] Figure 1 This is a flowchart illustrating the high dynamic range image tone mapping method based on diffusion model and zero-shot training according to the present invention.

[0035] Figure 2 A schematic diagram of the tone mapping network model designed for this invention.

[0036] Figure 3 Here is a set of processing examples, where from left to right are HDR images after linear scaling of pixel values, the results output by references [3][1][2] and the results output by this invention. Detailed Implementation

[0037] The present invention will be further described below with reference to the embodiments and accompanying drawings.

[0038] The high dynamic range image tone mapping method based on diffusion model and zero-shot training provided by this invention includes the following steps:

[0039] Step 1: Construct as follows Figure 2 The network structure is shown. TSD (Tone-Structure Decomposition) represents decomposing an image into structural and tonal information, while TSR (Tone-Structure Reconstruction) represents reconstructing the structural and tonal information into an image. The pre-trained generative model uses version v1.5 of the SD model, which includes cascaded encoding modules at four scales, an intermediate module, and cascaded decoding modules at four scales. The control module includes the same four-scale encoding modules as the SD model, but adds a zero convolutional layer at the end of the neural network at each scale. The output of each scale is added to the feature map output of the corresponding scale of the generative model.

[0040] The networks described above are integrated into the framework of the diffusion model. The backsampling algorithm of the diffusion model adopts DDIM[6], and the default number of backsampling steps is set to 20 steps (i.e., 20 iterations).

[0041] Step 2: The network's conditional inputs include structural information and a general brightness distribution map. The structural information is calculated using equations (1) to (3), where ω uses a 7×7 Gaussian filter kernel with a standard deviation of 7 / 6, and ε is set to 0.0000001. The calculation method for the general brightness distribution map differs between the training and testing phases. During the training phase, the calculation method involves mean filtering of the image's brightness channel; during the testing phase, equations (4) to (5) are used to obtain the initial tone mapping result for the brightness channel, and then mean filtering is applied to this result. The hyperparameter λ in equation (4) is set to 0.05 by default. The size of the mean filter kernel can be appropriately set to a larger value, such as 0.02 times the smaller of the original image's length and width. The structural information is input into the first layer of the first scale of the control module, and the general brightness distribution map is input into the first layer of the second scale of the control module.

[0042] Step 3: Select certain iterations in the DDIM backsampling process and correct the iteration results using equations (6) to (9). Under the default configuration, correct the results calculated for the first 10 iterations and the last iteration. The hyperparameter γ in equation (7) is set to 2 to enhance the texture of the image.

[0043] Step 4: Using the network structure described in Step 1 and the input described in Step 2, freeze the parameters of the SD model and train only the parameters of the control module. The training set can be a high-quality LDR image dataset, such as DIV2K[7]. During testing, only the luminance component of the image generated by the network is taken, and then converted into an RGB image using Equation (8), where the hyperparameter s is set to 0.5.

[0044] A set of examples of tone mapping of HDR images using this invention, and comparisons with results from other algorithms. Figure 3 As shown, compared with other tone mapping methods, the images output by this invention contain richer information in overly bright and overly dark areas of the scene (marked in the figures and enlarged in the lower right corner of each image), demonstrating a significant advantage. Therefore, this invention is more suitable for tone mapping of HDR images under extreme lighting conditions.

[0045] References

[0046] [1].Aakanksha Rana, Praveer Singh, Giuseppe Valenzise, ​​Frederic Dufaux, Nikos Komodakis, and Aljosa Smolic. Deep tone mapping operator for high dynamic range images. IEEE Transactions on Image Processing, 29:1285–1298, 2020.

[0047] [2].Yael Vinker,Inbar Huberman-Spiegelglas,and Raanan Fattal.Unpairedlearning for high dynamic range image tone mapping.In 2021IEEE / CVFInternational Conference on Computer Vision(ICCV),pages 14637–14646,2021.

[0048] [3].Kai-Fu Yang,Cheng Cheng,Shi-Xuan Zhao,Hong-Mei Yan,Xian-ShiZhang,and Yong-Jie Li.Learning to adapt to light.Int.J.Comput.Vision,131(4):1022–1041,2023

[0049] [4].R.Rombach,A.Blattmann,D.Lorenz,P.Esser,and B.Ommer,“High-resolution image synthesis with latent diffusion models,”in Proceedings ofthe IEEE / CVF Conference on Computer Vision and Pattern Recognition(CVPR),June2022,pp.10 684–10 695

[0050] [5].J.Ho,A.Jain,and P.Abbeel,“Denoising diffusion probabilisticmodels,”in Advances in Neural Information Processing Systems,H.Larochelle,M.Ranzato,R.Hadsell,M.Balcan,and H.Lin,Eds.,vol.33.Curran Associates,Inc.,2020,pp.6840–6851.[Online].Available:

[0051] https: / / proceedings.neurips.cc / paperfiles / paper / 2020 / file / 4c5bcfec8584af0d967f1ab10179ca4b-Paper.pdf

[0052] [6].J.Song,C.Meng,and S.Ermon,“Denoising diffusion implicit models,”arXiv preprint arXiv:2010.02502,2022

[0053] [7].E.Agustsson and R.Timofte,"NTIRE 2017Challenge on Single ImageSuper-Resolution:Dataset and Study,"2017IEEE Conference on Computer Visionand Pattern Recognition Workshops(CVPRW),Honolulu,HI,USA,2017,pp.1122-1131。

Claims

1. A high dynamic range image tone mapping method based on a diffusion model and zero-shot training, characterized in that, This involves constructing a tone mapping network model and integrating it into the diffusion model framework. Structural information and a rough brightness distribution map are used as inputs to the diffusion model. The structural information of the original image is used to correct the intermediate results of the diffusion model's backsampling. The specific steps are as follows: (1) Construct a tone mapping network model for each step in the diffusion model; the network model structure includes a generative model for pre-training and a control module for the model; The generative model employs a stable diffusion (SD) network structure, comprising: cascaded encoding modules at four scales, an intermediate module, and cascaded decoding modules at four scales. The control module has the same four-scale encoding modules as the SD model, but adds a zero-convolutional layer at the end of each scale's neural network. The output of the control module at each scale is added to the feature map output of the corresponding scale of the generative model, thereby controlling the final output of the generative model. The input to the control module includes structural information and a rough brightness distribution map for each pixel of the input image. The formula for calculating the structural information at each pixel is: In the formula, I ori Represents the original image, μ ori and σ ori Let represent the local Gaussian mean and local standard deviation of the image, respectively; ε is a small value introduced to prevent division by zero; and (i,j) are the coordinates of the current pixel in the image; μ ori and σ ori The specific calculation formula is as follows: Where K represents the window size of the Gaussian filter, and ω represents the Gaussian filter kernel; The method for calculating the approximate brightness distribution map is as follows: During the training phase, mean filtering is applied to the brightness channel of the image; during the testing phase, the image is preprocessed first, specifically using global tone mapping, and then mean filtering is applied. The preprocessing calculation formula is as follows: Where log represents the logarithmic function, max is the operation of taking the maximum value, Y represents the brightness channel of the input image, ε is a small value set to prevent division by zero, and the parameter λ is obtained by solving the optimization problem of equation (5): Among them, H l (Y c ) represents Y c Histogram, H l (LDR) is the average histogram of the brightness channels of a large number of high-quality natural LDR images; During the training phase, the above network structure is integrated into the Denoising Diffusion Probabilistic Model (DDPM); during the testing phase, the above network structure is integrated into the Denoising Diffusion Implicit Model (DDIM); in both diffusion models, the network structure is used to predict the noise added at each step of the diffusion model. (2) Based on the network model described in step (1), the structural information is input into the first scale encoding module of the control module, and the brightness approximate distribution map is input into the second scale encoding module of the control module to further enhance the detail information of the output image. (3) Based on the DDIM described in step (1), the intermediate results generated by each iteration or part of the iteration of its backsampling iteration process are corrected, that is: take the hue information of the generated result and combine it with the structural information of the original image to further protect the structural information of the original image; the corrected result is then used to perform subsequent calculations of the backsampling process according to the formula of the diffusion model. Specifically, for any iteration of DDIM backsampling, the initial estimate of the latent space feature embedding z0 for the noiseless data is first calculated using the formula of the original DDIM algorithm. The calculation formula is as follows: in, This is the initial estimate of the latent space feature embedding z0 of the noiseless data obtained in this iteration of DDIM. t This is the result of the previous iteration. For DDIM predefined weights, ∈ θ (z t ,t,c) is the network model described in step (1), where t is the current iteration time step and c is the conditional input of the network model; Then, using the structural information of the original image... Make corrections, the correction method is as follows: in, It is The operation of transferring from the latent space of the SD model to the image domain, TSD represents the tonal-structural decomposition of the image, that is, calculating the local Gaussian mean, local variance and structural information of the image according to formulas (1) to (3) in step (1). γ is a hyperparameter that controls the degree of image detail enhancement. The structural information of the original image is given, and ε is the operation that transforms the image into the latent space of the SD model; thus, we obtain... Then, this is used as the final estimate of z0 for this iteration. Following the principles of the original DDIM algorithm, subsequent calculations for this iteration are performed to obtain the result of this iteration. The calculation formula for the subsequent calculations in this iteration is: In the formula These are also predefined weights in DDIM; (4) Network training and testing; train the network on the LDR image dataset; freeze the parameters of the generative model during training and train only the control module; the loss function is the L2 loss of the noise predicted by the network and the noise actually added to the noise-free image. Once trained, the network model can be directly used for tone mapping of HDR images or other similar tasks. During testing, only the luminance component of the image generated by the network was taken, and then converted into an RGB image using the following formula: Among them, Y pred Y and are the brightness components of the network-generated image and the original image, respectively, i∈{R,G,B}. and These represent one of the three RGB channels of the input image and the final output result, respectively, and s is a hyperparameter.

Images