Training methods and apparatus for image generation models, as well as image generation methods and apparatus.

Abstract

This application discloses a training method and apparatus for an image generation model, as well as an image generation method and apparatus. The training method includes: training a main branch using a first defect image and a coarse mask, while masking the processing of a sub-branch; after completing the main branch training, training a sub-branch using a second defect image, a coarse mask, and a fine mask; during sub-branch training, the second defect image and the fine mask are input into the sub-branch, processed by a second VAE encoder and a second UNET encoder to obtain the output of the second UNET encoder; the second defect image and the coarse mask are input into the main branch, and the output of the second UNET encoder is input into a first UNET decoder. After processing, the predicted noise output by the UNET decoder is obtained, and the model parameters of the sub-branch are updated based on the predicted noise. Applying this application, the trained model can generate defect images with diversity in batches, while enhancing the controllability of defects.

Description

Technical Field

[0001] This application relates to industrial defect detection technology, and in particular to training methods and apparatus for image generation models, as well as image generation methods and apparatus. Background Technology

[0002] Industrial defect detection is an intelligent quality inspection method that uses computer vision and artificial intelligence technologies to automatically identify surface or internal defects (such as cracks, scratches, deformations, etc.) in products through trained deep learning models. However, training deep learning models requires a large number of image samples containing "defects," and in industrial defect detection, the scarcity of negative samples (defect samples) is a common challenge.

[0003] To address the scarcity of defect samples, image generation models are commonly used to generate images containing "defects." Most of the latest solutions are based on diffusion models as their fundamental structure. Within diffusion model technology, existing solutions fall into two categories: the first is a coarse masking scheme, which only requires a coarse mask of the defect location during inference, thus reducing annotation difficulty and allowing for the mass production of more diverse defect images; however, it lacks precise control over defect size and shape. The second is a fine masking scheme, which inputs a fine, pixel-level segmentation mask during inference, resulting in generated defects that closely adhere to the mask, offering high accuracy and controllability. However, the annotation cost of fine masks is relatively high, hindering the mass production of diverse defect images when annotation costs are limited. Summary of the Invention

[0004] This application provides a training method and apparatus for an image generation model, as well as an image generation method and apparatus. The trained model can generate images with diverse defects in batches, while enhancing the controllability of defects.

[0005] To achieve the above objectives, this application adopts the following technical solution:

[0006] A training method for an image generation model, the image generation model including a main branch and a sub-branch, the training method comprising:

[0007] The main branch is trained using the first defect image and its corresponding coarse mask, and the training process of the main branch masks the processing of the sub-branch; wherein, the main branch is a diffusion model, including a first VAE encoder, a first UNET encoder, a UNET decoder and a VAE decoder;

[0008] After the main branch is trained, the sub-branch is trained using the second defect image and its corresponding coarse and fine masks; the sub-branch includes a second VAE encoder and a second UNET encoder.

[0009] The training of the secondary branch includes:

[0010] The second defect image and its corresponding fine mask are input into the sub-branch, and after processing by the second VAE encoder and the second UNET encoder, the output of the second UNET encoder is obtained.

[0011] The second defect image and its corresponding coarse mask are input into the main branch, and the output of the second UNET encoder is input into the UNET decoder. The predicted noise output by the UNET decoder is obtained after processing by the main branch, and the model parameters of the sub-branch are updated based on the comparison between the predicted noise and the standard noise.

[0012] Preferably, the image generation model includes a text processing branch;

[0013] When training the main branch, the defect description text of the first defect image is input into the text processing branch for text feature extraction, and the extracted text features are used to guide the processing of the first UNET encoder.

[0014] When training the sub-branch, the defect description text of the second defect image is input into the text processing branch for text feature extraction, and the extracted text features are used to guide the processing of the second UNET encoder and the first UNET encoder.

[0015] Preferably, the initial parameters of the second UNET encoder are the same as the parameters of the first UNET encoder after the main branch training is completed.

[0016] Preferably, the method further includes: processing the output of the second UNET encoder with a zero convolutional layer and then inputting it into the UNET decoder.

[0017] Preferably, inputting the output of the second UNET encoder to the UNET decoder includes: inputting the output of each layer of the second UNET encoder to the UNET decoder;

[0018] In the UNET decoder, the outputs of each layer of the second UNET encoder are merged with the outputs of the same layer of the first UNET encoder for decoding processing of the corresponding layer of the UNET decoder.

[0019] Preferably, for the second defect image, its corresponding fine mask is randomly expanded to obtain a corresponding coarse mask.

[0020] An image generation method, comprising:

[0021] Input the normal image and the mask image describing the defect location into the image generation model trained by any of the above-mentioned methods for inference processing to generate a defect image;

[0022] Among them, the output of the second UNET encoder is multiplied by scale before inputting into the UNET decoder, and scale is the participation ratio of the fine mask set in advance according to requirements, where 0 ≤ scale ≤ 1.

[0023] Preferably, when scale = 0, the mask image input into the image generation model is a rough mask image;

[0024] When 0 < scale ≤ 1, the mask image input into the image generation model is a fine mask image and a rough mask image, and the masked part in the rough mask image includes the masked part in the fine mask image.

[0025] A training device for an image generation model, the image generation model includes a main branch and a sub-branch, and the training device includes: a main branch training unit and a sub-branch training unit;

[0026] The main branch training unit is used to train the main branch by using the first defect image and its corresponding rough mask, and the training process of the main branch shields the processing of the sub-branch; among them, the main branch is a diffusion model, including a first VAE encoder, a first UNET encoder, a UNET decoder and a VAE decoder;

[0027] The sub-branch training unit is used to train the sub-branch by using the second defect image and its corresponding rough mask and fine mask after the training of the main branch is completed; the sub-branch includes a second VAE encoder and a second UNET encoder;

[0028] Among them, the training of the sub-branch includes:

[0029] Input the second defect image and its corresponding fine mask into the sub-branch, and after being processed by the second VAE encoder and the second UNET encoder, obtain the output of the second UNET encoder;

[0030] Input the second defect image and its corresponding rough mask into the main branch, input the output of the second UNET encoder into the UNET decoder, and obtain the predicted noise output by the UNET decoder after being processed by the main branch. Update the model parameters of the sub-branch based on the comparison between the predicted noise and the standard noise.

[0031] Preferably, the training device includes a text processing unit, configured to input the defect description text of the first defect image into the text processing branch for text feature extraction, and use the extracted text features to guide the processing of the first UNET encoder and the UNET decoder in the main branch training unit; and further configured to input the defect description text of the second defect image into the text processing branch for text feature extraction, and use the extracted text features to guide the processing of the second UNET encoder, the first UNET encoder, and the UNET decoder in the sub-branch training unit.

[0032] Preferably, the training device further includes a zero-convolutional layer processing unit for processing the output of the second UNET encoder with a zero-convolutional layer before inputting it into the UNET decoder.

[0033] An image generation apparatus includes: an input module and an inference module for an image generation model;

[0034] The input module is used to input a normal image and a mask image describing the location of the defect into the image generation model output by the training device of the above image generation model;

[0035] The inference module of the image generation model is used to perform inference based on the normal image and the masked image using the image generation model to generate a defective image.

[0036] The output of the second UNET encoder is multiplied by a scale before being input to the UNET decoder. The scale is the participation ratio of a fine mask that is set in advance according to the requirements, where 0 ≤ scale ≤ 1.

[0037] As can be seen from the above technical solution, the image generation model in this application has a dual-branch structure, including a main branch and a sub-branch. The main branch includes a first VAE encoder, a first UNET encoder, a UNET decoder, and a VAE decoder, corresponding to a complete diffusion model. The sub-branch includes a second VAE encoder and a second UNET encoder, used to introduce fine mask information into the main branch, guiding the UNET decoding process of the main branch and improving the controllability of defects. A two-stage training strategy is adopted during model training. First, the main branch is trained using defect images and their corresponding coarse masks, enabling the model to generate diverse defect images in batches. After the main branch training is completed, the sub-branch is trained using defect images and their corresponding coarse and fine masks. When training the sub-branch, the defect image and fine mask are input into the sub-branch for processing to obtain the output of the second UNET encoder. The defect image and coarse mask are then input into the main branch for processing, and the output of the second UNET encoder is input into the UNET decoder of the main branch, so that the fine mask participates in the UNET decoding process of the main branch. The sub-branch parameters are updated based on the comparison between the predicted noise and standard noise output by the UNET decoder. Because the training of the secondary branch incorporates fine-grained masking information during the UNET decoding process of the main branch, the generated model effectively enhances the controllability of the generated defects. In summary, this application, through a dual-branch structure combined with a two-stage training strategy, enables the trained model to generate diverse defect images in batches, while also enhancing the controllability of the defects.

[0038] Furthermore, using the image generation model trained above, inference can be performed on the input normal image and the mask image describing the defect location to generate a defect image. During inference, the output of the second UNET encoder in the sub-branch is multiplied by a scale before entering the UNET decoder in the main branch, thus controlling the degree of involvement of the fine mask in the model inference through the scale value. Therefore, the above image model generation method can balance the diversity and accuracy of generated defects using scale. Reducing the scale factor, combined with a coarse mask, can generate diverse defect samples; while increasing the scale factor, combined with a fine mask, generates defect samples of a specified shape and size. Attached Figure Description

[0039] Figure 1 This is a schematic diagram of the basic structure of the image generation model in this application;

[0040] Figure 2 This is a schematic diagram of the basic process of the image generation model training method in this application;

[0041] Figure 3 This is a schematic diagram of the basic training process of the sub-branch in the image generation model training method of this application;

[0042] Figure 4 This is a schematic diagram of the basic process of the image generation method in this application;

[0043] Figure 5 This is a schematic diagram illustrating the specific process of training the image generation model in Embodiment 1 of this application;

[0044] Figure 6 This is a schematic diagram of the training structure of the main branch in Embodiment 1 of this application;

[0045] Figure 7 This is a schematic diagram of the training structure of the sub-branch in Embodiment 1 of this application;

[0046] Figure 8 This is a schematic diagram illustrating the reasoning of the image generation model in Embodiment 2 of this application;

[0047] Figure 9 This is a schematic diagram of the basic process of the image generation method in Embodiment 2 of this application;

[0048] Figure 10a , Figure 10b and Figure 10c These are schematic diagrams illustrating the defect image generation status for three different scale values.

[0049] Figure 11 This is a schematic diagram of the basic structure of the training device for the image generation model in this application;

[0050] Figure 12 This is a schematic diagram of the basic structure of the image generation device in this application. Detailed Implementation

[0051] To make the objectives, technical means, and advantages of this application clearer, the following detailed description is provided in conjunction with the accompanying drawings.

[0052] Current defect image generation schemes based on coarse masks typically reuse defect detection label data, using detected bounding boxes as coarse masks to ensure the mask covers the defect. However, the shape and size of the mask do not need to strictly match the defect's shape and size. Furthermore, defect description text can be combined to generate defects within the bounding boxes, thus introducing a degree of randomness. The advantages of this scheme are that it can reuse defect detection label data to train the model, eliminating the need for manually annotating fine segmentation labels, and generating diverse defects. During inference, only the bounding box mask is needed, making annotation relatively easy and allowing for large-scale generation. The disadvantages of this scheme are that it cannot precisely control the size and shape of defects, often requiring manual selection to obtain the desired defect samples.

[0053] In defect image generation schemes based on fine-grained masks, a fine-grained, pixel-level segmentation mask is input. Therefore, the defects generated by the model closely adhere to the mask, resulting in high accuracy and controllability. The advantage of this scheme is its ability to precisely control the shape and size of defects. However, its disadvantage is the need for manual annotation of the fine-grained segmentation mask. Given limited annotation costs, it is difficult to generate a large number of diverse defect image samples.

[0054] Based on the problems of coarse and fine masks, this application innovatively proposes a dual-branch model structure that includes coarse and fine mask processing, combined with a two-stage training strategy, to train a defect image generation model.

[0055] Figure 1 This is a schematic diagram of the basic structure of the image generation model in this application. Figure 2 This is a schematic diagram illustrating the basic workflow of an image generation model training method. (Example:) Figure 1 As shown, the image generation model, in its most basic form, includes a main branch and a sub-branch. Additionally, to better control defect generation, the image generation model may include a text processing section (shown in the dashed box) for providing textual descriptions of defects and offering extra defect information to guide model generation. The main branch is specifically a diffusion model, including at least a VAE encoder, a UNET encoder, a UNET decoder, and a VAE decoder. To distinguish them from the VAE encoder and UNET encoder in the sub-branch, the VAE encoder and UNET encoder in the main branch will be referred to as the first VAE encoder and the first UNET encoder, respectively. The sub-branch includes at least a VAE encoder and a UNET encoder, which will be referred to as the second VAE encoder and the second UNET encoder, respectively. Optionally, the first VAE encoder and the second VAE encoder can be the same VAE encoder. Figure 1 and Figure 2 As shown, the basic process of training the image generation model in this application includes:

[0056] Step 201: Train the main branch using the first defect image and its corresponding coarse mask.

[0057] The training in this application adopts a two-stage training strategy: first, the main branch is trained separately, and the secondary branch is not involved in the training process, that is, the processing of the secondary branch is shielded; after the main branch is trained, the model parameters of the main branch are fixed and the model parameters of the secondary branch are trained using the entire model.

[0058] This step involves training the main branch. The training samples for the main branch use defect images and their corresponding coarse masks. To distinguish them from the training samples used in the sub-branch training process, the defect image used for the main branch training is referred to as the first defect image. The coarse mask covers the defect portion in the defect image, typically the smallest bounding polygon of the defect portion, such as the smallest rectangle covering the defect portion. Since the main branch is a diffusion model, and its training process does not require the participation of the sub-branch, existing diffusion model training methods can be used. Specifically, the first VAE encoder, first UNET encoder, and UNET decoder of the main branch can be used to process the input defect image and coarse mask to obtain predicted noise. The predicted noise is compared with standard noise to determine the loss function, and then the model parameters on the main branch are updated based on the loss function. After several rounds of training, when the training termination condition is met, the training of the main branch ends, and the parameters of the main branch remain unchanged. During model training in the main branch, the update of model parameters includes at least the update of parameters of the first UNET encoder and UNET decoder, that is, the update of parameters of the UNET structure; the parameters of the VAE structure part (i.e. the first VAE encoder and VAE decoder) may not be updated during the training of the main branch, but pre-trained or classic VAE structure parameters can be used.

[0059] In another preferred embodiment, in order to enable the model to generate defect images better according to user needs, when the image generation model includes a text processing part, the above training process also needs to introduce a text processing part to extract features from the input defect description text and input the extracted text features into the UNET structure part of the main branch (i.e., the first UNET encoder and UNET decoder). The specific processing can be the same as the existing method.

[0060] Step 202: After completing the training of the main branch, the sub-branch is trained using the second defect image and its corresponding coarse and fine masks.

[0061] In this application, the sub-branch is mainly used to understand the fine mask, introducing its feature information into the main branch to assist its processing and guide the main branch's defect image generation process. This allows the main branch's diffusion model to reference the fine mask's feature information, enhancing the accuracy and controllability of defect images while generating diverse defect images. Since the sub-branch only generates auxiliary information and does not need to independently complete the defect image generation process, it can include a second VAE encoder and a second UNET encoder, but may not include a UNET decoder and a VAE decoder. Furthermore, the training of the sub-branch requires processing by the entire model. Because the parameters of the main branch have already been obtained during previous training, the model parameters of the main branch remain unchanged during the training of the sub-branch; only the model parameters of the sub-branch are updated.

[0062] Specifically, the training process for the secondary branch is as follows: Figure 3 As shown, it includes:

[0063] Step 202a: Input the second defect image and its corresponding fine mask into the sub-branch, and after processing by the second VAE encoder and the second UNET encoder, obtain the output of the second UNET encoder;

[0064] The input to the secondary branch is the defect image and its corresponding fine mask. The fine mask is a mask that has the same size and shape as the defect in the defect image. To distinguish it from the training samples of the main branch, the defect image used during the training of the secondary branch is referred to as the second defect image.

[0065] The processing of the second VAE encoder and the second UNET encoder can be implemented in the existing manner to obtain the output of the second UNET encoder. As mentioned earlier, the sub-branch itself does not need to independently complete the defect image generation process, so it is sufficient to utilize only the information encoding capability of the UNET encoder. It is not necessary to introduce a UNET decoder. Introducing a UNET decoder would significantly increase the number of parameters, resulting in redundancy and reducing training and inference efficiency.

[0066] Step 202b: Input the second defect image and its corresponding coarse mask into the main branch, input the output of the second UNET encoder into the UNET decoder, and obtain the prediction noise output by the UNET decoder after processing by the main branch;

[0067] The second defect image and coarse mask are input into the main branch for processing, and the output of the second UNET encoder from the sub-branch is also input into the UNET decoder to guide the UNET decoding process in the main branch. The main branch processes all inputs and obtains the output of the UNET decoder, i.e., the predicted noise, for subsequent parameter updates.

[0068] Step 202c: Update the model parameters of the sub-branch based on the comparison between the predicted noise and the standard noise.

[0069] This step involves noise comparison to determine the loss function, and the specific processing is the same as in the training process of existing diffusion models. Since only the model parameters of the sub-branch are trained here, while the model parameters of the main branch remain unchanged, the model parameters of the sub-branch are updated based on the loss function after obtaining it. Here, similar to the main branch, the model parameter update of the sub-branch includes at least the parameter update of the second UNET encoder; however, the parameters of the second VAE encoder may not be updated during the training of the sub-branch, and pre-trained or classic VAE encoder parameters may be used instead.

[0070] In another preferred embodiment, in order to enable the model to generate defect images better according to user needs, when the image generation model includes a text processing part, the above training process also needs to introduce a text processing part to extract features from the input defect description text, and input the extracted text features into the second UNET encoder of the sub-branch and the UNET structure part of the main branch (i.e., the first UNET encoder and the UNET decoder), respectively. The specific processing can be the same as the existing method.

[0071] As shown above, the training of the secondary branch is completed through steps 202a to 202c. After several rounds of training, the training of the entire model ends when the training termination condition is met, and a trained image generation model is obtained, which can be used to generate defect images.

[0072] This concludes the basic workflow of the image generation model training method of this application. Through the processing described above, the image generation model adopts a dual-branch structure, introducing fine mask information on the basis of a diffusion model based on coarse masks. Simultaneously, the model training employs a two-stage training strategy: first, the secondary branches are disabled, and the main branch is trained separately to ensure the model has the ability to generate diverse defect images in batches. After the main branch training is completed, the parameters of the secondary branch model are trained using the entire model's processing, introducing fine control over defects, thereby training and generating a complete model to improve the accuracy and controllability of the model's defect generation.

[0073] Based on the training method provided in this application, this application also provides an image generation method that can use the trained image generation model to generate defect images. Figure 4 This is a schematic diagram illustrating the basic workflow of an image generation method. (Example:) Figure 4 As shown, the method includes:

[0074] Step 401: Input the normal image and the mask image describing the defect location into the image generation model;

[0075] The image generation model is the aforementioned Figure 1 model shown, and is generated by the Figure 2 and Figure 3 methods shown.

[0076] Step 402, the image generation model performs inference processing to generate defective images.

[0077] Among them, to flexibly control the participation of the processing of the sub-branch in the entire model, the output of the second UNET encoder of the sub-branch can be further multiplied by scale before being input into the UNET decoder of the main branch. This scale is the participation ratio of the fine mask set in advance according to requirements in image generation, and 0 ≤ scale ≤ 1.

[0078] More specifically, when scale = 0, that is, the sub-branch does not participate in image generation, then the mask image input to the image generation model is only the rough mask image, and the model generates images completely according to the main branch. The processing performed at least includes the processing of the first VAE encoder, the processing of the first UNET encoder, the processing of the UNET decoder, and the processing of the VAE decoder. That is, image generation is performed using the rough mask, and the generated defective images have strong diversity; in this way, the precise mask is not required at all, so large-scale image generation can still be carried out when the labor cost is limited;

[0079] When 0 < scale ≤ 1, the sub-branch participates in image generation in a certain proportion. The mask image input to the image generation model is the fine mask image and the rough mask image, and the masked part in the rough mask image includes the masked part in the fine mask image; by introducing the sub-branch, the accuracy and controllability of the generated defective images can be enhanced. At the same time, the participation ratio of the sub-branch in image generation can be controlled by the value of scale, that is, it can be controlled how much accuracy and controllability the generated defective images have, and whether to select to generate more diverse defective image samples or more precisely controllable defective image samples. At the same time, for the same model input, multiple different defective images can be generated by setting different scale values; that is, compared with the existing method of using a fine mask for image generation, more output images can be generated by adjusting the scale value for each input mask image. The processing of the model at least includes: the processing of the second VAE encoder and the second UNET encoder for the image and its fine mask, the processing of the first VAE encoder, the first UNET encoder, the UNET decoder, and the VAE decoder for the image and its rough mask. Of course, it also includes the processing of the UNET decoder for the output of the second UNET encoder.

[0080] So far, Figure 4The image generation method flowchart shown is now complete. The above image generation method can generate images entirely based on coarse masks, or it can generate images based on a set ratio of coarse and fine masks. This allows for the batch generation of defect images with greater diversity, and also enables the generation of more precise and controllable defect images as needed.

[0081] The following specific embodiments illustrate the training method of the image generation model in this application and the method of image generation using the corresponding model.

[0082] Example 1:

[0083] Figure 5 This is a schematic diagram illustrating the specific process of training the image generation model in Embodiment 1 of this application. Figure 5 As shown, in this embodiment, the image generation model includes not only the main branch and sub-branch, but also a text processing part to further introduce a text description of the generated defects. Specifically, it can describe the defect type, such as a metal part with a crack. Steps 501-504 are the main branch training part, and steps 505-510 are the sub-branch training part. The entire training method specifically includes:

[0084] Step 501: Obtain the first defect image and its corresponding coarse mask.

[0085] The defect image sample used for main branch training is called the first defect image. A coarse mask corresponding to the first defect image covers the defect portion of the first defect image. In this embodiment, the coarse mask is obtained by randomly expanding a fine mask outwards; it can be a rectangle or other arbitrary shape, but the coarse mask must completely contain the fine mask. The fine mask is the same size and shape as the defect portion of the first defect image. The corresponding area of ​​the first defect image is covered by the coarse mask to obtain the Masked Image, which is then input into the main branch for processing.

[0086] Step 502: The text processing part processes the input text and inputs the text processing result into the UNET structure of the main branch.

[0087] The text processing part can be a text feature extraction structure, such as a text embedding model like CLIP, which extracts text features from the input defect description text, obtains text features, and inputs the text features into the UNET structure of the main branch (that is, the structure composed of the first UNET encoder and the UNET decoder).

[0088] Step 503: The main branch processes the Masked Image and text processing results to obtain the predicted noise.

[0089] Figure 6The diagram illustrates the training structure of the main branch, which includes a first VAE encoder, a first UNET encoder, a UNET decoder, and a VAE decoder. Since the training process of the main branch does not require a VAE decoder (similar to the training of existing diffusion models), therefore... Figure 6 The VAE decoder is omitted; the sub-branch does not participate in the training of the main branch, so the structural part of the sub-branch is represented by a gray semi-transparent box.

[0090] Existing diffusion models can also include a text processing component. The training process of the main branch can be the same as that of existing diffusion models, and the way the text processing results are introduced into the main branch for processing can also be the same as that of existing diffusion models.

[0091] The following is an example of main branch processing steps:

[0092] 1) Input the Masked Image into the first VAE encoder to obtain the latent variable Masked ImageLatents, and concatenate it with the mask and random latent variables (Noisy Latents, which are obtained by sampling through the Noisy Scheduler, the same as the existing diffusion model) in the channel dimension to obtain the concatenated vector;

[0093] 2) Input the text processing result concatenated vector together with the current time step t into the first UNET encoder; wherein, the first UNET encoder and UNET decoder form the UNET model. In this embodiment, the Stable-Diffusion Inpainting model is used as the initial UNET model, which inherits its Unet structure and parameters, thereby ensuring that it can fit quickly when fine-tuning with defective samples.

[0094] 3) The UNET model processes the input and outputs predicted noise. ,in, Let t be the noisy latent representation at step t, which is the concatenated vector mentioned above, and c be the vector representation of the text processing result, such as the text feature vector.

[0095] Step 504: Compare the predicted noise obtained in step 503 with the standard noise, calculate the MSE loss function, and update the model parameters of the main branch based on the value of the loss function.

[0096] This step involves noise comparison and loss function calculation; in this embodiment, the MSE loss function is specifically used. The model parameters of the main branch are updated based on the value of the loss function, mainly including the parameters of the UNET model. Standard noise is typically standard Gaussian noise, consistent with the noise introduced after processing by the first VAE encoder.

[0097] As described above, steps 501-504 complete one round of main branch training. After multiple rounds of training are executed and the training termination condition is met, step 505 is executed to train the secondary branch. After the main branch training, the main branch model can serve as an independent defect image generation model with diverse generation capabilities.

[0098] Step 505: Obtain the second defect image and its corresponding coarse and fine masks.

[0099] The defect image sample used for subbranch training is called the second defect image. The meanings of coarse mask and fine mask have been described in step 501, and will not be repeated here.

[0100] A coarse mask is used to cover the corresponding area of ​​the second defect image to obtain Masked Image A, which is then processed by the main branch. A fine mask is used to cover the corresponding area of ​​the second defect image to obtain Masked Image B, which is then processed by the sub-branch.

[0101] Step 506: The text processing part processes the input text and inputs the text processing results into the UNET structure of the main branch and the second UNET encoder of the sub-branch, respectively.

[0102] In the processing of the sub-branch, the text processing results are used to guide the processing of the UNET structure in the main branch and the sub-branch. Therefore, the text processing results need to be input into the UNET structure of the main branch and the sub-branch. The UNET structure of the sub-branch is the second UNET encoder, and the UNET structure of the main branch includes the first UNET encoder and the UNET decoder. The specific processing of the text processing part is the same as in step 502, and will not be repeated here.

[0103] Step 507: The sub-branch processes the Masked Image B and the text processing results to obtain the output of the second UNET encoder.

[0104] Figure 7 This is a schematic diagram of the training structure for the secondary branch, which includes the second VAE encoder and the second UNET encoder. The structure of the main branch is similar to... Figure 6 The same applies, so I won't go into details here.

[0105] In this embodiment, the parameter values ​​of the second UNET encoder in the sub-branch are the same as those of the first UNET encoder in the main branch in the initial state. That is, the second UNET encoder in the sub-branch reuses the first UNET encoder obtained after training in the main branch as a pre-trained model. Specifically, the processing of the second VAE encoder and the second UNET encoder can adopt the same processing as in the existing diffusion model, and the way the text processing results are introduced into the sub-branch for processing can also be the same as in the existing diffusion model.

[0106] The following is an example of the steps for handling a secondary branch:

[0107] The Masked Image B is input into the second VAE encoder to obtain the latent variable Masked ImageLatents. This latent variable is then concatenated with the mask and the random latent variable (Noisy Latents, obtained by sampling through the Noisy Scheduler) in the channel dimension to obtain the concatenated variable B. The co-processing text result and the current time step t are input into the second UNET encoder. The second UNET encoder then processes the data to obtain the output of the second UNET encoder.

[0108] Step 508: After processing the output of the second UNET encoder using a zero convolutional layer, the output is added to the UENT decoder of the main branch.

[0109] In this application, when the output of the second UNET encoder of the sub-branch is introduced into the UNET decoder, the output of each layer of the second UNET encoder of the sub-branch can be input into the UNET decoder of the main branch. In the UNET decoder, the output of each layer input by the second UNET encoder is merged with the corresponding layer output input by the first UNET encoder in the main branch, and the merged result is used by the UNET decoder to perform the decoding operation of the corresponding layer.

[0110] In this embodiment, before introducing the output of the second UNET encoder into the UNET decoder of the main branch, the output of the second UNET encoder is first processed by a zero convolutional layer, and then the introduction operation is performed.

[0111] Specifically, a zero-convolutional layer is a special type of convolutional layer. Its core characteristic is that the initial weights are all zero, or initialized to near-zero minimum values. During training, it gradually learns effective parameters through gradient updates to achieve progressive training control of the module. In this embodiment, the output of each layer of the second UNET encoder is introduced into the UNET decoder. More specifically, firstly, the outputs of each layer of the second UNET encoder are processed by zero-convolutional layers to stabilize the training process. Then, the outputs of the zero-convolutional layers are injected into the UNET decoder of the main branch and merged with the outputs of the corresponding layers of the first UNET encoder for decoding processing of the corresponding layers. The zero-convolutional layer processing ensures that the participation of the sub-branch in the main branch processing is progressive.

[0112] Step 509: The main branch processes Masked Image A, the text processing result, and the zero convolutional layer output introduced by the sub-branch to obtain the prediction noise.

[0113] After the zero-convolutional layer output from the secondary branch is introduced into the main branch, the main branch, as described above, merges the zero-convolutional layer output with the output of the first UNET encoder. Other processing remains the same as the training process for existing diffusion models, and will not be repeated here. The resulting processing yields the predicted noise.

[0114] Step 510: Compare the predicted noise obtained in step 509 with the standard noise, calculate the MSE loss function, and update the model parameters of the sub-branch based on the value of the loss function.

[0115] During the training of the secondary branch, the MSE loss function is calculated. ,in, To predict noise, The parameters of the main branch model, For the parameters of the secondary branch model, E represents the values ​​at t, x0, ... Take the expected value, where t represents the time step and x0 is the original defect map. It follows a standard Gaussian distribution.

[0116] When updating parameters based on the loss function value, only the parameters in the sub-branch are updated, which may include the parameters of the second UNET encoder and the parameters of the zero convolutional layer.

[0117] As described above, steps 505-510 complete one round of training for the secondary branch. After repeatedly executing multiple rounds of training and meeting the training termination condition, the training of the entire model ends. After training the complete model as described above, the resulting image generation model can simultaneously utilize coarse and fine masks to generate defective images with diversity and controllability.

[0118] At this point, Figure 5 The training process for the image generation model in Example 1 is now complete.

[0119] Example 2:

[0120] Figure 8 The following is a specific flowchart of the image generation method in the second embodiment of this application. This image generation method is based on Figure 5 the image generation model trained by the training method shown below, that is, the inference process of this image generation model, Figure 9 which is the specific inference structure of the image generation model. In this embodiment, similar to Figure 4 the process shown in the figure, scale is introduced to represent the participation ratio of the fine mask in image generation. As Figure 8 and Figure 9 shown, this image generation method includes:

[0121] Step 801: Obtain a normal image and a mask image describing the defect location.

[0122] Among them, according to different values of scale, the selection of the mask image will be different.

[0123] Specifically, when scale = 0, the mask image input to the image generation model is a rough mask image, which can be determined within the main body range of the normal image to determine the mask position and generate a rough mask image. For example, in the normal image, the main body range of the image is determined through object detection, image segmentation, etc. Then, within this main body range, the mask position is determined in a certain order or randomly. The mask shape can be the default, user-specified or randomly generated shape, and thus a rough mask can be determined and a rough mask image can be generated.

[0124] When 0 < scale ≤ 1, the mask images input to the image generation model are a fine mask image and a rough mask image. Among them, the fine mask can be marked based on user needs, and the rough mask can be obtained by randomly expanding the fine mask outward.

[0125] Use the rough mask to cover the corresponding area of the normal image to obtain Masked Image A and input it into the main branch for processing; use the fine mask to cover the corresponding area of the normal image to obtain Masked Image B and input it into the sub-branch for processing.

[0126] Step 802: The text processing part processes the input text and inputs the text processing results into the UNET structure of the main branch and the second UNET encoder of the sub-branch respectively.

[0127] The processing of this step is the same as the aforementioned step 506 and will not be elaborated here.

[0128] Step 803: The sub-branch processes the Masked Image B and the result of text processing to obtain the output of the second UNET encoder. After passing through the zero convolutional layer, it is multiplied by the scale factor scale and then input into the UNET decoder of the main branch.

[0129] Among them, by adjusting the value of scale, the proportion of the sub-branch participating in image generation can be controlled. For example, Figure 10a , 10b and 10c show the corresponding defective image generation situations under three cases of scale values.

[0130] When scale = 0, the sub-branch does not participate in the generation of defective images, and the model is equivalent to only retaining the main branch. Therefore, through model processing, defects can be randomly generated within the rough mask area, ensuring the maximum diversity of defect generation.

[0131] When 0 < scale < 1, the sub-branch partially participates in the generation of defective images. The defective images generated through model processing have both randomness and certainty. By adjusting the size of Scale, it is possible to choose to generate more diverse defective images or more precisely controllable defective images.

[0132] When scale = 1, the sub-branch fully participates in the generation of defective images. Due to the training of the sub-branch, the defective images generated by the model are deterministic and can generate defects that fit the fine mask, so it has strong controllability.

[0133] Users can select an appropriate scale according to their needs. Or, for the same image, multiple different scale values can be used to generate defective images, so as to generate more defective images.

[0134] Step 804: The main branch processes Masked Image A, the result of text processing, and the result after multiplying the output of the zero convolutional layer introduced by the sub-branch by scale to obtain the generated defective image.

[0135] In the inference process of this embodiment, the processing of the first VAE encoder, the first UNET encoder, the UNET decoder, and the VAE decoder is included, and finally a defective image is generated. Among them, the UNET decoder merges the output of the first UNET encoder and the product result input into the UNET decoder in step 803 (that is, the result of multiplying the output of the second UNET encoder by the scale after passing through the zero convolutional layer) for corresponding layer UNET decoding processing. Except for the merging operation in the above UNET decoder, the inference processing of other modules can be the same as the existing method and will not be elaborated here.

[0136] So far, Figure 8The illustrated process flow is now complete. The image generation method described above enables the generation of diverse and controllable defect images tailored to user needs.

[0137] The above describes the specific implementation of the training method and image generation method of the image generation model in this application. It employs an asymmetric dual-branch model structure combined with a two-stage training strategy, requiring only a small number of training samples to complete model fine-tuning training. This allows for the batch generation of defect samples, reducing data acquisition and annotation costs. The trained model can simultaneously utilize both coarse and fine masks to generate defect images with diversity and controllability. Furthermore, the model trained in this manner can adjust and balance the diversity and accuracy of generated defects through a scaling factor (Scale). Reducing the Scale factor, combined with a coarse mask, generates diverse defect images; while increasing the Scale factor, combined with a fine mask, generates defect images of a specified shape and size.

[0138] This application also provides a training apparatus for an image generation model and an image generation apparatus, which can be used to implement the training method and image generation method of the image generation model described above.

[0139] Figure 11 This is a schematic diagram of the basic structure of the training device for the image generation model provided in this application. The image generation model includes a main branch and a sub-branch, such as... Figure 11 As shown, the training device includes a main branch training unit and a sub-branch training unit.

[0140] The main branch training unit is used to train the main branch using the first defect image and its corresponding coarse mask, and the training process of the main branch masks the processing of the sub-branch; the main branch is a diffusion model, including a first VAE encoder, a first UNET encoder, a UNET decoder and a VAE decoder.

[0141] The sub-branch training unit is used to train the sub-branch using the second defect image and its corresponding coarse and fine masks after the main branch training is completed; the sub-branch includes the second VAE encoder and the second UNET encoder.

[0142] The training of the secondary branch specifically includes:

[0143] The second defect image and its corresponding fine mask are input into the sub-branch, and after processing by the second VAE encoder and the second UNET encoder, the output of the second UNET encoder is obtained.

[0144] The second defect image and its corresponding coarse mask are input into the main branch, and the output of the second UNET encoder is input into the UNET decoder. After processing by the main branch, the predicted noise output by the UNET decoder is obtained. The model parameters of the sub-branch are updated based on the comparison between the predicted noise and the standard noise.

[0145] Optionally, the training device may further include a text processing unit, which is used to input the defect description text of the first defect image into the text processing branch for text feature extraction, and use the extracted text features to guide the processing of the first UNET encoder in the main branch training unit; and is also used to input the defect description text of the second defect image into the text processing branch for text feature extraction, and use the extracted text features to guide the processing of the second UNET encoder and the first UNET encoder in the sub-branch training unit.

[0146] Optionally, the initial parameters of the second UNET encoder can be the same as those of the first UNET encoder after the main branch training is completed.

[0147] Optionally, the training device may further include a zero-convolutional layer processing unit for processing the output of the second UNET encoder with a zero-convolutional layer before inputting it into the UNET decoder.

[0148] Optionally, in the sub-branch training unit, the process of inputting the output of the second UNET encoder to the UNET decoder may specifically include: inputting the output of each layer of the second UNET encoder to the UNET decoder;

[0149] In the UNET decoder, the outputs of each layer of the second UNET encoder are merged with the outputs of the same layer of the first UNET encoder for decoding processing of the corresponding layer in the UNET decoder.

[0150] Optionally, for the second defect image, its corresponding fine mask can be randomly expanded to obtain the corresponding coarse mask.

[0151] Figure 12 A schematic diagram of the basic structure of the image generation apparatus provided in this application. Figure 12 As shown, the image generation device includes an input module and an inference module for the image generation model.

[0152] The input module is used to input a normal image and a mask image describing the location of the defect into the image generation model output by the training device of the aforementioned image generation model;

[0153] The inference module of the image generation model is used to generate defective images by inferring from normal images and masked images using the image generation model.

[0154] The output of the second UNET encoder is multiplied by a scale before being input to the UNET decoder. The scale is the participation ratio of a fine mask that is set in advance according to requirements, where 0 ≤ scale ≤ 1.

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

Claims

1. A method for training an image generation model, characterized in that, The image generation model comprises a main branch and a secondary branch, and the training method comprises: training the main branch by using a first defect image and a corresponding coarse mask, and the training process of the main branch shields the processing of the secondary branch; wherein the main branch is a diffusion model comprising a first VAE encoder, a first UNET encoder, a UNET decoder and a VAE decoder; after the training of the main branch is completed, the secondary branch is trained by using a second defect image and corresponding coarse and fine masks; the secondary branch comprises a second VAE encoder and a second UNET encoder; wherein the training of the secondary branch comprises: inputting the second defect image and the corresponding fine mask into the secondary branch to obtain a second UNET encoder output through the processing of the second VAE encoder and the second UNET encoder; inputting the second defect image and the corresponding coarse mask into the main branch, inputting the second UNET encoder output into the UNET decoder, obtaining a predicted noise of the UNET decoder output through the processing of the main branch, and updating the model parameters of the secondary branch based on the comparison between the predicted noise and a standard noise.

2. The method of claim 1, wherein, The image generation model comprises a text processing branch; when training the main branch, inputting a defect description text of the first defect image into the text processing branch for text feature extraction, and guiding the processing of the first UNET encoder by using the extracted text features; when training the secondary branch, inputting a defect description text of the second defect image into the text processing branch for text feature extraction, and guiding the processing of the second UNET encoder and the first UNET encoder by using the extracted text features.

3. The method according to claim 1 or 2, characterized in that, The initial parameters of the second UNET encoder are the same as the parameters of the first UNET encoder after the training of the main branch is completed.

4. The method according to claim 1 or 2, characterized in that, The method further comprises: inputting the second UNET encoder output into the UNET decoder after zero convolution layer processing.

5. The method according to claim 1 or 2, characterized in that, The inputting of the second UNET encoder output into the UNET decoder comprises: inputting the outputs of each layer of the second UNET encoder into the UNET decoder; In the UNET decoder, the outputs of each layer of the second UNET encoder are combined with the outputs of the same layer of the first UNET encoder for decoding processing of the corresponding layer of the UNET decoder.

6. The method of claim 1 or 2, wherein, For the second defect image, the corresponding coarse mask is obtained by randomly expanding the corresponding fine mask.

7. An image generation method characterized by, Comprises: inputting a normal image and a mask image describing a defect position into the image generation model trained by the method of any one of claims 1 to 6 for inference processing to generate a defect image; wherein the second UNET encoder output is multiplied by scale before inputting into the UNET decoder, and the scale is a participation ratio of the fine mask set in advance according to requirements, 0 ≤ scale ≤ 1.

8. The method of claim 7, wherein, When the scale = 0, the mask image input into the image generation model is a coarse mask image; When 0 < scale ≤ 1, the mask image input into the image generation model is a fine mask image and a coarse mask image, and the mask part in the coarse mask image includes the mask part in the fine mask image.

9. An apparatus for training an image generation model, comprising: The image generation model comprises a main branch and a secondary branch, and the training device comprises a main branch training unit and a secondary branch training unit. The main branch training unit is configured to train the main branch by using a first defect image and a corresponding coarse mask, and the training process of the main branch shields the processing of the secondary branch; wherein the main branch is a diffusion model comprising a first VAE encoder, a first UNET encoder, a UNET decoder and a VAE decoder. The secondary branch training unit is configured to train the secondary branch by using a second defect image, a corresponding coarse mask and a fine mask after the training of the main branch is completed; the secondary branch comprises a second VAE encoder and a second UNET encoder. The training of the secondary branch comprises: inputting the second defect image and the corresponding fine mask into the secondary branch to obtain a second UNET encoder output through the processing of the second VAE encoder and the second UNET encoder; inputting the second defect image and the corresponding coarse mask into the main branch, inputting the second UNET encoder output into the UNET decoder, obtaining a predicted noise of the UNET decoder output through the processing of the main branch, and updating the model parameters of the secondary branch based on the comparison between the predicted noise and a standard noise.

10. The training device of claim 9, wherein, The training device comprises a text processing unit configured to input a defect description text of the first defect image into the text processing branch to extract text features, and guide the processing of the first UNET encoder and the UNET decoder in the main branch training unit by using the extracted text features; and input a defect description text of the second defect image into the text processing branch to extract text features, and guide the processing of the second UNET encoder, the first UNET encoder and the UNET decoder in the secondary branch training unit by using the extracted text features.

11. Training device according to claim 9 or 10, characterized in that The training device further comprises a zero convolution layer processing unit configured to input the second UNET encoder output into the UNET decoder after zero convolution layer processing.

12. An image generation apparatus characterized by comprising: The training device comprises: an input module and an inference module of an image generation model; the input module is configured to input a normal image and a mask image describing a defect position into the image generation model output by the training device of the image generation model according to claim 9; the inference module of the image generation model is configured to use the image generation model to perform inference based on the normal image and the mask image to generate a defect image; wherein the second UNET encoder output is multiplied by scale before being input into the UNET decoder, and the scale is a participation ratio of the fine mask set in advance according to requirements, and 0 ≤ scale ≤ 1.

Images