Diffusion model fine-tuning method and device for few-sample defect image generation

Abstract

The present invention relates to the field of computer vision technology, and provides a method and device for fine-tuning a diffusion model for generating a few-sample defect image. The method includes: combining multiple defect words and a single product word to obtain an unbalanced abnormal text prompt; selecting a few-sample defect image and a defect image mask as inputs of a defect image branch, aligning the defect word to the defect area in the image according to the defect image mask, and optimizing the first objective function and the second objective function; selecting a fixed number of normal images as inputs of a normal image sample branch, and using the product-related part of the unbalanced abnormal text prompt to optimize the third objective function; a hybrid training strategy is used to enable defect images of different defect categories to participate in training together. The present invention can still achieve a relatively ideal training effect when the number of defect images of each defect category in an industrial scenario is small, thereby generating defect images using the trained diffusion model.

Description

Technical Field

[0001] The present invention relates to the field of computer vision technology, and in particular to a diffusion model fine-tuning method and device for generating a few-sample defect image. Background Art

[0002] Existing image segmentation methods typically require a large number of defect images annotated with masks, but the scarcity of defect images limits their application in industrial scenarios. To address this challenge, defect image generation methods for industrial scenarios have emerged, aiming to expand the training set of image segmentation models. Currently, generation methods in industrial scenarios can be roughly divided into two categories: defect generation and defect image generation. Defect generation methods simulate defects by pasting cropped normal regions onto normal images, generate defects only and then merge them into real normal images through different strategies, or generate defects using diffusion models and edit them into normal images based on the guidance of defect masks.

[0003] However, defect generation methods rely on defect masks as input. If the mask placement is inappropriate, the generated defects are prone to low fidelity and poor consistency. In contrast, defect image generation methods can directly generate complete industrial defect images, which include both industrial defects and the industrial products that exhibit them. This method avoids reliance on masks, potentially generating more natural and diverse defect instances. However, defect image generation methods require pre-training the diffusion model with a large number of defect images. Few-shot industrial defect image generation refers to training an industrial defect image generation model based on a small number of training images (generally less than five images). For the same industrial product, there are often multiple defect images of different defect categories. Because the occurrence of defects in industrial scenarios is normally a low-probability event, it is impossible to quickly collect a large number of defect images for each defect. Therefore, few-shot training is a common practice in industrial scenarios. This results in a small number of samples used for training a single defect category, making it difficult to achieve ideal training results.

[0004] In view of this, overcoming the defects of the prior art is an urgent problem to be solved in this technical field. Summary of the Invention

[0005] The technical problem to be solved by the present invention is that the number of samples used for training a single defect category is small, making it difficult to achieve an ideal training effect.

[0006] The present invention adopts the following technical solutions:

[0007] In a first aspect, the present invention provides a diffusion model fine-tuning method for generating a few-shot defect image. The defect image generation model includes a U-Net encoder, a U-Net decoder, an attention module, a first objective function, a second objective function, and a third objective function. The U-Net encoder, the U-Net decoder, the attention module, the first objective function, and the second objective function constitute a defect image branch, and the U-Net encoder, the U-Net decoder, and the third objective function constitute a normal image branch. The training process of the two branches is performed separately but implemented on a shared U-Net. The defect image generation model is trained using a hybrid training strategy, and the method includes:

[0008] Combining multiple defect words and a single product word to obtain an unbalanced abnormal text prompt, and inputting the unbalanced abnormal text prompt and the product-related portion of the unbalanced abnormal text prompt as guidance conditions into the defect image branch and the normal image branch respectively;

[0009] Selecting a few-sample defect image and a corresponding defect image mask as input to a defect image branch, the defect image branch obtains features of the defect image and aligns defect words in the imbalanced abnormal text prompt to the defect area in the image according to the defect image mask corresponding to the defect image, so as to optimize the first objective function and the second objective function;

[0010] A fixed number of normal images are selected as input to the normal image sample branch, which obtains features of the normal images and uses the product-related parts of the unbalanced abnormal text prompts to optimize the third objective function; wherein the hybrid training strategy is used to enable defect images of different defect categories to participate in training together.

[0011] Preferably, the first objective function is in, is the cross attention map corresponding to the nth defective word, N is the number of defective words, L is the number of U-Net layers corresponding to the cross attention map used in the alignment process, and M l It is a binary mask with a size of r×r pixels, with a value of 1 in the defect area and a value of 0 in the non-defect area. is the cross attention map corresponding to the product word, and ⊙ is the multiplication operation per pixel.

[0012] Preferably, the second objective function is Among them, E df represents the conditional vector obtained after the unbalanced abnormal text prompt passes through the text encoder, z df is the defect image x df The latent vector obtained by the VAE encoder, ∈ dfis random noise, t df is the number of noise adding steps, To use ∈ df and ∈ df z df The noise-added latent vector obtained by adding noise is the predicted noise, is the value of the first objective function.

[0013] Preferably, the third objective function is

[0014]

[0015] Among them, e ob represents the conditional vector obtained after the product-related part of the unbalanced abnormal text prompt passes through the text encoder, z ob is a normal image x ob The latent vector obtained by the VAE encoder, ∈ ob is random noise, t ob is the number of noise adding steps, To use ∈ ob and t ob z ob The noise latent vector obtained by adding noise is is the predicted noise, is the value of the first objective function.

[0016] Preferably, the step of combining multiple defective words with a single product word to obtain an imbalance abnormality text prompt specifically includes:

[0017] Combine the product word of a single product with the N defect words of the product to obtain the imbalanced abnormal text prompt 'a <ob>with <df1> <df2>… <df N >'; <ob>is the product word, <df1> 、 <df2>、…、 <df N > is N defect words, where N is a preset value. The N defect words for each product correspond to one defect category. Using K*N defect words, we obtain unbalanced abnormal text prompts of K shared product words, covering all defect categories of the product. K is a preset value, corresponding to K defect categories of a product.

[0018] Preferably, for each defect category, an unbalanced abnormal text prompt consisting of different defect tokens is used.

[0019] Preferably, the hybrid training strategy is specifically:

[0020] Generate a defect image set based on defect images of all defect categories of the corresponding product in, is the i-th defect image;

[0021] Select a preset number of normal images of the product to generate a normal image set in, is the i-th normal image;

[0022] Construct defect image set X df and the normal image set X ob ;

[0023] From the defect image set X df Select the first number of defective images from the normal image set X ob A first number of normal images are selected, the selected first number of defect images and the first number of normal images are mixed into a batch of data, and the data of each batch is input into the defect image branch and the normal image branch for multi-batch training.

[0024] Preferably, the defect image set X is constructed df and the normal image set X ob , specifically including:

[0025] Generate a defect image set based on defect images of all defect categories of the corresponding product in, is the i-th defect image;

[0026] Select a preset number of normal images of the product to generate a normal image set in, is the i-th normal image.

[0027] In a second aspect, the present invention further provides a diffusion model for generating a few-sample defect image, which is obtained using the diffusion model fine-tuning method for generating a few-sample defect image described in the first aspect.

[0028] In a third aspect, the present invention further provides a diffusion model fine-tuning device for generating a few-sample defect image, which is used to implement the diffusion model fine-tuning method for generating a few-sample defect image described in the first aspect, and the device includes:

[0029] At least one processor; and a memory communicatively connected to the at least one processor; wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the processor to execute the diffusion model fine-tuning method for generating a few-sample defect image as described in the first aspect.

[0030] In a fourth aspect, the present invention further provides a diffusion model device for generating a few-sample defect image, which is used to implement the diffusion model for generating a few-sample defect image described in the first aspect, and the device includes:

[0031] At least one processor; and a memory communicatively connected to the at least one processor; wherein the memory stores instructions executable by the at least one processor, the instructions being executed by the processor to carry the diffusion model for generating a few-sample defect image as described in the second aspect, and to generate a defect image using the diffusion model for generating a few-sample defect image.

[0032] In a fifth aspect, the present invention further provides a non-volatile computer storage medium, wherein the computer storage medium stores computer-executable instructions, which are executed by one or more processors to complete the method described in the first aspect.

[0033] In a sixth aspect, a chip is provided, comprising: a processor and an interface, for calling and running a computer program stored in a memory to execute the method of the first aspect.

[0034] In a seventh aspect, a computer program product comprising instructions is provided, which, when executed on a computer or a processor, causes the computer or the processor to execute the method of the first aspect.

[0035] The present invention introduces unbalanced abnormal text prompts, introduces multiple defect words into the unbalanced abnormal text prompts, and under the guidance of the unbalanced abnormal text prompts, proposes a dual-branch diffusion model fine-tuning method, that is, using two branches, namely a defect image branch and a normal image branch, and further using a hybrid training strategy to introduce defect images of different categories, thereby alleviating the overfitting problem that may be caused by the limited number of samples of each defect category, and increasing the diversity of defect images while still maintaining the authenticity of the defect images. Therefore, even if the number of defect images of each defect category in industrial scenarios is small, the ideal training effect can still be achieved, and the defect image can be generated using the trained diffusion model. BRIEF DESCRIPTION OF THE DRAWINGS

[0036] To more clearly illustrate the technical solutions of the embodiments of the present invention, the following briefly introduces the drawings required for use in the embodiments of the present invention. Obviously, the drawings described below are only some embodiments of the present invention. Those skilled in the art can also derive other drawings based on these drawings without inventive effort.

[0037] Figure 1 1 is a schematic diagram of the architecture of a defect image generation model in a diffusion model fine-tuning method for generating defect images with a small number of samples provided by an embodiment of the present invention;

[0038] Figure 2 This is a flow chart of a first diffusion model fine-tuning method for generating a few-sample defect image provided by an embodiment of the present invention;

[0039] Figure 3 This is a flow chart of a second method for fine-tuning a diffusion model for generating a few-sample defect image provided by an embodiment of the present invention;

[0040] Figure 4 1 is a flow chart of a third method for fine-tuning a diffusion model for generating a few-sample defect image provided by an embodiment of the present invention;

[0041] Figure 5 4 is a flow chart of a fourth method for fine-tuning a diffusion model for generating a few-sample defect image provided by an embodiment of the present invention;

[0042] Figure 6 is a schematic diagram of a fifth diffusion model fine-tuning method for generating a few-sample defect image provided by an embodiment of the present invention;

[0043] Figure 7 1 is a schematic diagram of the architecture of a diffusion model fine-tuning device for generating a few-sample defect image, provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0044] In order to make the purpose, technical solutions and advantages of the present invention more clearly understood, the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention.

[0045] Unless the context requires otherwise, throughout the specification and claims, the term "including" is to be interpreted as meaning open inclusion, that is, "including, but not limited to". In the description of the specification, the terms "one embodiment", "some embodiments", "exemplary embodiments", "example", "specific example" or "some examples" are intended to indicate that specific features, structures, materials or characteristics associated with the embodiment or example are included in at least one embodiment or example of the present disclosure. The schematic representation of the above terms does not necessarily refer to the same embodiment or example. In addition, the specific features, structures, materials or characteristics may be included in any one or more embodiments or examples in any appropriate manner, that is, although they may be carried in the embodiments or examples of the above terms due to reasons such as the order and position of appearance, it is not limited to that they can be carried in combination by one embodiment or example.

[0046] In the description of the present invention, the terms "first" and "second" are used only for descriptive purposes and cannot be understood as indicating or implying relative importance or implicitly indicating the number of the indicated technical features. Therefore, the features defined as "first" and "second" may explicitly or implicitly include one or more of the features. In the description of the embodiments of the present disclosure, unless otherwise specified, "multiple" means two or more. In addition, for example, the description may also use the method of adding "A" and "B" at the end to describe the same type of nouns as two independent individuals. In this case, the corresponding features defined as "A" and "B" are only used to distinguish the description purposes of the same type of individuals, and cannot be understood as indicating or implying relative importance or implicitly indicating the number of the indicated technical features.

[0047] In the description of the present invention, the expression "A and / or B" (where A and B are used to formally represent specific characteristic contents) is involved, and the corresponding expressions include the following three combinations: only A, only B, and a combination of A and B.

[0048] As used herein, "about," "substantially," or "approximately" includes the stated value and an average value that is within an acceptable range of deviation from the particular value as determined by one of ordinary skill in the art taking into account the measurements in question and the errors associated with the measurement of the particular quantity (i.e., the limitations of the measurement system).

[0049] In addition, the technical features involved in the various embodiments of the present invention described below can be combined with each other as long as they do not conflict with each other.

[0050] Embodiment 1:

[0051] Embodiment 1 of the present invention provides a diffusion model fine-tuning method for generating a few-sample defect image. Figure 1 As shown in , the defect image generation model includes a U-Net encoder, a U-Net decoder, an attention module, a first objective function, a second objective function, and a third objective function; wherein, the U-Net encoder, U-Net decoder, the attention module, the first objective function, and the second objective function constitute the defect image branch, and the U-Net encoder, U-Net decoder, and the third objective function constitute the normal image branch; the training process of the two branches is carried out separately, but is implemented on a shared U-Net; the defect image generation model is trained using a hybrid training strategy, as shown in Figure 2 As shown, the method includes:

[0052] In step 201, multiple defect words and a single product word are combined to obtain an unbalanced abnormal text prompt, and the unbalanced abnormal text prompt and the product-related part of the unbalanced abnormal text prompt are respectively input into the defect image branch and the normal image branch as guiding conditions; the total number of defect categories used is obtained by technical personnel in this field based on demand analysis.

[0053] In step 202, a few-sample defect image and a corresponding defect image mask are selected as inputs of a defect image branch. The defect image branch obtains features of the defect image and, based on the defect image mask corresponding to the defect image, aligns defect words in the unbalanced abnormal text prompt to the defect area in the image to optimize the first objective function and the second objective function. The defect image may be a defect image corresponding to any defect category of the corresponding product.

[0054] In step 203, a fixed number of normal images are selected as input to the normal image sample branch, the normal image branch obtains the features of the normal images, and uses the product-related parts of the unbalanced abnormal text prompts to optimize the third objective function; wherein the hybrid training strategy is used to enable defect images of different defect categories to participate in training together.

[0055] In the defect image branch, the output of the U-Net encoder is input into the U-Net decoder. Based on the output of the U-Net decoder and the defect image, the value of the second objective function is calculated. The output of the U-Net encoder is also input into the attention module. Based on the output of the attention module and the corresponding defect image mask, the value of the first objective function is calculated. In actual use, the output of the attention module is actually the intermediate result of the attention module extracted during the process of passing through the U-Net; the attention module is embedded in the U-Net and is part of the U-Net.

[0056] The output result of the U-Net encoder is input to the U-Net decoder, and the value of the third objective function is calculated based on the output result of the U-Net decoder and the normal image.

[0057] The process of optimizing the first objective function and the second objective function and optimizing the third objective function can be understood as performing multiple iterations, calculating the values ​​of the first objective function, the second objective function and the third objective function in each iteration, and adjusting the values ​​of each parameter in the U-Net architecture according to the values ​​of the first objective function, the second objective function and the third objective function, so that the calculated values ​​of each objective function tend to the ideal value. The specific adjustment process of the parameters is well known to those skilled in the art and will not be elaborated here.

[0058] This embodiment introduces unbalanced abnormal text prompts, uses multiple defect words in the unbalanced abnormal text prompts, and under the guidance of the unbalanced abnormal text prompts, proposes a dual-branch diffusion model fine-tuning method, that is, using two branches, namely a defect image branch and a normal image branch, and further uses a hybrid training strategy to introduce defect images of different categories, thereby alleviating the overfitting problem that may be caused by the limited number of samples of each defect category, and increasing the diversity of defect images while still maintaining the authenticity of the defect images. Therefore, even when the number of defect images of each defect category in industrial scenarios is small, a relatively ideal training effect can still be achieved, so that defect images are generated using the trained diffusion model (that is, the trained defect image generation model).

[0059] It should be noted here that the above-mentioned use of the hybrid training strategy to train the U-Net encoder and the U-Net decoder can be understood as an overall training strategy, and the steps 201 to 203 can be understood as the actual actions performed in each training.

[0060] In actual application scenarios, the combination of multiple defect words and a single product word to obtain an imbalanced abnormal text prompt specifically includes:

[0061] Combine the product word of a single product with the N defect words of the product to obtain the imbalanced abnormal text prompt 'a <ob>with <df1> <df2>… <df N >'; <ob>is the product word, <df1> 、 <df2>、…、 <df N > is N defect tokens, where N is a preset value. Each product's N defect tokens correspond to a defect category. N is determined by those skilled in the art based on demand analysis. In an optional embodiment, N is 4. K*N defect tokens are used to generate K shared product tokens with unbalanced exception text prompts, covering all product defect categories. K is a preset value corresponding to K defect categories for a product. K is determined by those skilled in the art based on empirical analysis. For each defect category, an unbalanced exception text prompt composed of different defect tokens is used.

[0062] In an optional embodiment, as Figure 3 As shown, the hybrid training strategy is specifically as follows:

[0063] In step 301, a defect image set is generated based on defect images of all defect categories of the corresponding product. in, is the i-th defect image.

[0064] In step 302, a preset number of normal images of the product are selected to generate a normal image set. in, is the i-th normal image. Steps 301 and 302 can be understood as constructing a defect image set X df and the normal image set X ob process.

[0065] In step 303, from the defect image set X df Select the first number of defective images from the normal image set X ob A first number of normal images are selected from the training set, the first number of defective images and the first number of normal images are mixed to form a batch of data, and each batch of data is input into the defective image branch and the normal image branch for multi-batch training. The first number and the preset number are both determined by those skilled in the art based on empirical analysis. One batch of data is used for one batch of training, and the specific implementation process for each batch of training is as described in steps 201 to 203 above.

[0066] Based on the above defect image branch and normal image branch, this embodiment also provides a specific implementation of each objective function. For example, in an optional embodiment, the first objective function is in, is the cross attention map corresponding to the n-th defective word, N is the number of defective words, L is the number of U-Net layers corresponding to the cross attention map used in the alignment process, and M l It is a binary mask with a size of r×r pixels, with a value of 1 in the defect area and a value of 0 in the non-defect area. is the cross attention map corresponding to the product word, ⊙ is the pixel-by-pixel multiplication operation. In order to remove defective words <df n >The value of cross attention in defective areas is improved, This is to prevent product words from being <ob>Align to defect area.

[0067] The second objective function is Among them, e df represents the conditional vector obtained after the unbalanced abnormal text prompt passes through the text encoder, z df is the defect image x df The latent vector obtained by the VAE encoder, ∈ df is random noise, t df is the number of noise adding steps, To use ∈ df and ∈ df z df The noise-added latent vector obtained by adding noise is the predicted noise, is the value of the first objective function.

[0068] The third objective function is

[0069]

[0070] Among them, e ob represents the conditional vector obtained after the product-related part of the unbalanced abnormal text prompt passes through the text encoder, z ob is a normal image x ob The latent vector obtained by the VAE encoder, ∈ ob is random noise, t ob is the number of noise adding steps, To use ∈ ob and t ob z ob The noise latent vector obtained by adding noise is is the predicted noise.

[0071] Based on the above-mentioned embodiments, this embodiment also provides a diffusion model for generating a few-sample defect image. The diffusion model is obtained using the above-mentioned diffusion model fine-tuning method for generating a few-sample defect image. The diffusion model is used to generate a defect image. The diffusion model is also referred to as a defect image generation model in the above-mentioned embodiments.

[0072] Example 2:

[0073] The present invention is based on the method described in Example 1, combined with specific application scenarios, and uses technical descriptions in related scenarios to illustrate the implementation process of the present invention in characteristic scenarios.

[0074] like Figure 4 As shown, the diffusion model fine-tuning method for generating a few-sample defect image provided by this embodiment includes the following steps:

[0075] In step 401, a small number of defective images and corresponding defective image masks are selected as inputs of the defective image branch, and a fixed number of normal images are selected as inputs of the normal image sample branch.

[0076] In step 402 , an imbalanced abnormality text prompt is designed. Specifically, a plurality of defect words applicable to industrial scenarios and a single product word are combined to obtain an imbalanced abnormality text prompt, which is used as a guiding condition to control the generation of a diffusion model image.

[0077] In step 403, defect regions are aligned. Specifically, a defect image branch is designed to obtain features of the defect image. Based on the defect image mask corresponding to the defect image, the attention module aligns the defect tokens in the imbalanced abnormal text prompt to the defect region in the image. Finally, the noise prediction objective function for the imbalanced abnormal text prompt is optimized.

[0078] In step 404, the normal image is fine-tuned. Specifically, the normal image features are obtained in the design of the normal image branch. The product-related parts of the imbalanced abnormal text prompt are used to optimize the noise prediction objective function.

[0079] In step 405, a hybrid training strategy is implemented. Specifically, a unified U-Net is fine-tuned for each product containing different defect categories. The defect image set consists of a small number of defect images of all defect categories for the same product, while the normal image set consists of a fixed number of normal images of the product. Each time, an equal number of images are extracted from the defect image set and the normal image set, and mixed into a batch for training.

[0080] Among them, the imbalance abnormal text prompt design specifically includes: using the imbalance abnormal text prompt 'a <ob>with <df1> <df2>… <df N >', where <ob>is the product word, <df1> <df2>… <df N > is N defective tokens. The imbalanced abnormal text prompt of each product consists of a single product token and N defective tokens.

[0081] The defect area alignment specifically includes: decoupling the defect area as the foreground from the defect image through a cross attention mechanism and separating it from the defect word. <df n >Align and optimize the noise prediction objective function under unbalanced abnormal text prompts. Figure 5 As shown, specifically including:

[0082] In step 501, based on the features output by the U-Net encoder and the defect image mask, the defect word is divided into <df n >Align to the defect area in the defect image to optimize the first objective function:

[0083]

[0084] in, is the nth defective word <df n >corresponding cross attention map, N is the number of defective words. L represents the number of U-Net layers corresponding to the cross attention map used in the alignment process, M l It is a binary mask with a size of r×r pixels, where the value is 1 in the defect area and 0 in the non-defect area. Product word <ob>The corresponding cross attention map, ⊙ is the pixel-wise multiplication. This objective function decouples the defect area from the defect image. The first term decouples the defect word <df n >The cross attention value in the defect area is increased, and the second item is to prevent the product word from being <ob>Align to defect area.

[0085] In step 502, the defect image branch uses the imbalanced abnormal text prompt 'a <ob>with <df1> <df2>… <df N >', optimize and fine-tune the second objective function of U-Net:

[0086]

[0087] Among them, e df Represents the conditional vector obtained after the unbalanced abnormal text prompt passes through the CLIP text encoder, the defect image x df Get the latent vector z through the VAE encoder df , and use random noise ∈ df and the number of noise addition steps t df z df Add noise to get the noise latent vector Then, under condition e df Noise prediction is carried out under the guidance of and ∈ df Compute the loss for fine-tuning U-Net.

[0088] The normal image fine-tuning specifically includes: the normal image branch uses the product-related part of the unbalanced abnormal text prompt 'a <ob>', optimize and fine-tune the third objective function of U-Net:

[0089]

[0090] Among them, e ob Indicates the product-related part of the imbalance exception text prompt 'a <ob>'Conditional vector obtained after CLIP text encoder, normal image x ob Get the latent vector z through the VAE encoder ob , and use random noise ∈ ob and the number of noise addition steps t ob z ob Add noise to get the noise latent vector Then under condition e ob Noise prediction is carried out under the guidance of and ∈ ob Compute the loss for fine-tuning U-Net.

[0091] The hybrid training strategy specifically includes: using the hybrid training strategy to fine-tune a unified U-Net for each product. Given a product image set, which covers 15 products, each product corresponds to defect images with multiple defect categories and corresponding defect image masks and normal images. The defect images of all products are aggregated into a unified defect image set. For each defect category, an unbalanced abnormal text prompt consisting of different defect words is used. In addition, a fixed number of normal images are extracted to form a normal image set. In each fine-tuning step, from X df and X ob The same number of images are extracted from the dataset and combined into a batch for training. This not only alleviates the overfitting problem caused by the limited number of samples for each defect category, but also increases the diversity of defect images while still maintaining their authenticity.

[0092] The U-Net architecture obtained through the above training process is used in the diffusion model, which is used to achieve image generation under specific conditions. Given an input image x0, the diffusion model first converts it into a representation z=ε(x0) in the latent space, and then adds randomly sampled noise ∈~N(0,I) to z to obtain Where t is a randomly selected time step. Next, use U-Net to predict the noise ∈. Let c θ (P) is the CLIP text encoder, which maps the conditional text prompt P to the conditional vector e. The noise prediction objective function that the diffusion model needs to optimize is:

[0093]

[0094] in is the predicted noise.

[0095] This embodiment proposes an unbalanced abnormal text prompt generated for industrial defect images, which contains multiple defect words and a single product word. Furthermore, under the guidance of the unbalanced abnormal text prompt, a dual-branch diffusion model fine-tuning method is proposed. For the defect image branch, the present invention proposes a defect area alignment objective function to align the defect word and the defect area, and optimizes the noise prediction objective function under the unbalanced abnormal text prompt to achieve high diversity of defects. For the normal image branch, only the product-related part of the unbalanced abnormal text prompt is used to optimize the noise prediction objective function and synthesize product images with global consistency and local changes. The training process of the two branches is carried out separately, but is implemented on a shared U-Net, such as Figure 6 As shown, Representing textual prompts for each imbalanced anomaly, both the defect image branch and the normal image branch operate on the U-Net architecture. Furthermore, this embodiment employs a hybrid training strategy, fine-tuning a unified U-Net for each product. This allows for the generation of a large number of highly realistic and diverse defect images using a small number of defect images and a few normal images for training.

[0096] Example 3:

[0097] like Figure 7 FIG. 1 is a schematic diagram of the architecture of a diffusion model fine-tuning device for generating a small number of defect images according to an embodiment of the present invention. The diffusion model fine-tuning device for generating a small number of defect images according to this embodiment includes one or more processors 21 and a memory 22. Figure 7 A processor 21 is taken as an example.

[0098] The processor 21 and the memory 22 may be connected via a bus or other means. Figure 7 The bus connection is taken as an example.

[0099] Memory 22, as a nonvolatile computer-readable storage medium, can be used to store nonvolatile software programs and nonvolatile computer-executable programs, such as the diffusion model fine-tuning method for generating a few-shot defect image in Example 1. Processor 21 executes the diffusion model fine-tuning method for generating a few-shot defect image by running the nonvolatile software program and instructions stored in memory 22.

[0100] The memory 22 may include high-speed random access memory and non-volatile memory, such as at least one disk storage device, flash memory device, or other non-volatile solid-state memory device. In some embodiments, the memory 22 may optionally include a memory remotely located relative to the processor 21, and such remote memory may be connected to the processor 21 via a network. Examples of such networks include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0101] The program instructions / modules are stored in the memory 22 , and when executed by the one or more processors 21 , perform the diffusion model fine-tuning method for generating a few-sample defect image in the above-mentioned embodiment 1.

[0102] The embodiment of the present invention also provides a schematic diagram of the architecture of a diffusion model device for generating a small number of sample defect images. Figure 7 As shown, the processor carries a diffusion model for generating a few-sample defect image, and uses the diffusion model for generating a few-sample defect image to generate a defect image.

[0103] It is worth noting that the information interaction, execution process, etc. between the modules and units within the above-mentioned devices and systems are based on the same concept as the processing method embodiment of the present invention. The specific content can be found in the description of the method embodiment of the present invention and will not be repeated here.

[0104] Those skilled in the art will understand that all or part of the steps in the various methods of the embodiments can be completed by instructing related hardware through a program, and the program can be stored in a computer-readable storage medium, which may include: read-only memory (ROM), random access memory (RAM), a disk or an optical disk, etc.

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

Claims

1. A diffusion model fine-tuning method for generating few-shot defect images, characterized in that: The defect image generation model includes a U-Net encoder, a U-Net decoder, an attention module, a first objective function, a second objective function, and a third objective function. The U-Net encoder, U-Net decoder, attention module, the first objective function, and the second objective function constitute a defect image branch, while the U-Net encoder, U-Net decoder, and the third objective function constitute a normal image branch. The training process of the two branches is performed separately but implemented on a shared U-Net. The defect image generation model is trained using a hybrid training strategy, and the method includes: Combining multiple defect words and a single product word to obtain an unbalanced abnormal text prompt, and inputting the unbalanced abnormal text prompt and the product-related portion of the unbalanced abnormal text prompt as guidance conditions into the defect image branch and the normal image branch respectively; Selecting a few-sample defect image and a corresponding defect image mask as input to a defect image branch, the defect image branch obtains features of the defect image and aligns defect words in the imbalanced abnormal text prompt to the defect area in the image according to the defect image mask corresponding to the defect image, so as to optimize the first objective function and the second objective function; A fixed number of normal images are selected as inputs to a normal image sample branch. The normal image branch obtains features of the normal images and uses the product-related portion of the unbalanced abnormal text prompts to optimize a third objective function. The hybrid training strategy is used to allow defect images of different defect categories to participate in the training together. The first objective function is ;in, This is for the The cross attention map corresponding to the defective word, is the number of defective words, Indicates the number of U-Net layers corresponding to the cross attention map used in the alignment process, is a binary mask of size Pixel, the defect area value is 1, the non-defect area value is 0, is the cross attention map corresponding to the product word, To perform pixel-wise multiplication; The second objective function is ;in, represents the conditional vector obtained after the unbalanced abnormal text prompt passes through the text encoder, Defect image The latent vector obtained by the VAE encoder, is random noise, is the number of noise adding steps, For use and right The noise-added latent vector obtained by adding noise , is the predicted noise, is the value of the first objective function; The third objective function is ;in, represents the conditional vector obtained after the product-related part of the unbalanced abnormal text prompt passes through the text encoder, Normal image The latent vector obtained by the VAE encoder, is random noise, is the number of noise adding steps, For use and right The noise latent vector obtained by adding noise is is the predicted noise.

2. The diffusion model fine-tuning method for generating a few-sample defect image according to claim 1, characterized in that: The step of combining multiple defective words with a single product word to obtain an imbalanced abnormal text prompt specifically includes: Combine the product tokens of a single product with the product tokens of that product The defective word units are combined to obtain the unbalanced abnormal text prompt of the product 'a < > with < > < > … < >'; where < > is the product word, < >、< >, …, < >Yes defective words, For the default value, each product Defective words correspond to a defect category, using defective words Unbalanced abnormal text prompts of shared product words, covering all product defect categories, is the preset value, corresponding to a product defect categories.

3. The diffusion model fine-tuning method for generating a few-sample defect image according to claim 1, characterized in that: For each defect category, an imbalanced abnormal text prompt consisting of different defect tokens is used.

4. The diffusion model fine-tuning method for generating a few-sample defect image according to claim 1, characterized in that: The hybrid training strategy is specifically as follows: Building defect image sets and normal image set ; From the defect image set Select the first number of defective images from the normal image set A first number of normal images are selected, the selected first number of defect images and the first number of normal images are mixed into a batch of data, and the data of each batch is input into the defect image branch and the normal image branch for multi-batch training.

5. The diffusion model fine-tuning method for generating a few-sample defect image according to claim 4, characterized in that: The defect image set is constructed and normal image set , specifically including: Generate a defect image set based on defect images of all defect categories of the corresponding product ;in, is the i-th defect image; Select a preset number of normal images of the product to generate a normal image set ;in, is the i-th normal image.

6. A diffusion model fine-tuning device for generating a few-sample defect image, characterized in that: include: at least one processor; And, a memory communicatively connected to the at least one processor; wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the processor to execute the diffusion model fine-tuning method for generating a few-sample defect image as described in any one of claims 1-5.

7. A non-volatile computer storage medium, characterized in that The computer storage medium stores computer-executable instructions, which are executed by one or more processors to implement the method according to any one of claims 1 to 5.

Images