Model training method, image reconstruction method, device, electronic device and medium

Abstract

The present application discloses a model training method, an image reconstruction method, an apparatus, an electronic device, and a medium, relating to the field of neural network technology. The method first inputs an initial training image into an initial image reconstruction model; performs image segmentation processing through a pre-trained segmentation network module to obtain a segmentation mask probability training image; inputs the segmentation mask probability training image and the initial training image into the denoising unit of the diffusion network module to obtain a noisy training image; inputs the noisy training image into the denoising unit of the diffusion network module to obtain predicted training noise and a reconstructed training image; calculates the total loss based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image; iteratively updates the initial image reconstruction model based on the total loss to obtain a trained target image reconstruction model. This method can improve the reconstruction quality of the image reconstruction model for images with occluded areas.

Description

Technical Field

[0001] The present application relates to the field of neural network technology, and in particular to a model training method, image reconstruction method, device, electronic equipment and medium. Background Art

[0002] Image object removal is a technique designed to remove specific objects from an image and restore the integrity of the occluded background. It is widely used in surveillance, security, medical image analysis, and other fields. By improving image integrity and usability, this technology provides important support for the application of computer vision systems in various fields.

[0003] With the rapid development of deep learning, generative adversarial networks (GANs) and their derivatives have been widely used in image object removal tasks. These networks and their derivatives have achieved remarkable results in processing objects with clear boundaries and distinct outlines. However, current networks and their derivatives suffer from poor image reconstruction quality for images with blurred boundaries and translucent properties, such as flames and smoke. Summary of the Invention

[0004] The present application aims to solve at least one of the technical problems existing in the prior art. To this end, the present application proposes a model training method, image reconstruction method, device, electronic device and medium that can improve the reconstruction quality of an image reconstruction model for images with occluded areas.

[0005] To achieve the above objectives, the first embodiment of the present application provides an image reconstruction model training method, comprising:

[0006] Obtaining an initial training image with an occluded area and a true label image corresponding to the initial training image;

[0007] Inputting the initial training image into an initial image reconstruction model; the initial image reconstruction model includes an initial diffusion network module and a pre-trained segmentation network module;

[0008] Performing image segmentation processing on the initial training image using the segmentation network module to obtain a segmentation mask probability training image; wherein a first pixel of the segmentation mask probability training image corresponds one-to-one to a second pixel of the initial training image, and a value of the first pixel represents a probability that the corresponding second pixel is an occluded area;

[0009] Inputting the segmentation mask probability training image and the initial training image into the noise adding unit of the diffusion network module to obtain a noisy training image;

[0010] Inputting the noisy training image into the denoising unit of the diffusion network module to obtain predicted training noise and reconstructed training image;

[0011] Calculating a total loss based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image;

[0012] The initial image reconstruction model is iteratively updated based on the total loss to obtain a trained target image reconstruction model.

[0013] According to some embodiments of the first aspect of the present application, inputting the segmentation mask probability training image and the initial training image into the noise addition unit of the diffusion network module to obtain the noisy training image includes:

[0014] Performing a pixel-by-pixel weighting operation on the initial training image and the segmentation mask probability training image to obtain a weighted training image;

[0015] The weighted training image is subjected to noise addition processing by the noise addition unit to obtain the noisy training image.

[0016] According to some embodiments of the first aspect of the present application, inputting the noisy training image into the denoising unit of the diffusion network module to obtain predicted training noise and reconstructed training image includes:

[0017] performing noise prediction processing on the noisy training image based on the denoising unit to obtain predicted training noise;

[0018] Image reconstruction processing is performed based on the predicted training noise and the noisy training image to obtain the reconstructed training image.

[0019] According to some embodiments of the first aspect of the present application, the calculating the total loss based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image includes:

[0020] Calculating an image reconstruction loss based on the reconstructed training image and the initial training image;

[0021] Calculating a denoising loss based on the predicted training noise;

[0022] Calculating a mask consistency loss based on the predicted training noise and the segmentation mask probability training image;

[0023] Calculating a segmentation loss based on the segmentation mask probability training image and the true label image;

[0024] Calculating a sparse loss based on the segmentation mask probability training image;

[0025] The total loss is obtained based on the image reconstruction loss, the denoising loss, the mask consistency loss, the segmentation loss and the sparsity loss.

[0026] According to some embodiments of the first aspect of the present application, calculating the image reconstruction loss based on the reconstructed training image and the initial training image includes:

[0027] The image reconstruction loss is calculated based on the reconstructed training image, the initial training image and a reconstruction loss function; the reconstruction loss function is:

[0028] ;

[0029] in, L recon Characterize the image reconstruction loss, I output characterizing the reconstructed training image, I org The initial training image is characterized.

[0030] According to some embodiments of the first aspect of the present application, the calculating the denoising loss based on the predicted training noise includes:

[0031] The denoising loss is calculated based on the predicted training noise and the denoising loss function; the denoising loss function is:

[0032] ;

[0033] in, L denoise represents the denoising loss; E represents the mathematical expectation; is Gaussian noise; represents the predicted training noise, y0 represents the noisy training image; y t represents the noisy training image obtained by the diffusion network module at time step t; Mprob represents the segmentation mask probability training image.

[0034] According to some embodiments of the first aspect of the present application, the calculating the mask consistency loss based on the predicted training noise and the segmentation mask probability training image includes:

[0035] The mask consistency loss is calculated based on the predicted training noise, the segmentation mask probability training image and the mask consistency loss function; the mask consistency loss function is:

[0036] ;

[0037] in, L mask is the mask consistency loss, E represents the mathematical expectation; is Gaussian noise, Characterize the prediction training noise, y t represents the noisy training image obtained by the diffusion network module at time step t; Mprob represents the segmentation mask probability training image, and y0 represents the noisy training image when the time step is 0.

[0038] According to some embodiments of the first aspect of the present application, the calculating the segmentation loss based on the segmentation mask probability training image and the true label image includes:

[0039] The segmentation loss is calculated based on the segmentation mask probability training image, the true label image and the segmentation loss function; the segmentation loss function is:

[0040] ;

[0041] in, L seg Characterize the segmentation loss, Characterize the probability of pixel i in the segmentation mask probability training image; represents the probability of pixel i in the true label image; N is the total number of pixels in the segmentation mask probability training image.

[0042] To achieve the above objectives, a second embodiment of the present application provides an image reconstruction method, comprising:

[0043] Acquire an image to be reconstructed with an occluded area;

[0044] The image to be reconstructed is input into a target image reconstruction model to obtain a target reconstructed image; wherein, the target image reconstruction model is trained according to the image reconstruction model training method of any one of the embodiments of the first aspect.

[0045] To achieve the above objectives, the third embodiment of the present application provides an image reconstruction model training device, comprising:

[0046] An acquisition module, configured to acquire an initial training image having an occluded area and a true label image corresponding to the initial training image;

[0047] A first input module is used to input the initial training image into an initial image reconstruction model; the initial image reconstruction model includes an initial diffusion network module and a pre-trained segmentation network module;

[0048] a segmentation module configured to perform image segmentation processing on the initial training image using the segmentation network module to obtain a segmentation mask probability training image; wherein a first pixel of the segmentation mask probability training image corresponds one-to-one to a second pixel of the initial training image, and a value of the first pixel represents a probability that the corresponding second pixel is an occluded area;

[0049] A second input module is configured to input the segmentation mask probability training image and the initial training image into the noise adding unit of the diffusion network module to obtain a noisy training image;

[0050] A third input module is used to input the noisy training image into the denoising unit of the diffusion network module to obtain predicted training noise and reconstructed training image;

[0051] A calculation module, configured to calculate a total loss based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image;

[0052] An iterative update module is used to iteratively update the initial image reconstruction model based on the total loss to obtain a trained target image reconstruction model.

[0053] To achieve the above-mentioned purpose, the fourth aspect embodiment of the present application provides an electronic device, which includes a memory and a processor, and the memory stores a computer program. When the processor executes the computer program, it implements the image reconstruction model training method described in any one of the first aspect embodiments, or the image reconstruction method described in the second aspect embodiment.

[0054] To achieve the above-mentioned purpose, the fifth aspect embodiment of the present application provides a computer-readable storage medium, which stores a computer program. When the computer program is executed by a processor, it is the image reconstruction model training method described in any one of the first aspect embodiments, or the image reconstruction method described in the second aspect embodiment.

[0055] According to the model training method, image reconstruction method, device, electronic device and medium of the embodiments of the present application, the initial training image is first input into the initial image reconstruction model; the initial image reconstruction model includes an initial diffusion network module and a pre-trained segmentation network module; image segmentation processing is performed by the pre-trained segmentation network module to obtain a segmentation mask probability training image, and the segmentation mask probability training image and the initial training image are input into the denoising unit of the diffusion network module to obtain a noisy training image; the noisy training image is input into the denoising unit of the diffusion network module to obtain predicted training noise and a reconstructed training image; based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image, the total loss is calculated, and the initial image reconstruction model is iteratively updated based on the total loss to obtain a trained target image reconstruction model. In this process, since the segmentation network module is pre-trained, the accuracy of the segmentation network module is relatively high. In the segmentation mask probability training image output by the segmentation network module, the first pixel of the segmentation mask probability training image corresponds one-to-one to the second pixel of the initial training image, and the value of the first pixel represents the probability that the corresponding second pixel is an occluded area. Therefore, the information represented by the segmentation mask probability training image can be used as prior knowledge to help the diffusion network module learn the intrinsic mapping mechanism between the transparency of the occluded area and the information represented by the segmentation mask probability training image based on prior knowledge to accurately distinguish the edge and background pixels of the occluded area, thereby improving the ability to recognize occluded areas and overcome texture interference. Under the guidance of prior knowledge, the diffusion network module can deeply explore the background structure information that may exist in the occluded area, and while removing blurring objects such as flames and smoke, reasonably integrate this background information into the image reconstruction process. In this way, the diffusion network module can better preserve the background characteristics and logical structure of the original scene of the initial training image, help the reconstruction model solve the edge blur dilemma, improve the accuracy and effect of edge processing of the image reconstruction model, and thus significantly enhance the quality of image reconstruction. In this process, the information represented by the segmentation mask probability training image can divide the initial training image into occluded areas and non-occluded areas. In the denoising unit, the occluded areas are guided by noise, prompting the image reconstruction model to generate natural and continuous results in the subsequent denoising process.

[0056] Additional aspects and advantages of the present application will be given in part in the description below, and in part will become obvious from the description below, or will be learned through practice of the present application. BRIEF DESCRIPTION OF THE DRAWINGS

[0057] The present application is further described below with reference to the accompanying drawings and embodiments, wherein:

[0058] Figure 1 This is a flowchart of the steps of the image reconstruction model training method according to an embodiment of the present application;

[0059] Figure 2 A schematic diagram of the framework of the image reconstruction model according to an embodiment of the present application;

[0060] Figure 3 This is a specific flowchart of step S140 in the embodiment of the present application;

[0061] Figure 4 This is a specific flowchart of step S150 in an embodiment of the present application;

[0062] Figure 5 This is a specific flowchart of step S160 in the embodiment of the present application;

[0063] Figure 6 This is a schematic flow chart of the steps of the image reconstruction method according to an embodiment of the present application;

[0064] Figure 7 Schematic diagram of the structure of the image reconstruction device according to an embodiment of the present application;

[0065] Figure 8 This is a schematic diagram of the structure of the image reconstruction model training device according to an embodiment of the present application;

[0066] Figure 9 This is a schematic structural diagram of an electronic device according to an embodiment of the present application. DETAILED DESCRIPTION

[0067] The following describes in detail embodiments of the present application. Examples of the embodiments are shown in the accompanying drawings, wherein the same or similar reference numerals throughout represent the same or similar elements or elements having the same or similar functions. The embodiments described below with reference to the accompanying drawings are exemplary and are only used to explain the present application and are not to be construed as limiting the present application.

[0068] In the description of this application, it should be understood that descriptions involving orientations, such as up, down, front, back, left, right, etc., indicating orientations or positional relationships, are based on the orientations or positional relationships shown in the accompanying drawings. They are only for the convenience of describing this application and simplifying the description, and do not indicate or imply that the device or element referred to must have a specific orientation, be constructed and operated in a specific orientation. Therefore, they cannot be understood as limitations on this application.

[0069] In the description of this application, "several" means more than one, "plurality" means more than two, "greater than," "less than," and "exceed" are understood to exclude the number itself, while "above," "below," and "within" are understood to include the number itself. The use of "first" and "second" in the description is solely for the purpose of distinguishing technical features and should not be construed as indicating or implying relative importance, implicitly specifying the number of the indicated technical features, or implicitly specifying the order of the indicated technical features.

[0070] In the description of this application, unless otherwise clearly defined, terms such as setting, installing, and connecting should be understood in a broad sense, and technicians in the relevant technical field can reasonably determine the specific meanings of the above terms in this application based on the specific content of the technical solution.

[0071] In the description of this application, reference to the terms "one embodiment," "some embodiments," "illustrative embodiments," "examples," "specific examples," or "some examples" means that the specific features, structures, materials, or characteristics described in conjunction with the embodiment or example are included in at least one embodiment or example of this application. In this specification, the schematic representations of the above terms do not necessarily refer to the same embodiment or example. Moreover, the specific features, structures, materials, or characteristics described can be combined in any appropriate manner in any one or more embodiments or examples.

[0072] Some methods in the related art have many limitations when processing translucent objects such as flames. As far as edge processing is concerned, when removing translucent objects such as flames, the edge blurring phenomenon is extremely prominent. This is mainly due to the fact that the translucent characteristics cause complex pixel changes in the transition area between the edge of the object and the background, and the texture information of the occluded area is intertwined with the background, resulting in serious texture interference, which greatly affects the accuracy and effect of the removal. In terms of occlusion processing and completion, the currently commonly used strategy is to directly remove the occluded area of ​​translucent objects such as flames, and then fill in the gaps. However, this method has obvious drawbacks. It does not fully consider the background structure information contained in the translucent area. This information is lost during the removal and completion process, resulting in the destruction of the logical correlation between the reconstructed image and the original scene, ultimately resulting in poor image reconstruction quality.

[0073] Based on this, embodiments of the present application provide a model training method, image reconstruction method, apparatus, electronic device, and medium that can improve the reconstruction quality of an image reconstruction model for images with occluded areas. The information representing the segmentation mask probability training image output by the segmentation network module in the model can serve as prior knowledge. Under the guidance of this prior knowledge, the diffusion network module can deeply explore the background structure information that may exist in the occluded area. In this way, the diffusion network module can better preserve the background features and logical structure of the original scene of the initial training image, thereby significantly enhancing the quality of image reconstruction.

[0074] The first aspect of the present application provides an image reconstruction model training method. The image reconstruction model training method of the embodiment of the present application can be applied to a terminal or a server. In some embodiments, the terminal can be a smart phone, a tablet computer, a laptop computer, a desktop computer, a router, a programmable switch, a network card, etc.; the server can be configured as an independent physical server, or as a server cluster or distributed system composed of multiple physical servers. It can also be configured as a cloud server that provides basic cloud computing services such as cloud services, cloud databases, cloud computing, cloud functions, cloud storage, network services, cloud communications, middleware services, domain name services, security services, CDN, and big data and artificial intelligence platforms; the software can be an application that implements the image reconstruction model training method, etc., but is not limited to the above forms.

[0075] The present application can be used in many general or special computer system environments or configurations. For example: personal computers, server computers, handheld or portable devices, tablet devices, multi-processor systems, microprocessor-based systems, set-top boxes, programmable consumer electronic devices, network PCs, minicomputers, mainframe computers, distributed computing environments including any of the above systems or devices, etc. The present application can be described in the general context of computer-executable instructions executed by a computer, such as program modules. Generally, program modules include routines, programs, objects, components, data structures, etc. that perform specific services or implement specific abstract data types. The present application can also be practiced in distributed computing environments, in which services are performed by remote processing devices connected through a communication network. In a distributed computing environment, program modules can be located in local and remote computer storage media including storage devices.

[0076] Reference Figure 1 , Figure 1 1 is a flowchart of the steps of the image reconstruction model training method according to an embodiment of the present application. The image reconstruction model training method may include but is not limited to steps S110 to S170.

[0077] Step S110, obtaining an initial training image with an occluded area and a true label image corresponding to the initial training image;

[0078] It is worth noting that an object with flames or smoke can be photographed by a camera to obtain an initial training image with an occluded area, wherein the occluded area is a flame part or a smoke part. In addition, the occluded area can also be a cloud part, a water vapor part, a translucent glass part, a mosaic part, etc. After obtaining the initial training image, the initial training image is manually annotated to obtain a real label image corresponding to the initial training image. For example, the real label image is obtained by annotating the initial training image, and the pixels in the occluded area in the initial training image are annotated as 1, and the pixels in the non-occluded area are annotated as 0. In this way, the real label image corresponding to the initial training image is obtained. In another embodiment, the initial training image and the corresponding real label image are obtained from a public material library. There are many ways to obtain from the material library. For example, the material library can be accessed through a custom application to obtain the initial training image and the corresponding real label image.

[0079] Step S120, inputting the initial training image into the initial image reconstruction model; the initial image reconstruction model includes an initial diffusion network module and a pre-trained segmentation network module;

[0080] Step S130: performing image segmentation processing on the initial training image using a segmentation network module to obtain a segmentation mask probability training image; wherein a first pixel of the segmentation mask probability training image corresponds one-to-one to a second pixel of the initial training image, and the value of the first pixel represents the probability that the corresponding second pixel is an occluded area;

[0081] It is worth noting that, referring to Figure 2 , Figure 2 This is a schematic diagram of the framework of the image reconstruction model of an embodiment of the present application. The segmentation network module in the image reconstruction model adopts an encoder-decoder architecture with an attention mechanism. Specifically, in step S130, the initial training image is first converted into an input sequence. The encoder is responsible for converting the input sequence into a fixed-length context vector, and the decoder generates an output sequence based on this context vector. The processing process of the attention mechanism is expressed as follows:

[0082] ;

[0083] Among them, Q, K, and V represent the query vector (Query), key vector (Key), and value vector (Value) of the initial training image, respectively. kRepresenting model dimensions. Through the attention mechanism, the segmentation network module is able to capture long-range dependencies in the initial training image, thereby better identifying occluded areas. The self-attention mechanism is integrated into the encoder part as a feature enhancement module. This module captures the global context information of the occluded area through global feature modeling, while enhancing the feature saliency of the flame area through the channel attention mechanism. The processing of the segmentation network module can be expressed as:

[0084] ;

[0085] in, Mprob represents the segmentation mask probability training image, f seg represents the segmentation network module, I org Characterize the initial training image. The segmentation mask probability training image represents the spatial distribution of occluded regions (i.e., flames or smoke). The value of each first pixel in the segmentation mask probability training image ranges from [0, 1], where values ​​close to 1 indicate a high probability that the corresponding second pixel in the initial training image is an occluded region, while values ​​close to 0 indicate a low probability that the corresponding second pixel in the initial training image is an occluded region. The value of the first pixel is greater than or equal to 0 and less than or equal to 1. For example, if a first pixel value is 0.9, the probability that the corresponding second pixel in the initial training image is an occluded region is 90%; if a first pixel value is 0.01, the probability that the corresponding second pixel in the initial training image is an occluded region is 1%. The segmentation mask probability training image not only provides explicit locations of flame regions but also provides spatial guidance for the next stage of denoising learning. Occluded regions typically only occupy a portion of the image. Therefore, sparsity constraints are used to constrain the distribution of the segmentation mask probability training image to prevent the segmentation network module from generating unnecessary, large-scale masks.

[0086] In one embodiment, the segmentation network module uses the Attention U-Net network, a deep learning model that incorporates an attention mechanism based on the U-Net. The U-Net is a convolutional neural network architecture commonly used for medical image segmentation. It consists of an encoder and a decoder, which transmits feature maps from the encoder to the decoder via skip connections to preserve more spatial information.

[0087] Step S140, inputting the segmentation mask probability training image and the initial training image into the noise adding unit of the diffusion network module to obtain a noisy training image;

[0088] Step S150, inputting the noisy training image into the denoising unit of the diffusion network module to obtain predicted training noise and reconstructed training image;

[0089] In some embodiments, the diffusion network module includes a denoising unit and a denoising unit. The diffusion network module uses the DDPM (Denoising Diffusion Probabilistic Models) diffusion model, a generative model based on variational inference and Markov chains, primarily used for image generation tasks. It generates high-quality images by learning to gradually add noise to the data (forward diffusion process) and then gradually remove the noise to restore the original data (backward diffusion process). The denoising unit performs the forward diffusion process, while the denoising unit performs the backward diffusion process. The forward diffusion process begins with the original data and gradually converts it into Gaussian noise through a series of denoising steps. This process can be considered a destructive process, with each step adding noise, making the data increasingly similar to random noise. The backward diffusion process begins with pure noise and gradually restores the original data through a series of denoising steps. This process effectively restores a clear image from the noise by learning a neural network to predict each denoising step.

[0090] Step S160, calculating a total loss based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image;

[0091] Step S170 , iteratively updating the initial image reconstruction model based on the total loss to obtain a trained target image reconstruction model.

[0092] It is worth noting that steps S110 to S170 are repeatedly performed, and each time the total loss is calculated, the parameter set of the image reconstruction model is updated. For example, after the total loss is obtained, the parameter set is updated using a gradient descent algorithm. This iterative update of the image reconstruction model is achieved until the number of iterations reaches a preset iteration threshold.

[0093] The image reconstruction model training method of the embodiment of the present application completes the training of the initial image reconstruction model through the above-mentioned steps S110 to S170, thereby obtaining a trained target image reconstruction model. First, the initial training image is input into the initial image reconstruction model; the initial image reconstruction model includes an initial diffusion network module and a pre-trained segmentation network module; image segmentation processing is performed by the pre-trained segmentation network module to obtain a segmentation mask probability training image, the segmentation mask probability training image and the initial training image are input into the denoising unit of the diffusion network module to obtain a noisy training image; the noisy training image is input into the denoising unit of the diffusion network module to obtain predicted training noise and a reconstructed training image; based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image, the total loss is calculated, and the initial image reconstruction model is iteratively updated based on the total loss to obtain a trained target image reconstruction model. In this process, since the segmentation network module is pre-trained, the accuracy of the segmentation network module is relatively high. In the segmentation mask probability training image output by the segmentation network module, the first pixel of the segmentation mask probability training image corresponds one-to-one to the second pixel of the initial training image, and the value of the first pixel represents the probability that the corresponding second pixel is an occluded area. Therefore, the information represented by the segmentation mask probability training image can be used as prior knowledge to help the diffusion network module learn the intrinsic mapping mechanism between the transparency of the occluded area and the information represented by the segmentation mask probability training image based on prior knowledge to accurately distinguish the edge and background pixels of the occluded area, thereby improving the ability to recognize occluded areas and overcome texture interference. Under the guidance of prior knowledge, the diffusion network module can deeply explore the background structure information that may exist in the occluded area, and while removing blurring objects such as flames and smoke, reasonably integrate this background information into the image reconstruction process. In this way, the diffusion network module can better preserve the background characteristics and logical structure of the original scene of the initial training image, help the reconstruction model solve the edge blur dilemma, improve the accuracy and effect of edge processing of the image reconstruction model, and thus significantly enhance the quality of image reconstruction. In this process, the information represented by the segmentation mask probability training image can divide the initial training image into occluded areas and non-occluded areas. In the denoising unit, the occluded areas are guided by noise, prompting the image reconstruction model to generate natural and continuous results in the subsequent denoising process.

[0094] In some embodiments, reference Figure 3 , Figure 3 This is a specific flow chart of step S140 in an embodiment of the present application. Step S140 may include but is not limited to step S310 and step S320.

[0095] Step S310, performing a pixel-by-pixel weighting operation on the initial training image and the segmentation mask probability training image to obtain a weighted training image;

[0096] In one embodiment, pixels in the initial training image correspond one-to-one with pixels in the segmentation mask probability training image. A pixel-by-pixel weighting operation is performed by multiplying the value of a first pixel in the initial training image by the value of a corresponding second pixel in the segmentation mask probability training image. This pixel-by-pixel weighting operation can reduce the contribution of occluded areas and add noise to background areas. The initial training image is segmented into known areas (non-occluded areas) and unknown areas (occluded areas) by the segmentation mask probability training image. In subsequent noise addition and denoising processes, the value of the first pixel in the segmentation mask probability training image in the weighted training image serves as a weight for the noise addition process, adding more noise to the unknown areas to learn how to complete them. Known areas are only slightly adjusted, ensuring that semantic information in the background of the initial training image is not lost due to excessive noise. The areas to be removed are guided by high noise, prompting the model to produce natural and continuous results in the subsequent denoising process.

[0097] Step S320: performing noise addition processing on the weighted training image by a noise addition unit to obtain a noisy training image.

[0098] In one embodiment, the noise adding process is expressed as:

[0099] ;

[0100] in is the noise attenuation coefficient; It is the Gaussian noise used to remove the occluded area. is the low-weight noise in the non-occluded area, Mprob represents the segmentation mask probability training image, I org Represent the initial training images, Characterize the noisy training image. In this way, more noise is added to the occluded area, causing it to quickly degenerate into pure noise, retaining more image information in the background area and reducing unnecessary noise. Through the noise addition process, since the first pixel of the segmentation mask probability training image corresponds one-to-one to the second pixel of the initial training image, and the value of the first pixel represents the probability that the corresponding second pixel is an occluded area, therefore, combined with the expression of the noise addition process, it can be seen that the value of the first pixel is equivalent to the weight of the noise. The noise addition process realizes the addition of Gaussian noise to the occluded area, and the greater the probability of being an occluded area, the greater the noise added, while low-weight noise is added to the non-occluded area. The noisy training image contains mixed information of the occluded area and the background area (i.e., the non-occluded area), so that the diffusion network module can better learn the features of eliminating the occluded area.

[0101] The present application first performs a pixel-by-pixel weighted operation on the initial training image and the segmentation mask probability training image to obtain a weighted training image through steps S310 to S320, and then performs noise processing on the weighted training image. In this way, the weighted training image can cover the probability information of the occluded area in the segmentation mask probability training image, so that the information represented by the segmentation mask probability training image output by the segmentation network module can be used as prior knowledge to help the diffusion network module learn the intrinsic mapping mechanism between the transparency of the occluded area and the information represented by the segmentation mask probability training image based on the prior knowledge to accurately distinguish the edge of the occluded area and the background pixels, so as to improve the occluded area recognition ability and overcome texture interference ability.

[0102] In some embodiments, reference Figure 4 , Figure 4 1 is a specific flow chart of step S150 of the embodiment of the present application. Step S150 may include but is not limited to steps S410 to S420.

[0103] Step S410, performing noise prediction processing on the noisy training image based on the denoising unit to obtain predicted training noise;

[0104] Step S420 : performing image reconstruction processing based on the predicted training noise and the noisy training image to obtain a reconstructed training image.

[0105] In one embodiment, the image reconstruction process is expressed as follows:

[0106] ;

[0107] ;

[0108] It is worth noting that Mprob represents the segmentation mask probability training image, I org Representing the initial training image, in the diffusion network module, the image reconstruction process is to perform t-step denoising processing, where The diffusion network module generates content based on the non-occluded area in step t-1; t-1 Represents the denoised image generated at step t-1, y t represents the denoised image generated at step t, and y t The image reconstruction process adds more noise to unknown (occluded) areas to learn how to complete them, while making only slight adjustments to known (non-occluded) areas to ensure that semantic information in the background is not lost due to excessive noise.

[0109] ,

[0110] The diffusion network module generates content based on the occluded area at step t-1. Characterizes the variance of the predicted training noise, Characterize the mean of the predicted training noise. The mean of the predicted training noise is expressed as:

[0111] ;

[0112] in, represents the prediction training noise, where is the noise attenuation coefficient; Mprob represents the segmentation mask probability training image, β t is the preset coefficient.

[0113] It is worth noting that, through steps S410 to S420, the denoising unit of the diffusion network module, after predicting the predicted training noise, performs denoising on the noisy training image based on the predicted training noise. This denoising process is also the image reconstruction process, thereby obtaining a reconstructed training image. The noise prediction process and image reconstruction process are based on the structure of the diffusion network module, which is not specifically limited in this application. The diffusion network module adopts the diffusion model DDPM.

[0114] In some embodiments, reference Figure 5 , Figure 5 This is a specific flow chart of step S160 of the embodiment of the present application. Step S160 may include but is not limited to steps S510 to S560.

[0115] Step S510, calculating the image reconstruction loss based on the reconstructed training image and the initial training image;

[0116] Step S520, calculating a denoising loss based on the predicted training noise;

[0117] Step S530 , calculating a mask consistency loss based on the predicted training noise and the segmentation mask probability training image;

[0118] Step S540, calculating the segmentation loss based on the segmentation mask probability training image and the true label image;

[0119] Step S550, calculating a sparse loss based on the segmentation mask probability training image;

[0120] In step S560 , a total loss is obtained based on the image reconstruction loss, the denoising loss, the mask consistency loss, the segmentation loss, and the sparsity loss.

[0121] It is worth noting that the embodiment of the present application calculates the image reconstruction loss, denoising loss, mask consistency loss, segmentation loss and sparse loss respectively through steps S510 to S560, thereby calculating the total loss, and performing iterative updates based on the total loss to obtain a trained target image reconstruction model. In this way, during the iterative update process, the sparse loss can constrain the sparsity of the intermediate representation and output of the image reconstruction model; the denoising loss can improve the prediction accuracy of the image reconstruction model for noise; the image reconstruction loss can improve the accuracy of the image reconstruction model; and the segmentation loss can improve the accuracy of the image reconstruction model in identifying occluded areas. In this way, by iteratively updating the image reconstruction model through multiple losses, the reconstruction quality of the image reconstruction model can be improved.

[0122] In one embodiment, the image reconstruction loss is calculated based on the reconstructed training image and the initial training image, including:

[0123] Based on the reconstructed training image, the initial training image and the reconstruction loss function, the image reconstruction loss is calculated; the reconstruction loss function is:

[0124] ;

[0125] in, L recon Characterize the image reconstruction loss, I output Representation reconstruction training images, I org Characterize the initial training images.

[0126] In one embodiment, the denoising loss is calculated based on the predicted training noise, including:

[0127] Based on the predicted training noise and the denoising loss function, the denoising loss is calculated; the denoising loss function is:

[0128] ;

[0129] in, L denoise represents the denoising loss; E represents the mathematical expectation; is Gaussian noise; Represents the predicted training noise, y0 represents the noisy training image; y t Mprob represents the noisy training image obtained by the diffusion network module at time step t; Mprob represents the segmentation mask probability training image. The denoising loss guides the model to gradually remove noise from the noisy image, which can improve the image reconstruction model's prediction accuracy for noise.

[0130] In one embodiment, the mask consistency loss is calculated based on the predicted training noise and the segmentation mask probability training image, including:

[0131] Based on the predicted training noise, segmentation mask probability training image and mask consistency loss function, the mask consistency loss is calculated; the mask consistency loss function is:

[0132] ;

[0133] in, L mask is the mask consistency loss, E represents the mathematical expectation; is Gaussian noise, Characterize the prediction training noise, y t Represents the noisy training image obtained by the diffusion network module at time step t; Mprob represents the segmentation mask probability training image, and y0 represents the noisy training image when time step is 0.

[0134] In one embodiment, the segmentation loss is calculated based on the segmentation mask probability training image and the true label image, including:

[0135] The segmentation loss is calculated based on the segmentation mask probability training image, the true label image and the segmentation loss function; the segmentation loss function is:

[0136] ;

[0137] in, L seg Representation segmentation loss, Characterize the probability of pixel i in the segmentation mask probability training image; represents the probability of pixel i in the true label image; N is the total number of pixels in the segmentation mask probability training image.

[0138] In one embodiment, the sparse loss is calculated based on the segmentation mask probability training image, including:

[0139] Based on the segmentation mask probability training image and the sparse loss function, the sparse loss is calculated; the sparse loss function is:

[0140] ;

[0141] in, L sparisty represents the sparse loss, and Mprob represents the segmentation mask probability training image.

[0142] In one embodiment, the total loss is obtained based on the image reconstruction loss, denoising loss, mask consistency loss, segmentation loss, and sparsity loss, which is specifically expressed as:

[0143] ;

[0144] Where L is the total loss, L sparisty Representation sparsity loss, L seg Representation segmentation loss, L mask is the mask consistency loss, L denoise Characterize denoising loss; L recon Characterize the image reconstruction loss. is the weight of denoising loss, is the weight of the mask consistency loss, is the weight of image reconstruction loss, is the weight of the sparse loss, is the weight of the segmentation loss. Those skilled in the art can set 、 、 、 、 The value of .

[0145] In some embodiments, the target image reconstruction model of the present application is used to assist in fire accident investigations and to restore and reconstruct fire images. Generally, flames and smoke coexist, so fire images typically have both flame and smoke regions. In this embodiment, an RGB image having both flame and smoke regions is selected as the initial training image. Specifically, before step S120, the initial training image is preprocessed. The preprocessing steps include:

[0146] Step S111, converting the initial training image from the three-primary color space to the hexagonal pyramid model color space;

[0147] It should be noted that a three-primary color space is a color space composed of three basic colors, typically red, green, and blue, corresponding to the three channels in the RGB color model. In the RGB color space, the intensity of each color can be represented by mixing different degrees of red, green, and blue. The Hexagonal Cone Color (HSV) color space is a color space created based on the intuitive nature of color, in which each color is represented by hue, saturation, and value.

[0148] Step S112: a region in the initial training image whose hue is within a first preset region, whose saturation is greater than a first preset value, and whose color brightness is greater than a second preset value is defined as a first region;

[0149] It is noteworthy that the first region represents the flame region in the initial training image. In one embodiment, the first preset interval is [0, 30], the first preset value is 150, and the second preset value is 200. Flames are characterized by a red or dark red color with high saturation and brightness. Therefore, the flame region can be determined using the first preset interval, the first preset value, and the second preset value. It should be noted that those skilled in the art may adjust the first preset interval, the first preset value, and the second preset value based on actual circumstances.

[0150] Step S113, a region in the initial training image whose hue is within the second preset range, whose saturation is less than the third preset value, and whose color lightness is less than the fourth preset value is used as the second region;

[0151] It is worth noting that the second region represents the smoke region in the initial training image. In one embodiment, the second preset interval is [90, 120], the third preset value is 50, and the fourth preset value is 50. The smoke region is characterized by gray, low saturation, and medium-low brightness. Therefore, based on the second preset interval, the third preset value, and the fourth preset value, the smoke region can be determined. It should be noted that those skilled in the art may adjust the second preset interval, the third preset value, and the fourth preset value according to actual circumstances.

[0152] Step S114, subtracting the first region and the second region from the initial training image to obtain a third region;

[0153] It is worth noting that the third area is the background area outside the flame area and the smoke area.

[0154] Step S115, increasing the saturation and color brightness of the first region to obtain a first updated region;

[0155] Specifically, increasing the saturation and color brightness of the first area can make the flame area appear more vivid red / orange, and the saturation is significantly improved.

[0156] Exemplarily, the saturation of a pixel in the first update region = the saturation of the corresponding pixel in the first region + 30. The color brightness of a pixel in the first update region = the color brightness of the corresponding pixel in the first region + 10.

[0157] Step S116, reducing the hue and saturation of the second region and increasing the color brightness of the second region to obtain a second updated region;

[0158] Specifically, reducing the hue and saturation of the second area can weaken the color interference in the smoke area and reduce color interference, while increasing the color brightness of the second area can highlight the smoke texture.

[0159] For example, the hue of a pixel in the second update region = the hue of the corresponding pixel in the second region - 10; the saturation of a pixel in the second update region = the saturation of the corresponding pixel in the second region * 0.3; and the lightness of a pixel in the second update region = the lightness of the corresponding pixel in the second region + 10.

[0160] Step S117, merging the third area, the first updated area, and the second updated area to obtain a merged image;

[0161] Step S118 : converting the merged image from the hexagonal pyramid model color space into the three-primary color space to obtain a new initial training image.

[0162] During the preprocessing process of this embodiment, steps S111 to S118 are performed to achieve a more vivid red / orange color in the flame region of the initial training image, significantly improving saturation. Color interference in occluded regions is reduced, resulting in clearer textures. Inputting this new initial training image into the image reconstruction model helps improve the segmentation network module's ability to distinguish between smoke and flame regions, thereby improving the quality of the reconstructed image output by the target image reconstruction model.

[0163] The second embodiment of the present application provides an image reconstruction method. Figure 6 , Figure 6 1 is a flow chart of the steps of the image reconstruction method according to an embodiment of the present application. The image reconstruction method includes but is not limited to step S610 and step S620.

[0164] Step S610, obtaining an image to be reconstructed with an occluded area;

[0165] It is worth noting that an object with flames or smoke can be photographed by a camera to obtain an image to be reconstructed with an occluded area, where the occluded area is the flame part or the smoke part.

[0166] Step S620: input the image to be reconstructed into a target image reconstruction model to obtain a target reconstructed image; wherein, the target image reconstruction model is trained according to the image reconstruction model training method of any one of the embodiments of the first aspect.

[0167] It is worth noting that the image reconstruction model can be Figure 1Specifically, after the image to be reconstructed is input into the trained target image reconstruction model, the image to be reconstructed is segmented by the segmentation network module of the image reconstruction model to obtain a segmentation mask image; the segmentation mask image and the image to be reconstructed are weighted pixel by pixel to obtain a weighted image; the weighted image is denoised by the denoising unit of the diffusion network module of the image reconstruction model to obtain a noisy image; the noisy image is noise-predicted by the denoising unit of the diffusion network module to obtain predicted noise; and then the noisy image is denoised based on the predicted noise and the noisy image to obtain a target reconstructed image.

[0168] It should be noted that since the image reconstruction model in the image reconstruction method of the second embodiment of the application is obtained by the image reconstruction model training method of the first embodiment, the specific implementation method of the image reconstruction method of the second embodiment of the application is basically the same as the specific implementation method of the above-mentioned image reconstruction model training method, and has the same beneficial effects, and will not be repeated here.

[0169] Some other embodiments of the present application further provide an image reconstruction device, referring to Figure 7 , Figure 7 Schematic diagram of the structure of an image reconstruction device according to an embodiment of the present application. The image reconstruction device includes:

[0170] The to-be-reconstructed image acquisition module 710 is used to acquire the to-be-reconstructed image with the occluded area;

[0171] The reconstruction module 720 is used to input the image to be reconstructed into the target image reconstruction model to obtain a target reconstructed image; wherein the target image reconstruction model is trained according to the image reconstruction model training method of any one of the embodiments of the first aspect.

[0172] The image reconstruction device of the embodiment of the present application is used to execute the image reconstruction method of the second embodiment of the present application. When executing the method, after the image to be reconstructed is input into the trained target image reconstruction model, the image to be reconstructed is subjected to image segmentation processing by the segmentation network module of the image reconstruction model to obtain a segmentation mask image; the segmentation mask image and the image to be reconstructed are subjected to pixel-by-pixel weighting processing to obtain a weighted image; the weighted image is subjected to noise processing by the noise adding unit of the diffusion network module of the image reconstruction model to obtain a noisy image; the noisy image is subjected to noise prediction processing by the denoising unit of the diffusion network module to obtain predicted noise, and then the noisy image is subjected to denoising processing based on the predicted noise and the noisy image to obtain a target reconstructed image.

[0173] The third embodiment of the present application provides an image reconstruction model training device. Figure 7 , Figure 7Schematic diagram of the structure of the image reconstruction model training device of the embodiment of the present application. The image reconstruction model training device includes:

[0174] An acquisition module 810 is configured to acquire an initial training image having an occluded area and a true label image corresponding to the initial training image;

[0175] A first input module 820 is used to input the initial training image into the initial image reconstruction model; the initial image reconstruction model includes an initial diffusion network module and a pre-trained segmentation network module;

[0176] A segmentation module 830 is configured to perform image segmentation processing on the initial training image using the segmentation network module to obtain a segmentation mask probability training image; wherein a first pixel of the segmentation mask probability training image corresponds one-to-one to a second pixel of the initial training image, and the value of the first pixel represents the probability that the corresponding second pixel is an occluded area;

[0177] The second input module 840 is used to input the segmentation mask probability training image and the initial training image into the noise adding unit of the diffusion network module to obtain a noisy training image;

[0178] The third input module 850 is used to input the noisy training image into the denoising unit of the diffusion network module to obtain the predicted training noise and the reconstructed training image;

[0179] A calculation module 860 is configured to calculate a total loss based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image;

[0180] The iterative update module 880 is used to iteratively update the initial image reconstruction model based on the total loss to obtain a trained target image reconstruction model.

[0181] It is worth noting that the image reconstruction model training device is used to execute the image reconstruction model training method of the first aspect of the embodiment of the present application. When executing the method, the initial training image is first input into the initial image reconstruction model; the initial image reconstruction model includes an initial diffusion network module and a pre-trained segmentation network module; image segmentation processing is performed by the pre-trained segmentation network module to obtain a segmentation mask probability training image, and the segmentation mask probability training image and the initial training image are input into the denoising unit of the diffusion network module to obtain a noisy training image; the noisy training image is input into the denoising unit of the diffusion network module to obtain predicted training noise and reconstructed training image; based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image, the total loss is calculated, and the initial image reconstruction model is iteratively updated based on the total loss to obtain a trained target image reconstruction model. In this process, since the segmentation network module is pre-trained, the accuracy of the segmentation network module is relatively high. In the segmentation mask probability training image output by the segmentation network module, the first pixel of the segmentation mask probability training image corresponds one-to-one to the second pixel of the initial training image, and the value of the first pixel represents the probability that the corresponding second pixel is an occluded area. Therefore, the information represented by the segmentation mask probability training image can be used as prior knowledge to help the diffusion network module learn the intrinsic mapping mechanism between the transparency of the occluded area and the information represented by the segmentation mask probability training image based on prior knowledge to accurately distinguish the edge and background pixels of the occluded area, thereby improving the ability to recognize occluded areas and overcome texture interference. Under the guidance of prior knowledge, the diffusion network module can deeply explore the background structure information that may exist in the occluded area, and while removing blurring objects such as flames and smoke, reasonably integrate this background information into the image reconstruction process. In this way, the diffusion network module can better preserve the background characteristics and logical structure of the original scene of the initial training image, help the reconstruction model solve the edge blur dilemma, improve the accuracy and effect of edge processing of the image reconstruction model, and thus significantly enhance the quality of image reconstruction.

[0182] In some embodiments, the second input module 840 includes a pixel-by-pixel weighting submodule and a noise adding submodule.

[0183] The pixel-by-pixel weighting submodule is used to perform pixel-by-pixel weighting operations on the initial training image and the segmentation mask probability training image to obtain a weighted training image;

[0184] The denoising submodule is used to perform denoising processing on the weighted training image through the denoising unit to obtain a noisy training image.

[0185] In some embodiments, the third input module 850 includes a prediction submodule and a reconstruction submodule.

[0186] The prediction submodule is used to perform noise prediction processing on the noisy training image based on the denoising unit to obtain predicted training noise;

[0187] The reconstruction submodule is used to perform image reconstruction processing based on the predicted training noise and the noisy training image to obtain a reconstructed training image.

[0188] In some embodiments, the computing module 860 includes a first computing submodule, a second computing submodule, a third computing submodule, a fourth computing submodule, a fifth computing submodule, and a sixth computing submodule.

[0189] The first calculation submodule is used to calculate the image reconstruction loss based on the reconstructed training image and the initial training image;

[0190] The second calculation submodule is used to calculate the denoising loss based on the predicted training noise;

[0191] The third calculation submodule is used to calculate the mask consistency loss based on the predicted training noise and the segmentation mask probability training image;

[0192] The fourth calculation submodule is used to calculate the segmentation loss based on the segmentation mask probability training image and the true label image;

[0193] The fifth calculation submodule is used to calculate the sparse loss based on the segmentation mask probability training image;

[0194] The sixth calculation submodule is used to obtain the total loss based on image reconstruction loss, denoising loss, mask consistency loss, segmentation loss and sparse loss.

[0195] In some embodiments, the first calculation submodule is specifically used to calculate the image reconstruction loss based on the reconstructed training image, the initial training image and the reconstruction loss function.

[0196] In some embodiments, the second calculation submodule is specifically configured to calculate the denoising loss based on the predicted training noise and the denoising loss function.

[0197] In some embodiments, the third calculation submodule is specifically configured to calculate the mask consistency loss based on the predicted training noise, the segmentation mask probability training image, and the mask consistency loss function.

[0198] In some embodiments, the fourth calculation submodule is specifically used to calculate the segmentation loss based on the segmentation mask probability training image, the true label image and the segmentation loss function.

[0199] A fourth embodiment of the present application provides an electronic device, comprising a memory and a processor, wherein the memory stores a computer program, and the processor, when executing the computer program, implements the image reconstruction model training method of any one of the first embodiment or the image reconstruction method of the second embodiment. The electronic device can be any smart terminal, such as a tablet computer or an in-vehicle computer.

[0200] Reference Figure 9 , Figure 9 FIG. 1 is a schematic diagram of the structure of an electronic device according to an embodiment. The electronic device includes:

[0201] The processor 901 can be implemented as a general-purpose CPU (Central Processing Unit), a microprocessor, an application-specific integrated circuit (ASIC), or one or more integrated circuits, and is used to execute relevant programs to implement the technical solutions provided in the embodiments of the present application;

[0202] The memory 902 can be implemented in the form of a read-only memory (ROM), a static storage device, a dynamic storage device, or a random access memory (RAM). The memory 902 can store an operating system and other application programs. When the technical solutions provided in the embodiments of this specification are implemented through software or firmware, the relevant program code is stored in the memory 902 and is called by the processor 901 to execute the image reconstruction model training method or image reconstruction method of the embodiments of this application;

[0203] Input / output interface 903, used to implement information input and output;

[0204] Communication interface 904, used to implement communication interaction between this device and other devices, which can be achieved through wired means (such as USB, network cable, etc.) or wireless means (such as mobile network, WiFi, Bluetooth, etc.);

[0205] Bus 905 , which transmits information between various components of the device (e.g., processor 901 , memory 902 , input / output interface 903 , and communication interface 904 );

[0206] The processor 901 , the memory 902 , the input / output interface 903 and the communication interface 904 are connected to each other in communication within the device via a bus 905 .

[0207] The fifth aspect of the present application is a computer-readable storage medium, which stores a computer program. When the computer program is executed by a processor, it implements the image reconstruction model training method of any one of the first aspect embodiments or the image reconstruction method of the second aspect embodiments.

[0208] The memory, as a non-transient computer-readable storage medium, can be used to store non-transient software programs and non-transient computer executable programs. In addition, the memory may include a high-speed random access memory and may also include a non-transient memory, such as at least one disk storage device, a flash memory device, or other non-transient solid-state storage device. In some embodiments, the memory may optionally include a memory remotely arranged relative to the processor, and these remote memories may be connected to the processor via a network. Examples of the above-mentioned network include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0209] The embodiments described in the embodiments of this application are intended to more clearly illustrate the technical solutions of the embodiments of this application and do not constitute a limitation on the technical solutions provided by the embodiments of this application. Those skilled in the art will appreciate that with the evolution of technology and the emergence of new application scenarios, the technical solutions provided in the embodiments of this application are also applicable to similar technical problems.

[0210] Those skilled in the art will understand that the technical solutions shown in the figures do not constitute a limitation on the embodiments of the present application, and may include more or fewer steps than shown in the figures, or a combination of certain steps, or different steps.

[0211] The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate, i.e., they may be located in one place or distributed across multiple network units. Some or all of the modules may be selected based on actual needs to achieve the objectives of this embodiment.

[0212] Those skilled in the art will appreciate that all or some of the steps in the methods, systems, and functional modules / units in the devices disclosed above may be implemented as software, firmware, hardware, or appropriate combinations thereof.

[0213] The terms "first", "second", "third", "fourth", etc. (if any) in the specification of the present application and the above-mentioned drawings are used to distinguish similar objects and are not necessarily used to describe a specific order or sequential order. It should be understood that the data used in this way can be interchangeable where appropriate, so that the embodiments of the present application described herein can be implemented in an order other than those illustrated or described herein. In addition, the terms "including" and "having" and any variations thereof are intended to cover non-exclusive inclusions, for example, a process, method, system, product or device that includes a series of steps or units is not necessarily limited to those steps or units clearly listed, but may include other steps or units that are not clearly listed or inherent to these processes, methods, products or devices.

[0214] It should be understood that in the present application, "at least one (item)" means one or more, and "plurality" means two or more. "And / or" is used to describe the mapping relationship of the mapping objects, indicating that three relationships may exist. For example, "A and / or B" can mean: only A exists, only B exists, and A and B exist at the same time, where A and B can be singular or plural. The character " / " generally indicates that the previous and next mapping objects are in an "or" relationship. "At least one of the following items" or similar expressions refers to any combination of these items, including any combination of single items or plural items. For example, at least one of a, b or c can mean: a, b, c, "a and b", "a and c", "b and c", or "a and b and c", where a, b, c can be single or multiple.

[0215] In the several embodiments provided in this application, it should be understood that the disclosed systems and methods can be implemented in other ways. For example, the system embodiments described above are merely illustrative. For example, the division of the above units is merely a logical function division. In actual implementation, there may be other division methods, such as multiple units or components can be combined or integrated into another system, 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 an indirect coupling or communication connection through some interfaces, devices or units, which can be electrical, mechanical or other forms.

[0216] The units described above as separate components may or may not be physically separate, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed across multiple network units. Some or all of these units may be selected according to actual needs to achieve the purpose of the solution of this embodiment.

[0217] In addition, the functional units in the various embodiments of the present application may be integrated into a single processing unit, or each unit may exist physically separately, or two or more units may be integrated into a single unit. The aforementioned integrated units may be implemented in the form of hardware or software functional units.

[0218] If the integrated unit is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present application, or the part that contributes to the prior art, or all or part of the technical solution can be embodied in the form of a software product, which is stored in a storage medium and includes multiple instructions for enabling a computer device (which can be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods of various embodiments of the present application. The aforementioned storage medium includes: various media that can store programs, such as a USB flash drive, a mobile hard disk, a read-only memory (ROM), a random access memory (RAM), a magnetic disk or an optical disk.

[0219] The preferred embodiments of the present invention are described above with reference to the accompanying drawings, but are not intended to limit the scope of the present invention. Any modifications, equivalent substitutions, and improvements made by those skilled in the art without departing from the scope and essence of the present invention should be within the scope of the present invention.

Claims

1. A method for training an image reconstruction model, characterized in that: include: Obtaining an initial training image with an occluded area and a true label image corresponding to the initial training image; Inputting the initial training image into an initial image reconstruction model; the initial image reconstruction model includes an initial diffusion network module and a pre-trained segmentation network module; Performing image segmentation processing on the initial training image using the segmentation network module to obtain a segmentation mask probability training image; wherein a first pixel of the segmentation mask probability training image corresponds one-to-one to a second pixel of the initial training image, and a value of the first pixel represents a probability that the corresponding second pixel is an occluded area; Inputting the segmentation mask probability training image and the initial training image into the noise adding unit of the diffusion network module to obtain a noisy training image; Inputting the noisy training image into the denoising unit of the diffusion network module to obtain predicted training noise and reconstructed training image; Calculating a total loss based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image; Iteratively updating the initial image reconstruction model based on the total loss to obtain a trained target image reconstruction model; The step of inputting the segmentation mask probability training image and the initial training image into the noise adding unit of the diffusion network module to obtain the noise adding training image comprises: Performing a pixel-by-pixel weighting operation on the initial training image and the segmentation mask probability training image to obtain a weighted training image; Performing noise addition processing on the weighted training image by the noise addition unit to obtain the noisy training image; The step of inputting the noisy training image into the denoising unit of the diffusion network module to obtain predicted training noise and reconstructed training image comprises: performing noise prediction processing on the noisy training image based on the denoising unit to obtain predicted training noise; Performing image reconstruction processing based on the predicted training noise and the noisy training image to obtain the reconstructed training image; The total loss is calculated based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image, including: Calculating an image reconstruction loss based on the reconstructed training image and the initial training image; Calculating a denoising loss based on the predicted training noise; Calculating a mask consistency loss based on the predicted training noise and the segmentation mask probability training image; Calculating a segmentation loss based on the segmentation mask probability training image and the true label image; Calculating a sparse loss based on the segmentation mask probability training image; The total loss is obtained based on the image reconstruction loss, the denoising loss, the mask consistency loss, the segmentation loss and the sparsity loss.

2. The image reconstruction model training method according to claim 1, characterized in that: The calculating the image reconstruction loss based on the reconstructed training image and the initial training image includes: The image reconstruction loss is calculated based on the reconstructed training image, the initial training image and a reconstruction loss function; the reconstruction loss function is: ； in, L recon Characterize the image reconstruction loss, I output characterizing the reconstructed training image, I org The initial training image is characterized.

3. The image reconstruction model training method according to claim 1, characterized in that: The calculating the denoising loss based on the predicted training noise includes: The denoising loss is calculated based on the predicted training noise and the denoising loss function; the denoising loss function is: ； in, L denoise represents the denoising loss; E represents the mathematical expectation; is Gaussian noise; represents the predicted training noise, y0 represents the noisy training image; y t represents the noisy training image obtained by the diffusion network module at time step t; Mprob represents the segmentation mask probability training image.

4. The image reconstruction model training method according to claim 1, characterized in that The calculating the mask consistency loss based on the predicted training noise and the segmentation mask probability training image includes: The mask consistency loss is calculated based on the predicted training noise, the segmentation mask probability training image and the mask consistency loss function; the mask consistency loss function is: ； in, L mask is the mask consistency loss, E represents the mathematical expectation; is Gaussian noise, Characterize the prediction training noise, y t represents the noisy training image obtained by the diffusion network module at time step t; Mprob represents the segmentation mask probability training image, and y0 represents the noisy training image when the time step is 0.

5. The image reconstruction model training method according to claim 1, characterized in that: The calculating the segmentation loss based on the segmentation mask probability training image and the true label image includes: The segmentation loss is calculated based on the segmentation mask probability training image, the true label image and the segmentation loss function; the segmentation loss function is: ； in, L seg Characterize the segmentation loss, Characterize the probability of pixel i in the segmentation mask probability training image; represents the probability of pixel i in the true label image; N is the total number of pixels in the segmentation mask probability training image.

6. An image reconstruction method, characterized in that: include: Acquire an image to be reconstructed with an occluded area; The image to be reconstructed is input into a target image reconstruction model to obtain a target reconstructed image; wherein the target image reconstruction model is trained according to the image reconstruction model training method according to any one of claims 1 to 5.

7. An image reconstruction model training device, characterized in that: include: An acquisition module, configured to acquire an initial training image having an occluded area and a true label image corresponding to the initial training image; A first input module is used to input the initial training image into an initial image reconstruction model; the initial image reconstruction model includes an initial diffusion network module and a pre-trained segmentation network module; a segmentation module configured to perform image segmentation processing on the initial training image using the segmentation network module to obtain a segmentation mask probability training image; wherein a first pixel of the segmentation mask probability training image corresponds one-to-one to a second pixel of the initial training image, and a value of the first pixel represents a probability that the corresponding second pixel is an occluded area; A second input module is configured to input the segmentation mask probability training image and the initial training image into the noise adding unit of the diffusion network module to obtain a noisy training image; A third input module is used to input the noisy training image into the denoising unit of the diffusion network module to obtain predicted training noise and reconstructed training image; A calculation module, configured to calculate a total loss based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image; an iterative update module, configured to iteratively update the initial image reconstruction model based on the total loss to obtain a trained target image reconstruction model; The step of inputting the segmentation mask probability training image and the initial training image into the noise adding unit of the diffusion network module to obtain the noise adding training image comprises: Performing a pixel-by-pixel weighting operation on the initial training image and the segmentation mask probability training image to obtain a weighted training image; Performing noise addition processing on the weighted training image by the noise addition unit to obtain the noisy training image; The step of inputting the noisy training image into the denoising unit of the diffusion network module to obtain predicted training noise and reconstructed training image comprises: performing noise prediction processing on the noisy training image based on the denoising unit to obtain predicted training noise; Performing image reconstruction processing based on the predicted training noise and the noisy training image to obtain the reconstructed training image; The total loss is calculated based on the reconstructed training image, the initial training image, the predicted training noise, the segmentation mask probability training image, and the true label image, including: Calculating an image reconstruction loss based on the reconstructed training image and the initial training image; Calculating a denoising loss based on the predicted training noise; Calculating a mask consistency loss based on the predicted training noise and the segmentation mask probability training image; Calculating a segmentation loss based on the segmentation mask probability training image and the true label image; Calculating a sparse loss based on the segmentation mask probability training image; The total loss is obtained based on the image reconstruction loss, the denoising loss, the mask consistency loss, the segmentation loss and the sparsity loss.

8. An electronic device, characterized in that: The electronic device includes a memory and a processor, the memory stores a computer program, and when the processor executes the computer program, it implements the image reconstruction model training method according to any one of claims 1 to 5, or the image reconstruction method according to claim 6.

9. A computer-readable storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, it implements the image reconstruction model training method according to any one of claims 1 to 5, or the image reconstruction method according to claim 6.

Images