Method and system for image conversion based on diffusion model, equipment and medium

Abstract

The invention relates to the technical field of computational graphics and computer vision, and discloses an image conversion method and system based on a diffusion model, and the method comprises the steps: obtaining two public data sets; constructing a diffusion network model, wherein the diffusion network model comprises a feature extraction module, a texture perception diffusion module, a color consistency module and a decoder which are sequentially connected from front to back; training a diffusion network model on the two public data sets, guiding the diffusion network model to learn by using a loss function, and outputting a trained diffusion network model; and collecting a target image, and inputting the target image into the trained diffusion network model to generate a trained conversion image. The system comprises a data acquisition unit, a network model construction unit, a training unit and a conversion unit. The invention further discloses electronic equipment and a computer readable storage medium.

Description

Technical Field

[0001] The present invention relates to the fields of computational graphics and computer vision technology, and in particular to a method, system, device and medium for image conversion based on a diffusion model, which are used to convert a single Bayer RAW image into an sRGB image using a diffusion network model. Background Art

[0002] Bayer RAW images are single-channel images captured under the Bayer array using CCD or CMOS devices as light sensors, retaining the most original data of the image. The Image Signal Processor (ISP) converts the RAW data captured by the camera sensor into sRGB images that can be perceived by the human eye. The traditional ISP pipeline consists of a series of discrete processes, including demosaicing, denoising, white balance, gamma correction, and color correction, each of which requires complex and extensive manual parameter adjustments. Complex parameter adjustments can also cause cumulative errors. In recent years, using deep learning methods to implement partial or complete ISP pipelines has become a key technology that needs to be studied urgently.

[0003] At the same time, smartphones have gradually become mainstream photography tools due to their convenience and portability. However, due to the limitations of aperture and sensor size, the image quality captured by mobile devices is difficult to match that of SLR cameras. Unlike the traditional ISP process, deep learning-based methods can achieve end-to-end conversion of RAW data captured by mobile device sensors into sRGB images with SLR quality. However, most deep learning-based ISP methods use convolutional neural networks, either focusing on solving the spatial misalignment problem in the RAW to sRGB mapping of dual-device captures, or treating the RAW to sRGB mapping as a color mapping task only, which is limited by the capabilities of convolutional neural networks, resulting in local detail differences and global color distortion.

[0004] In the present invention, we propose a method and system, device, and medium for image conversion based on a diffusion model. The ISP method for improving the quality of smartphone RAW data to sRGB images through a diffusion network model makes it possible for smartphone RAW data to obtain the same quality as sRGB images taken by SLR cameras after passing through ISP. This is achieved by decoupling the texture reconstruction diffusion network model and the color consistency module, and is trained by the proposed texture enhancement loss function and color consistency loss function. Experiments have shown that the ISP method (ISPDiffuser) for image conversion based on the diffusion network model proposed by us has surpassed existing research in the ISP method of converting mobile phone RAW to SLR camera sRGB images, and improved image quality. We believe that this work has laid a solid foundation for promoting the field of mobile phone ISP and demonstrated the potential of the diffusion network model in solving the ISP process. Summary of the invention

[0005] The present invention provides a method, system, device and medium for image conversion based on a diffusion model, and utilizes a diffusion network model to convert a single Bayer RAW image into an sRGB image.

[0006] The present invention is implemented by the following technical solution: A method for image conversion based on a diffusion model comprises the following steps:

[0007] Step S1, obtaining two public data sets;

[0008] Step S2, constructing a diffusion network model, the diffusion network model comprising a feature extraction module, a texture perception diffusion module, a color consistency module and a decoder connected in sequence from front to back, and training the diffusion network model on the two public data sets;

[0009] Step S3, using a loss function to guide the diffusion network model learning, and outputting a trained diffusion network model;

[0010] Step S4, collecting a target image, inputting the target image into a trained diffusion network model, and generating a trained conversion image.

[0011] In order to better implement the present invention, further, step S1 includes:

[0012] The two public datasets contain RAW-sRGB image pairs taken by different mobile devices and SLR cameras.

[0013] In order to better implement the present invention, further, step S2 includes:

[0014] The feature extraction module includes an encoder, which converts the RAW image I r Grayscale image I corresponding to the sRGB image g The image is input into the encoder, feature extraction is performed respectively, and then the image is downsampled and converted into latent space features, which is expressed as and ;

[0015] The texture perception diffusion module includes a connection module, a segmentation module and an attention module, and the latent space features generated by the feature extraction module are and As and x0 are input into the texture-aware diffusion module, and the predicted sRGB grayscale features are output after forward diffusion and reverse diffusion. ;

[0016] The color consistency module includes a convolution layer, a pooling layer, a SoftMax layer and a linear layer; the predicted sRGB grayscale feature RAW image features Input into the color consistency module, and output the predicted sRGB features after colorization ;

[0017] The sRGB feature Input to the decoder, add to the output RAW features of the corresponding layer of the encoder, and finally output the predicted sRGB image .

[0018] In order to better implement the present invention, further, the method for generating gray features includes:

[0019] In the forward diffusion process, Gaussian noise related to time t is added to x0 to obtain the noise grayscale feature x at time t. t ;

[0020] In the back diffusion process, the RAW image features are input and the noise grayscale feature x at time t t The position-encoded time t is spliced ​​together through the connection module, and then the predicted sRGB grayscale features are output through the segmentation module and the attention module. .

[0021] In order to better implement the present invention, further, the output predicts the sRGB features after colorization The methods include:

[0022] The RAW image features Input to the convolution layer, pooling layer, SoftMax layer and linear layer in turn, and output the predicted sRGB color histogram ;

[0023] The predicted sRGB color histogram And the predicted sRGB grayscale features Colorized sRGB features predicted by residual blocks and cross-attention outputs .

[0024] In order to better implement the present invention, further, step S3 includes:

[0025] The diffusion network model adopts a two-stage training method;

[0026] The encoder-decoder parameters are trained in the first stage, and the encoder-decoder parameters trained in the first stage are frozen in the second stage, and the texture-aware diffusion module and the color consistency module are trained.

[0027] In order to better implement the present invention, further, the loss function used in the first stage is: Where I is the RAW image, sRGB image and the corresponding grayscale image, ε(·) is the encoder, D(·) is the decoder, and ||·||2 is the l2 norm;

[0028] The texture-aware diffusion module in the second stage is subject to the loss function constraint, It is composed of l2 loss function and texture enhancement loss function, expressed as: Where T(·) represents the Canny texture information of the extracted input features, ||·||1 and ||·||2 represent the 1l norm and l2 norm respectively, and λ1 is 0.01;

[0029] The constraints that the color consistency module needs to satisfy are expressed as: in and They are the color histogram and feature representation of the target sRGB image, respectively, and λ2 is 0.01;

[0030] The two-stage training loss function is

[0031] The present invention also provides a system for image conversion based on a diffusion model, comprising a data acquisition unit, a network model construction unit, a training unit and a conversion unit, wherein:

[0032] A data acquisition unit, used to acquire two public data sets;

[0033] A network model construction unit, used to construct a diffusion network model, wherein the diffusion network model includes a feature extraction module, a texture perception diffusion module, a color consistency module and a decoder connected in sequence from front to back;

[0034] A training unit, used for training the diffusion network model on the two public data sets, guiding the learning of the diffusion network model using a loss function, and outputting a trained diffusion network model;

[0035] The conversion unit is used to collect a target image, input the target image into a trained diffusion network model, and generate a trained conversion image.

[0036] The present invention also provides an electronic device, which includes a processor and a memory; the processor includes the system for performing image conversion based on the diffusion model described in the second aspect above.

[0037] The present invention also provides a computer-readable storage medium, which includes instructions; when the instructions are executed on the electronic device described in the third aspect, the electronic device executes the method described in the first aspect.

[0038] Compared with the prior art, the present invention has the following advantages and beneficial effects:

[0039] (1) The present invention provides a method, system, device, and medium for image conversion based on a diffusion model, so that the RAW data of a smartphone can obtain the same quality as sRGB images taken by a SLR camera after passing through the ISP. BRIEF DESCRIPTION OF THE DRAWINGS

[0040] The present invention is further described in conjunction with the following drawings and embodiments, and all the concepts and innovations of the present invention should be regarded as the disclosed contents and the protection scope of the present invention.

[0041] Figure 1 A schematic diagram of a method for image conversion based on a diffusion model and a method flow diagram in a system, device, and medium provided in an embodiment of the present application;

[0042] Figure 2 A structural diagram of a diffusion network model in a method and system, a device, and a medium for performing image conversion based on a diffusion model provided in an embodiment of the present application;

[0043] Figure 3 A structural diagram of a method and system, device, and a color consistency module in a medium for performing image conversion based on a diffusion model provided in an embodiment of the present application;

[0044] Figure 4 A schematic diagram of the effects of a method, system, device, and medium for performing image conversion based on a diffusion model provided in an embodiment of the present application. DETAILED DESCRIPTION

[0045] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. It should be understood that the described embodiments are only part of the embodiments of the present invention, not all of the embodiments, and therefore should not be regarded as limiting the scope of protection. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technical personnel in this field without making creative work are within the scope of protection of the present invention.

[0046] In the description of the present invention, it should be noted that, unless otherwise clearly specified and limited, the terms "disposed", "connected", and "connected" should be understood in a broad sense, for example, it can be a fixed connection, a detachable connection, or an integral connection; it can be a mechanical connection or an electrical connection; it can be a direct connection, or it can be an indirect connection through an intermediate medium, or it can be the internal communication of two components. For ordinary technicians in this field, the specific meanings of the above terms in the present invention can be understood according to specific circumstances.

[0047] Embodiment 1:

[0048] The present embodiment provides a method, system, device, and medium for performing image conversion based on a diffusion model, such as Figures 1 to 4 As shown, in this invention, we propose a decoupled ISP method based on a diffusion network model, which is designed to separate the RAW to sRGB mapping into two parts: detail reconstruction in grayscale space and color consistency mapping from grayscale space to sRGB. Detail reconstruction in grayscale space is achieved by the designed texture-aware diffusion network module (TADM), which leverages the generative power of the diffusion network model and focuses on detail reconstruction without considering color information. The texture-aware diffusion network module (TADM) performs the diffusion process under the guidance of the encoded RAW features and the newly proposed texture enhancement loss. Color consistency mapping from grayscale space to sRGB is achieved by the designed color consistency module (HCCM), which adopts color histogram as a guide to learn accurate color information, converts the gray features generated by TADM into sRGB features with consistent colors with SLR camera images, and supervises the learned color information through color consistency loss.

[0049] Embodiment 2:

[0050] This embodiment is further optimized on the basis of Embodiment 1. We train the network on two public data sets. The two data sets contain a large number of RAW-sRGB image pairs taken by different mobile devices and SLR cameras, which can ensure that the model fully learns the mapping from mobile phone RAW images to camera sRGB images.

[0051] The other parts of this embodiment are the same as those of Embodiment 1, and thus will not be described in detail.

[0052] Embodiment 3:

[0053] This embodiment is further optimized on the basis of the above embodiment 1 or 2. Figure 2As shown in the figure, in the model construction and reasoning stage, the RAW image is input into the encoder for feature extraction, and then the image is downsampled and converted into latent space features. The latent space features generated by the feature extraction module are used as input to the texture-aware diffusion module together with randomly sampled Gaussian noise. After forward diffusion and backward diffusion, the predicted sRGB grayscale features are output, and finally the predicted sRGB image is output through the color consistency module and the decoder.

[0054] In the training phase of the model, the RAW image I r Grayscale image I corresponding to the sRGB image g Input encoder, perform feature extraction, downsample the image and convert it into latent space feature representation and As and x0 are input into the texture-aware diffusion network module (TADM). In the forward diffusion process, Gaussian noise related to time t is added to x0 to obtain the noise grayscale feature x at time t. t In the back diffusion process, the RAW image features are input and the noise grayscale feature x at time t t And the time t after position embedding is spliced ​​together through the Concatenate module (connection module), and then the predicted sRGB grayscale features are output through the U-net module (segmentation module) and the PatchAttention module (attention module) Predicted sRGB grayscale features RAW image features As input to the color consistency module (HCCM).

[0055] like Figure 3 As shown, Input to the convolution layer (Conv), pooling layer (Pooling), SoftMax layer and linear layer (Linear) in sequence, and output the predicted sRGB color histogram . Predicted sRGB color histogram And the predicted sRGB grayscale features , output the predicted colored sRGB features through resblock (residual block) and cross-attention (cross attention) . Input to the decoder, add to the output RAW features of the corresponding layer of the encoder, and finally output the predicted sRGB image In the inference phase of the model, a RAW image captured by a mobile device is input r, the corresponding feature representation is obtained through the encoder, and the sampling method of DDIM is used to obtain , and then the predicted sRGB image is obtained through the color consistency module (HCCM) and decoder (Decoder).

[0056] The other parts of this embodiment are the same as any of the above embodiments 1-2, so they are not repeated here.

[0057] Embodiment 4:

[0058] This embodiment is further optimized on the basis of any one of the above embodiments 1-3. The network adopts a two-stage training method. First, the encoder-decoder architecture is trained, and then the encoder-decoder parameters are frozen to train the texture-aware diffusion network model and the color consistency model. The loss function used in the first stage is as follows:

[0059]

[0060] Among them, I is the RAW image, sRGB image and the corresponding grayscale image, ε(·) is the encoder, D(·) is the decoder, and ||·||2 is the norm.

[0061] The two-stage network texture-aware diffusion network model is subject to the loss function constraint, It consists of l2 loss function and texture enhancement loss function:

[0062]

[0063] Where T(·) represents the Canny texture information extracted from the input features, ||·||1 and ||·||2 represent the l1 norm and l2 norm respectively, and λ1 is 0.01

[0064] The constraints that the color consistency module needs to satisfy are expressed as:

[0065]

[0066] in and They are the color histogram and feature representation of the target sRGB image, respectively, and λ2 is 0.01.

[0067] The two-stage training loss function is expressed as:

[0068]

[0069] The other parts of this embodiment are the same as any one of the above embodiments 1-3, so they are not repeated here.

[0070] Embodiment 5:

[0071] This embodiment is further optimized based on any one of the above embodiments 1-4. Figure 4 The results are shown on the public datasets ZRR (first row) and MAI (second row). (a) is the target sRGB image, (b)-(g) are the sRGB images predicted by the previous advanced methods, and (h) is the effect of our proposed method. Error Map is the difference heat map between the predicted image and the target image.

[0072] The other parts of this embodiment are the same as any one of the above embodiments 1-4, so they are not repeated here.

[0073] Embodiment 6:

[0074] The present invention also provides a system for image conversion based on a diffusion model that matches the method, comprising a data acquisition unit, a network model construction unit, a training unit and a conversion unit.

[0075] The present invention also provides an electronic device, which includes a processor and a memory; the processor includes the above-mentioned system for performing image conversion based on the diffusion model.

[0076] Embodiment 7:

[0077] The present invention also provides a computer-readable storage medium, which includes instructions; when the instructions are executed on the electronic device described in the above embodiment, the electronic device executes the method described in the above embodiment. Optionally, the computer-readable storage medium may be a memory.

[0078] The processor involved in the embodiments of the present application may be a chip. For example, it may be a field programmable gate array (FPGA), an application specific integrated circuit (ASIC), a system on chip (SoC), a central processor unit (CPU), a network processor (NP), a digital signal processor (DSP), a microcontroller unit (MCU), a programmable logic device (PLD), or other integrated chips.

[0079] The memory involved in the embodiments of the present application may be a volatile memory or a non-volatile memory, or may include both volatile and non-volatile memories. Among them, the non-volatile memory may be a read-only memory (ROM), a programmable read-only memory (PROM), an erasable programmable read-only memory (EPROM), an electrically erasable programmable read-only memory (EEPROM), or a flash memory. The volatile memory may be a random access memory (RAM), which is used as an external cache. By way of example and not limitation, many forms of RAM are available, such as static RAM (SRAM), dynamic RAM (DRAM), synchronous DRAM (SDRAM), double data rate SDRAM (DDR SDRAM), enhanced SDRAM (ESDRAM), synchronous link DRAM (SLDRAM), and direct RAM (DR RAM). It should be noted that the memory of the systems and methods described herein is intended to include, but is not limited to, these and any other suitable types of memory.

[0080] It should be understood that in the various embodiments of the present application, the size of the serial numbers of the above-mentioned processes does not mean the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiments of the present application.

[0081] Those of ordinary skill in the art will appreciate that the modules and algorithm steps of each example described in conjunction with the embodiments disclosed herein can be implemented in electronic hardware, or a combination of computer software and electronic hardware. Whether these functions are performed in hardware or software depends on the specific application and design constraints of the technical solution. Professional and technical personnel can use different methods to implement the described functions for each specific application, but such implementation should not be considered to be beyond the scope of this application.

[0082] Those skilled in the art can clearly understand that, for the convenience and brevity of description, the specific working processes of the systems, devices and modules described above can refer to the corresponding processes in the aforementioned method embodiments and will not be repeated here.

[0083] In the several embodiments provided in the present application, it should be understood that the disclosed systems, devices and methods can be implemented in other ways. For example, the device embodiments described above are only schematic. For example, the division of the modules is only a logical function division. There may be other division methods in actual implementation, such as multiple modules or components can be combined or integrated into another device, or some features can be ignored or not executed. Another point is that the mutual coupling or direct coupling or communication connection shown or discussed can be through some interfaces, indirect coupling or communication connection of devices or modules, which can be electrical, mechanical or other forms.

[0084] The modules described as separate components may or may not be physically separated, and the components shown as modules may or may not be physical modules, that is, they may be located in one device or distributed on multiple devices. Some or all of the modules may be selected according to actual needs to achieve the purpose of the present embodiment.

[0085] In addition, each functional module in each embodiment of the present application may be integrated into one device, or each module may exist physically separately, or two or more modules may be integrated into one device.

[0086] In the above embodiments, it can be implemented in whole or in part by software, hardware, firmware or any combination thereof. When implemented using a software program, it can be implemented in whole or in part in the form of a computer program product. The computer program product includes one or more computer instructions. When loading and executing a computer program instruction on a computer, the process or function described in the embodiment of the present application is generated in whole or in part. The computer may be a general-purpose computer, a special-purpose computer, a computer network, or other programmable devices. The computer instructions may be stored in a computer-readable storage medium, or transmitted from one computer-readable storage medium to another computer-readable storage medium, for example, the computer instructions may be transmitted from a website site, a computer, a server or a data center by wired (e.g., coaxial cable, optical fiber, digital subscriber line (Digital Subscriber Line, DSL)) or wireless (e.g., infrared, wireless, microwave, etc.) mode to another website site, computer, server or data center. The computer-readable storage medium may be any available medium that a computer can access or may contain one or more servers, data centers and other data storage devices that can be integrated with a medium. The available medium may be a magnetic medium (eg, a floppy disk, a hard disk, a magnetic tape), an optical medium (eg, a DVD), or a semiconductor medium (eg, a solid state disk (SSD)).

[0087] The above description is only a preferred embodiment of the present invention and does not limit the present invention in any form. Any simple modification or equivalent change made to the above embodiment based on the technical essence of the present invention shall fall within the protection scope of the present invention.

Claims

1. A method for image conversion based on a diffusion model, characterized in that: The following steps are involved: Step S1, obtaining two public data sets; Step S2, constructing a diffusion network model, the diffusion network model comprising a feature extraction module, a texture perception diffusion module, a color consistency module and a decoder connected in sequence from front to back, and training the diffusion network model on the two public data sets; Step S3, using a loss function to guide the diffusion network model learning, and outputting a trained diffusion network model; Step S4, collecting a target image, inputting the target image into a trained diffusion network model, and generating a trained conversion image.

2. The method for image conversion based on a diffusion model according to claim 1, characterized in that: The step S1 comprises: The two public datasets contain RAW-sRGB image pairs taken by different mobile devices and SLR cameras.

3. The method for image conversion based on a diffusion model according to claim 1, characterized in that: The step S2 comprises: The feature extraction module includes an encoder, which converts the RAW image I r Grayscale image I corresponding to the sRGB image g The image is input into the encoder, feature extraction is performed respectively, and then the image is downsampled and converted into latent space features, which is expressed as and The texture perception diffusion module includes a connection module, a segmentation module and an attention module, and the latent space features generated by the feature extraction module are and As and x0 are input into the texture-aware diffusion module, and the predicted sRGB grayscale features are output after forward diffusion and reverse diffusion. The color consistency module includes a convolution layer, a pooling layer, a SoftMax layer and a linear layer; the predicted sRGB grayscale feature RAW image features Input into the color consistency module, and output the predicted sRGB features after colorization The sRGB feature Input to the decoder, add to the output RAW features of the corresponding layer of the encoder, and finally output the predicted sRGB image 4. The method for image conversion based on diffusion model according to claim 3, characterized in that: The method for generating gray features comprises: In the forward diffusion process, Gaussian noise related to time t is added to x0 to obtain the noise grayscale feature x at time t. t ; In the back diffusion process, the RAW image features are input and the noise grayscale feature x at time t t The position-encoded time t is spliced ​​together through the connection module, and then the predicted sRGB grayscale features are output through the segmentation module and the attention module.

5. The method for image conversion based on diffusion model according to claim 3, characterized in that: The output predicts the sRGB features after colorization The methods include: The RAW image features Input to the convolution layer, pooling layer, SoftMax layer and linear layer in turn, and output the predicted sRGB color histogram The predicted sRGB color histogram And the predicted sRGB grayscale features Colorized sRGB features predicted by residual blocks and cross-attention outputs 6. The method for image conversion based on a diffusion model according to claim 1, characterized in that: The step S3 comprises: The diffusion network model adopts a two-stage training method; The encoder-decoder parameters are trained in one stage, and the encoder-decoder parameters are frozen in the second stage to train the texture-aware diffusion module and the color consistency module.

7. The method for image conversion based on diffusion model according to claim 6, characterized in that: include: The loss function used in the first stage is: Where I is the RAW image, sRGB image and the corresponding grayscale image, ε(·) is the encoder, D(·) is the decoder, and ||·||2 is the l2 norm; The texture-aware diffusion module in the second stage is subject to the loss function constraint, It is composed of l2 loss function and texture enhancement loss function, expressed as: Where T(·) represents the Canny texture information of the extracted input features, ||·||1 and ||·||2 represent the l1 norm and l2 norm respectively, and λ1 is 0.01; The constraints that the color consistency module needs to satisfy are expressed as: in and They are the color histogram and feature representation of the target sRGB image, respectively, and λ2 is 0.01; The two-stage training loss function is 8. A system for image conversion based on a diffusion model, characterized in that: It includes a data acquisition unit, a network model building unit, a training unit and a conversion unit, wherein: A data acquisition unit, used to acquire two public data sets; A network model construction unit, used to construct a diffusion network model, wherein the diffusion network model includes a feature extraction module, a texture perception diffusion module, a color consistency module and a decoder connected in sequence from front to back; A training unit, used for training the diffusion network model on the two public data sets, guiding the learning of the diffusion network model using a loss function, and outputting a trained diffusion network model; The conversion unit is used to collect a target image, input the target image into a trained diffusion network model, and generate a trained conversion image.

9. An electronic device, characterized in that: It comprises a processor and a memory; the processor comprises the system for performing image conversion based on a diffusion model as claimed in claim 8.

10. A computer-readable storage medium, characterized in that: The computer-readable storage medium includes instructions; when the instructions are executed on the electronic device as claimed in claim 9, the electronic device is caused to execute the method as claimed in any one of claims 1 to 7.

Images