Industrial defect sample generation method based on small sample fine tuning and controllable diffusion model

Abstract

The invention provides an industrial defect sample generation method based on small sample fine tuning and a controllable diffusion model, and the method comprises the steps: employing a surface defect image small sample data set, and screening out a model fine tuning data set through employing a pre-constructed image distribution evaluation module; finely adjusting the weight of the pre-trained stable diffusion model; an industrial defect sample conforming to the specified defect feature is generated by combining a trained space condition control neural network with a preset identifier as a cue word; and calculating a distribution difference value between the industrial defect generation sample and the original small sample data set, selecting images with the distribution difference value smaller than a preset threshold value, and adding the selected images into the original small sample data set to form an expanded training data set. According to the method, the small sample data set is utilized, the diffusion model weight is finely adjusted, the cue word containing the preset identifier, the space condition and the space condition control neural network are combined, the training samples are expanded, the calculation cost is reduced, and the application range is wide.

Description

TECHNICAL FIELD

[0001] Embodiments of the present disclosure relate to the technical field of computer vision, and particularly relate to an industrial defect sample generation method and system based on small sample fine-tuning and controllable diffusion model and a computer readable storage medium. BACKGROUND

[0002] In the process of industrial production, artificial detection is a traditional method for detecting surface defects of products. This method has low sampling rate, low accuracy, poor real-time performance, low efficiency, high labor intensity, and is greatly affected by human experience and subjective factors. Defect detection methods based on machine vision can largely overcome the above drawbacks.

[0003] Data plays a crucial role in deep learning-based defect detection algorithms. Detection models require a large amount of data to achieve good accuracy, robustness, and generalization. In the case of insufficient data and class imbalance, the model may overfit and other problems may occur. Industrial defect detection based on supervised deep learning may be hindered by multiple factors such as data quality, data scarcity, and data annotation difficulty.

[0004] Industrial defect datasets generally face the problems of small sample and class imbalance. At the same time, the type, location, and size of defects are often specific and unique. Existing image generation techniques based on stable diffusion models are generally applied in the natural and artistic fields. There is little research on defect image generation in the industrial field, and it is difficult to achieve precise control using current stable diffusion models. At the same time, training a large-capacity generation model from scratch requires high computational cost, data demand, and professional knowledge. SUMMARY

[0005] The purpose of the embodiments of the present disclosure is to provide an industrial defect sample generation method, system and computer readable storage medium based on small sample fine-tuning and controllable diffusion model, thereby solving the aforementioned problems existing in the prior art.

[0006] In order to achieve the above-mentioned purpose, the technical solutions adopted by the embodiments of the present disclosure are as follows: The embodiments of the present disclosure propose an industrial defect sample generation method based on small sample fine-tuning and controllable diffusion model, which comprises: According to the small sample dataset of surface defect images, an image distribution evaluation module is used to select a model fine-tuning dataset; The pre-trained stable diffusion model is weight-tuned according to the model fine-tuning data set, and a class prior loss function is used to jointly optimize an image reconstruction loss and a pre-training knowledge preservation loss during the fine-tuning process; and the defect feature is bound to the generation model weight through a special identifier, wherein the fine-tuning data set comprises defect images and corresponding labels, and the labels comprise a preset identifier and a defect type description. According to the fine-tuned stable diffusion model and the trained spatial condition control neural network, a prompt word containing a preset identifier and a spatial condition specifying a defect type, size and position are input, a hierarchical control signal is generated through the spatial condition control neural network, an intermediate feature representation of the stable diffusion model in the denoising process is adjusted, and an industrial defect sample image conforming to the specified defect feature is generated.

[0007] Optionally, after the industrial defect sample image conforming to the specified defect feature is generated, the method further comprises: calculating a distribution difference value of the industrial defect sample and an original small sample data set, selecting an image with a distribution difference value less than a preset threshold to add to the original small sample data set to form an expanded training data set.

[0008] Optionally, the image distribution evaluation module comprises a convolutional neural network and a Fréchet function. The model fine-tuning data set is selected from the surface defect image small sample data set by using a pre-constructed image distribution evaluation module, and comprises: A convolutional neural network (CNN) is used to extract feature vectors of each image in the small sample data set; The feature vectors are subjected to L2 normalization processing to obtain normalized feature vectors; A preset number of normalized feature vectors are randomly sampled to form a plurality of groups, the Fréchet distance between the feature distribution of each group and the overall distribution of the small sample data set is calculated, and the distribution difference value of each group of data is obtained, and the group data with a distribution difference value less than a preset difference is selected as the model fine-tuning data set.

[0009] Optionally, the pre-trained stable diffusion model is weight-tuned according to the model fine-tuning data set, and a class prior loss function is used to jointly optimize an image reconstruction loss and a pre-training knowledge preservation loss during the fine-tuning process; and the defect feature is bound to the generation model weight through a special identifier, wherein the fine-tuning data set comprises defect images and corresponding labels, and the labels comprise a preset identifier and a defect type description. The pre-trained stable diffusion model and the model fine-tuning data set are loaded, and the weights of the CLIP text encoder and the VAE are frozen; wherein the model fine-tuning data set comprises defect images and corresponding labels, and the labels comprise at least a preset identifier and a defect type description; Noise is added through a diffusion process, the model is trained to predict the noise, an image reconstruction loss and a class prior loss are jointly optimized, and the UNet partial layer weights are updated through back propagation; Save the weights to obtain the fine-tuned stable diffusion model.

[0010] Optionally, the formula of jointly optimizing the image reconstruction loss and the class prior loss is as follows:

[0011] In the formula: α t , σ t , and w t are respectively a term for controlling the noise diffusion schedule and the sample quality at the moment, α t x + σ t combining the real image with the noise according to the diffusion schedule, c is a text condition, λ is a scalar weight; [] represents the mathematical expectation of the entire loss on the distribution of the real image x , the text condition c , the noise sampling , and the time step t ; represents the reconstruction loss of the user input image; is the class prior loss; x is an image containing a target concept, i.e., a real industrial defect; x pr is a class prior image; c pr a class prior text; t and respectively represent the time step; 、 respectively represent the sampled Gaussian noise, which is subject to ; α t and σ t respectively represent the weight parameters related to the diffusion process; and w t respectively represent the time step and the corresponding weighting coefficients; represents the square of the two-norm; θ represents a UNet noise prediction model with parameters θ, and θ represents a set of parameters that need to be fine-tuned for the UNet noise prediction model.

[0012] Optionally, the spatial condition control neural network comprises: a condition encoding module, which converts the spatial condition into a structured tensor; an encoder comprising a four-level convolutional layer, which down-samples the condition graph to extract a feature graph; a triple residual block, which obtains a feature flow from the feature graph through the residual block to enhance gradient flow; a zero convolutional layer group, which converts the feature flow output by the residual block into a unified high-dimensional tensor.

[0013] Optionally, the generation of the hierarchical control signal by the spatial condition control neural network adjusts the intermediate feature representation of the stable diffusion model in the denoising process, comprising: encoding the spatial condition into a spatial condition graph, wherein the spatial condition graph comprises: defect type, and position and size; down-sampling the spatial condition graph by the encoder to obtain a feature graph, and outputting a hierarchical control signal from the feature graph through the residual module and the zero convolutional layer group; injecting the primary control signal into the shallow layer of the denoising network; injecting the intermediate control signal into the middle layer of the denoising network; and injecting the high-level control signal into the deep layer of the denoising network.

[0014] Optionally, the injection of the primary control signal into the shallow layer of the denoising network; the injection of the intermediate control signal into the middle layer of the denoising network; and the injection of the high-level control signal into the deep layer of the denoising network, comprise: the primary control signal is injected into the 3rd layer of UNet, i.e., the shallow layer feature, through a residual connection, and acts on the spatial detail dimension of the feature graph; the intermediate control signal is injected into the 5th layer of UNet, i.e., the middle layer feature, and acts on the intermediate abstract representation space; the high-level control signal is injected into the 7th layer of UNet, i.e., the deep semantic feature space, and acts on the global semantic representation.

[0015] In one aspect, the disclosure provides an industrial defect sample generation system based on small sample fine-tuning and controllable diffusion model, comprising: a selection module configured to select a model fine-tuning data set from a small sample data set of surface defect images by using a pre-constructed image distribution evaluation module; a fine-tuning module configured to fine-tune the weight of a pre-trained stable diffusion model according to the model fine-tuning data set, and jointly optimize an image reconstruction loss and a pre-training knowledge preservation loss by using a class prior loss function during the fine-tuning process; and bind the defect feature to the weight of the generation model through a special identifier, wherein the fine-tuning data set comprises: a defect image and a corresponding label, and the label comprises: a preset identifier and a defect type description; ​The generating module is configured to input the prompt word containing the preset identifier and the spatial condition specifying the defect type, size and position into the fine-tuned stable diffusion model and the trained spatial condition control neural network, generate a hierarchical control signal through the spatial condition control neural network, adjust the intermediate feature representation of the stable diffusion model in the denoising process, and generate an industrial defect sample image conforming to the specified defect feature.

[0016] In another aspect, the present disclosure provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the steps of the method described above.

[0017] In another aspect, the present disclosure provides a computer program product comprising a computer program, which, when executed by a processor, implements the steps of the method described above.

[0018] The present disclosure has the following beneficial effects: The method of the present disclosure utilizes small sample data to fine-tune the diffusion model weight, combines the prompt word containing the preset identifier, the spatial condition and the spatial condition control neural network, expands the training sample, reduces the calculation cost, and has a wide application range; the prior loss function supervises the consistency of the fine-tuned model and the pre-trained model. BRIEF DESCRIPTION OF DRAWINGS

[0019] Figure 1 is a flowchart of an industrial defect sample generation method based on small sample fine-tuning and controllable diffusion model according to an embodiment of the present disclosure; Figure 2 is a flowchart of an industrial defect sample generation method based on small sample fine-tuning and controllable diffusion model according to an embodiment of the present disclosure; Figure 3 is a schematic diagram of the role of the class loss function in the industrial defect sample generation method based on small sample fine-tuning and controllable diffusion model according to an embodiment of the present disclosure; Figure 4 is a flowchart of the spatial control neural network model in the industrial defect sample generation method based on small sample fine-tuning and controllable diffusion model according to an embodiment of the present disclosure; Figure 5 is a structural diagram of an industrial defect sample generation system based on small sample fine-tuning and controllable diffusion model according to an embodiment of the present disclosure. DETAILED DESCRIPTION

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

[0021] The embodiment of the present disclosure proposes a scheme capable of controllable sample generation through diffusion mode in the case of insufficient industrial product defect data and class imbalance. The embodiment of the present disclosure relates to an industrial defect sample generation technology based on small sample fine-tuning and controllable diffusion model, and belongs to the field of artificial intelligence generation, and is particularly applied to the detection model algorithm training process of industrial defects.

[0022] As shown in Figure 1 The embodiment of the present disclosure proposes a method for generating industrial defect samples based on small sample fine-tuning and controllable diffusion model, which comprises the following steps: Step S100, according to the pre-acquired small sample data set of surface defect images, a pre-constructed image distribution evaluation module is used to select a model fine-tuning data set.

[0023] The method of the embodiment of the present disclosure before step S100 further comprises: acquiring a small sample data set with surface defect images, which can be a small sample data set of industrial product surface defects. Small sample refers to a sample set with a sample capacity less than N. N can be 1500 or 2000, etc., and N is less than 10000.

[0024] The purpose of step S100 of the embodiment of the present disclosure is to evaluate the industrial defect small sample original data set by using the image distribution evaluation module according to the selected fine-tuning training samples, that is, the small sample data set of images with surface defects, to obtain the distribution difference value of each image, that is, the deviation degree of each image from the overall distribution. Select the image sample with smaller distribution difference value (that is, the sample most consistent with the global distribution) as the model fine-tuning data set.

[0025] Exemplarily, the image distribution evaluation module comprises a convolutional neural network and a Fréchet function.

[0026] The use of the pre-constructed image distribution evaluation module to select the model fine-tuning data set comprises: Step S110, a convolutional neural network CNN is used to extract the feature vector of each image in the small sample data set.

[0027] A Convolutional Neural Network (CNN) consists of convolutional layers, pooling layers, normalization and regularization layers, and activation layers. The final classification layer is removed, and the output is the feature vector of the intermediate layers, such as an N-dimensional vector after global average pooling. A CNN can be Inception v3, using a pre-trained Inception v3 model as a feature extractor, removing its final classification layer (Softmax) and retaining the global average pooling layer (outputting a 2048-dimensional feature vector). This model, pre-trained on ImageNet, effectively captures common visual features of industrial defects (such as texture and shape). Alternatively, the backbone network of a CNN (such as ResNet or VGG) pre-trained on a large dataset (such as ImageNet) can be used directly to extract image features (such as the 2048-dimensional vector output by ResNet-50), which will not be elaborated upon here.

[0028] Convolutional neural networks extract image features: v m = f cnn ( P m ); Among them, v m It is the extracted m-th image feature vector. P m This represents the m-th image in the model fine-tuning dataset. f cnn This represents the feature extraction process of a convolutional neural network. m ∈ M , M This represents the total number of samples in the model fine-tuning dataset.

[0029] Step S120: Perform L2 normalization on the feature vectors of each image to obtain the normalized feature vectors.

[0030] The normalization formula is as follows:

[0031] Where v is the original feature vector, v i The eigenvector of the eigenvector i Each dimension value; v norm These are the L2-normalized eigenvectors. The eigenvectors are scaled or shifted to a standard range to eliminate dimensional differences and improve the stability of the covariance matrix calculation.

[0032] The purpose of normalization is: 1. To eliminate dimensional differences: to prevent certain feature dimensions from dominating distance calculations due to excessively large numerical ranges (such as Fréchet distance and cosine similarity). 2. To improve numerical stability: to avoid ill-conditioned problems caused by inconsistent feature scales when calculating the covariance matrix.

[0033] The normalized feature data is obtained after standardization, which aims to eliminate the differences in dimension, range or distribution of the original features, so as to make different features comparable, thereby improving the stability and effect of subsequent calculation (such as distance measurement, model training).

[0034] For example, if the original feature vector is [3, 4], its L2 norm is , and after normalization, it is [0.6, 0.8].

[0035] Step S130, randomly sample a preset number of normalized feature vectors, divide them into multiple groups, calculate the Fréchet distance between the feature distribution of each group and the overall distribution of the small sample data set, and obtain the distribution difference value of each group of data. Select the group data with a distribution difference value less than the preset difference as the model fine-tuning data set.

[0036] For industrial defect small sample data set, feature extraction is performed by convolutional neural network CNN, and after normalization, random grouping is performed, the mean and covariance matrix of each group of feature vectors are calculated, and then the Fréchet distance between each group of feature vectors and the feature distribution of the small sample data set is calculated to obtain the distribution difference value. Select the small group of images with smaller distribution difference value as the model fine-tuning data set.

[0037] The convolutional neural network CNN used in step S100 is used for feature extraction. Specifically, a convolutional layer module is used for feature extraction. Since the deep convolutional neural network learns general features such as edges, textures, and corner points at the low level, this model can extract useful low-level features.

[0038] The process of FID evaluation is as follows: first, randomly sample the feature vectors of the data set into multiple groups, for example, 10-20 images per group. Calculate the Fréchet distance between the overall feature distribution (mean and covariance matrix) of each group and the small sample original data set. Select the image group with higher FID score as the model fine-tuning data set.

[0039] Then, the extracted feature vectors are calculated by Fréchet distance, which measures the distance between image distributions, providing a quantitative indicator for measuring the quality of the generated model. The smaller the value, the closer the generated sample to the real sample, and the more consistent the training sample with the overall sample distribution. The Fréchet distance distribution difference value used is d FID p 1, p 2), the specific formula is:

[0040] p ​1, p 2 Gaussian distribution of the randomly sampled group and the pre-acquired small sample dataset of surface defect images respectively; is the standard deviation of the data segment; denotes the trace of a matrix. ∑1, ∑2 are the feature covariance matrices of the randomly sampled group and the small sample dataset respectively; are the mean vectors of the feature vectors of the randomly sampled group and the small sample dataset respectively, i.e., the center of the feature space; denotes the trace of a matrix.

[0041] Selecting d FID ( p 1, p 2) Group image data less than the preset threshold value as a model fine-tuning dataset. Selecting the group image with a difference value less than the threshold value and adding it to the original dataset to form an expanded dataset (step S500).

[0042] Step S200, according to the model fine-tuning dataset, the weight of the pre-trained stable diffusion model is fine-tuned, and the prior loss function is used to jointly optimize the image reconstruction loss and the pre-training knowledge preservation loss in the fine-tuning process; the defect features are bound to the generation model weight through a special identifier, wherein the fine-tuning dataset includes: defect images and corresponding labels, and the labels include: a preset identifier and a defect type description.

[0043] As Figure 3 shown, the specific process includes: freezing the CLIP text encoder and VAE decoder weight of the pre-trained stable diffusion model; using the prior loss function to jointly optimize the following two items: image reconstruction loss: constraining the model to learn defect image features; pre-training knowledge preservation loss: constraining the output distribution to be consistent with the original model through the prior image; and binding the defect features to the UNet weight through a special identifier embedded in the text label.

[0044] The embodiment of the disclosure uses the model fine-tuning dataset in step S100 to fine-tune the weight of the existing pre-trained stable diffusion model in step S200, wherein the prior loss function is used in the fine-tuning process to preserve prior knowledge, and the defect features are bound through a special identifier to generate a personalized industrial defect model. In the process of generating images, through the prompt word containing the special identifier, an image with specific defect features can be generated. The special identifier is a preset identifier. The special identifier is set before the weight is fine-tuned. The defect features are bound to the special identifier.

[0045] The fine-tuning dataset includes: defect images and corresponding labels, and the labels include a special identifier and a defect type description, and the special identifier can be set according to actual conditions.

[0046] The fine-tuning dataset provides industrial defect samples, such as defect images and corresponding labels, for learning new features. The image reconstruction loss is used to calculate the image reconstruction loss, ensuring that the model learns to generate defects. The fine-tuning process is constrained by the class prior loss function, and the pre-trained model is used to generate general images unrelated to the industrial field as a distribution consistency comparison benchmark, preventing the model from forgetting the original capabilities. The class prior loss is calculated between the general image and the fine-tuning model generation result, and the feature distribution of the two is consistent. The pre-trained model is a pre-trained stable diffusion model. By jointly optimizing the loss of these two types of inputs, the pre-trained knowledge is retained while adapting to the industrial defect features. The special identifier is a unique text label that is strongly associated with the defect category during the fine-tuning process, and in the generation stage, the special defect generation is triggered by inputting the prompt word containing the identifier.

[0047] The pre-trained stable diffusion model (such as Stable Diffusion) is usually trained based on a large-scale general dataset (such as LAION). Fine-tuning refers to further adjusting the model weights based on the pre-trained model using data specific to a particular field (such as industrial defect images) to adapt to new tasks. The pre-trained stable diffusion model can be considered as an initial stable diffusion model.

[0048] During training, the defect feature is bound to the special identifier, uniquely binding the new defect concept and avoiding contamination of the pre-trained semantics. In the generation stage, the special defect generation is triggered by inputting the prompt word containing the special identifier and the defect type description.

[0049] For example, the pre-trained stable diffusion model is weight fine-tuned according to the model fine-tuning dataset, and the class prior loss function is used to jointly optimize the image reconstruction loss and the pre-trained knowledge preservation loss during the fine-tuning process. The defect feature is bound to the generation model weight through the special identifier, including: Step S210, load the pre-trained stable diffusion model and the model fine-tuning dataset, and freeze the weights of the CLIP text encoder and the VAE; wherein the model fine-tuning dataset includes: defect images and their corresponding labels, and the labels at least include: a preset identifier and a defect type description, so that the preset identifier is strongly associated with the defect category.

[0050] The model fine-tuning dataset in step S100 includes: defect images and their corresponding labels, and in the labels of the model fine-tuning dataset, the defect category is associated with the special identifier. For example, the label is "product surface has <gjjm>Crack”, where the product surface has a crack as a defect type description, GJJM as a special identifier. Each defect type corresponds to a special identifier. The data needs to be pre-processed into the model input format, such as 512x512 resolution. Load the open-source pre-trained Stable Diffusion model, such as the weights of stable-diffusion-v1-5, including: its UNet, CLIP text encoder, and VAE parts. Fine-tune the U-Net denoising network weights.

[0051] Special identifier (such as <defect-type>) is a placeholder to trigger the model's generative capabilities for specific defect features at the generation stage. Its role is similar to the [V] identifier in DreamBooth, which binds the new concept (scratch) with a unique identifier. Special identifiers (such as <defect>) is a key design in fine-tuning, to bind new concepts without disturbing the original knowledge.

[0052] After fine-tuning is completed, the user inputs a prompt word containing a special identifier, such as "product surface has <gjjm>"Cracks," the model generates defect images based on the features learned through fine-tuning.

[0053] Step S220: Add noise through the diffusion process, train the model to predict the noise, jointly optimize the image reconstruction loss and class prior loss, and backpropagate to update the weights of some UNet layers.

[0054] It should be noted that step S220 adds noise through the diffusion process and trains the model to predict the noise. The diffusion loss and the class-prior loss are jointly optimized to train the pre-trained stable diffusion model.

[0055] Specifically, for each batch, noise is added and predicted, and this is applied to the actual defect images. x Gaussian noise is added according to the diffusion time step t to generate a noisy image; a UNet network is trained to predict the noise based on the noisy image, time step t, and text conditions (including labels with preset identifiers); the loss is optimized using a joint loss function, the total loss function including: image reconstruction loss and... λ· Prior loss function L pr .

[0056] like Figure 3 As shown, the specific formula for the total loss function is:

[0057] in, α t , σ t ,and w t They are respectively Items that control the noise propagation timeline and sample quality at all times. α t x + σ t To combine real images with noise according to the diffusion timeline, c For text conditions, λ·L pr , L pr For prior preservation terms, λ Scalar weights. [] represents the total loss in the real image. x Text conditions c Noise sampling , and time step t The mathematical expectation on the distribution represents the loss over the entire real image. x Text conditions c Noise sampling , and time step t mathematical expectation (average) of the distribution at the top, The entire expression wrapped in [] is the weighted MSE loss function for user images and class prior images, which is used to stabilize model training and prevent overfitting to user images. The first term represents the reconstruction loss for user input images, which calculates the prediction error of the model for the real noise , by minimizing this error, the model learns how to reconstruct the target concept x from the noise, the MSE of the predicted noise and the real noise, by forcing the model to accurately predict the noise added to the real image, the model learns to reconstruct the unique features of the target image (user image) from the noise. The second term λ·L pr , i.e. is the class prior loss, which constrains the output to be consistent with the pre-trained model distribution, and constrains the model not to destroy the original class distribution when learning user concepts. x represents the image containing the target concept, i.e. the real industrial defect; x pr The class prior image (prior image) is a general class image generated by the original, un-tuned pre-trained model according to the class prior text c pr ; c The condition text is used to describe the concept provided by the user, for example "a photo of [V] dog"; c pr The class prior text is usually a general class text, for example "a photo of a dog"; t and respectively represent the time step (time step index for diffusion process); and respectively represent the sampled Gaussian noise, which is subject to ; α t and σ t represent the weight parameters related to the diffusion process, which determine the proportion of signal and noise at the current time step; and w t respectively represent the time step t´ and t corresponding weight coefficients (usually noise scheduler related terms); represents the square of the two-norm, i.e. mean square error (MSE); θ The UNet noise prediction model represented by a parameter θ, θ represents a set of parameters that need to be fine-tuned for the UNet noise prediction model, θ The input is a noisy latent variable α t x + σ t , a time step t, and a condition c, and the output is a predicted noise tensor. Before starting fine-tuning, training data needs to be prepared for the class prior loss. Using the pre-trained diffusion model with original weights and the sampler that is not fine-tuned, generate a batch of class prior image data according to the class prior text c pr Generate a batch of class prior image data x pr = (z t1 , c pr ) The model is provided with random initial noise and a condition vector. The second term is a prior protection term, which is used to supervise the model with self-generated images, λ control the relative weight of this term. The generated model represented by G is responsible for mapping the latent variable to the pixel space. z t1 represents the latent variable at time step t.

[0058] The role of the class prior loss function is to solve the catastrophic forgetting problem of the diffusion model when learning new concepts, so that the diffusion model can learn the concept of industrial defects while not forgetting the original knowledge. Under a small amount of data, both defect feature generation and original generation ability are learned.

[0059] The mechanism of the class prior loss is to constrain the noise prediction results of the fine-tuned model on the class prior image to be consistent with the pre-trained model: feature matching: constrain the feature distribution of the fine-tuned model to be close to the pre-trained model (such as measuring the difference of UNet intermediate layer features by KL divergence). Output consistency: the outputs of the model before and after fine-tuning should remain similar for the same input (text + noise) (such as calculating the LPIPS distance of the generated image).

[0060] Before fine-tuning the pre-trained stable diffusion model in step S220, the core part of the pre-trained model is frozen, such as the CLIP text encoder of Stable Diffusion and the weight of VAE. During the fine-tuning process of the model, the weights of CLIP and VAE will be frozen first.

[0061] Step S230, save the weights to get the fine-tuned stable diffusion model.

[0062] Fine-tuned UNet or LoRA weights can be saved separately to obtain personalized industrial defect models.

[0063] Pre-trained models are used to generate high-quality samples to assist the fine-tuning process (similar to self-distillation).

[0064] Frozen model: CLIP, VAE, etc. parts of the pre-trained model are frozen throughout the process, and only the UNet part can be trained. Pre-sampling: in each round of training, use the current model (in fine-tuning) and the frozen pre-trained model to generate samples No image generation. Loss calculation: task loss: difference between supervised real defect data and model prediction. Prior loss: consistency between supervised fine-tuned model and pre-trained model generation results. Parameter update: only update the weights of the unfrozen part (such as LoRA parameters of UNet).

[0065] The class prior loss used in step S200 limits the generated image from deviating from the common attributes of the class by comparing the difference between the feature distribution of the generated image and the distribution of the target class. This loss usually works together with other generation losses (such as image reconstruction loss, adversarial loss, etc.). The role of the class prior loss includes: 1) help the generation model maintain prior knowledge of a specific class during training, especially during fine-tuning, it avoids the model relying completely on the features of a small number of samples, ensuring that the generated images have stronger class consistency and universality; 2) the class prior loss encourages the model to generate images that conform to the general features of the class, rather than simply overfitting based on a small number of input samples. This means that even if the number of input training samples is very small, the model can still understand the overall distribution of the class and generate diversified images that conform to the characteristics of the class.

[0066] Step S300, according to the fine-tuned stable diffusion model and the trained spatial condition control neural network, input the prompt word containing the preset identifier and the spatial condition specifying the defect type, size and position, generate a hierarchical control signal through the spatial condition control neural network, adjust the intermediate feature representation of the stable diffusion model in the denoising process, and generate an industrial defect sample image conforming to the specified defect characteristics, wherein the spatial condition control neural network comprises: a condition encoding module, an encoder, a triple residual block and a zero convolution layer group.

[0067] Using the fine-tuned stable diffusion model combined with the control network weight, and using the special identifier in step S200 as the prompt word, an industrial defect sample conforming to the specified defect characteristics is obtained, and the special identifier is the preset identifier.

[0068] Based on the trained spatial condition control neural network, the intermediate feature representation of the stable diffusion model in the diffusion process can be adjusted by a hierarchical control signal according to the input geometric conditions such as defect type, position and size, so as to affect the image generated by the large neural network, that is, the stable diffusion model. In the embodiments of the present disclosure, the network to be controlled is the stable diffusion network after learning new concepts. In the whole process, first, the model weight is frozen to retain the initial weight of the model during training. Then, the spatial condition control neural network, the adapter network of ControlNet, which is a lightweight trainable branch attached to the frozen main model, is adjusted, and it is used as a conditional adapter to adjust the frozen copy. The spatial condition control neural network is a structural copy of the main model UNet, but the weight is independent and trainable. The structure of ControlNet is to copy part of the layers of the main model (such as UNet), then connect these copied layers with the corresponding layers of the main model, and inject control signals through the zero convolution mechanism. The input conditions are the rough geometric shapes of the defect content and the images of the defect categories and shapes.

[0069] It should be noted that the spatial condition control neural network is independently trained. A spatial condition control neural network is trained to adjust the step-by-step denoising process of generating an image according to the input conditions of the type, size and position of the specified defect, that is, to adjust the intermediate feature representation in the diffusion process, so as to control the step-by-step denoising process of the image, so that the generated image meets the predetermined conditions, thereby generating an industrial defect sample that meets the specified defect characteristics.

[0070] A trainable control path parallel to the pre-trained diffusion model is generated through external condition input to realize pixel-level spatial constraint. The essence is to construct a conditional adapter network to inject precise control signals without damaging the original model generation capability.

[0071] Step S200 outputs the fine-tuned stable diffusion model as a basic model for generating defect images.

[0072] The spatial condition control neural network is an independent lightweight network (such as ControlNet or similar structure) responsible for dynamically adjusting the generation process of the stable diffusion model according to the user input conditions such as defect type, size and position.

[0073] As shown in Figure 4 The implementation of the condition control of the spatial condition control neural network, that is, generating a hierarchical control signal through the spatial condition control neural network to adjust the intermediate feature representation of the stable diffusion model in the diffusion process, includes: Step S310, encode the spatial condition into a spatial condition graph, wherein the spatial condition graph includes: defect type, and position and size.

[0074] The conditional encoding module encodes the user-specified spatial conditions (such as "scratch, 5px long, located at the top left corner of the image") into a spatial condition map, for example: the defect type is represented by a class embedding vector; the location and size are marked by a binary mask or a heatmap, where the defect area is 1 and the background is 0. The defect location can be given by a bounding box or a segmentation mask, which is then converted into a binary map of fixed resolution.

[0075] In step S320, the spatial condition map is down-sampled by an encoder to obtain a feature map, and the feature map is outputted through a residual module and a zero convolution layer group to obtain a hierarchical control signal.

[0076] The spatial condition control neural network takes the above-mentioned condition map (mask map, class embedding) as input, and first enters an encoder for down-sampling. The encoder includes four convolution layers, which gradually compress the features through the four convolution layers: the first layer extracts the surface basic texture, the second layer captures the contour structure, the third layer uses a dilated convolution to expand the receptive field to 51x51 pixels to cover larger features, and the fourth layer compresses the feature map to 32x32 resolution and encapsulates the global defect distribution information. This process realizes the physical mapping from the pixel space to the high-dimensional feature space, and the four convolution layers of the encoder realize the spatial resolution compression (512 to 32), and the receptive field is gradually expanded, serving as the UNet regulation signal of the fine-tuned stable diffusion model.

[0077] The feature map outputted by the encoder is inputted into a residual module in sequence to obtain a feature stream, and the gradient flow is enhanced. The residual module includes a triple residual block.

[0078] The feature stream enters a zero convolution layer group, which converts the feature stream outputted by the residual block into a unified high-dimensional tensor as the original control signal. The weights and biases of these 1x1 convolutions are initialized to zero. In the initial state, the module always outputs a zero tensor, ensuring that it does not interfere with the knowledge system of the pre-trained UNet in the early stage of training. As the training progresses, the zero convolution layer gradually learns to convert the control signal into a physical offset in the feature space. The tensor is then decomposed into shallow, medium, and deep control signals by a hierarchical projection module (composed of three parallel 1x1 convolutions), which correspond to the 3rd, 5th, and 7th layers of the denoising network of UNet, respectively.

[0079] The shallow signal (64x64 resolution), i.e., the primary signal, is converted into a spatial attention mask, and a binary heat map is generated by Sigmoid activation to guide the model to enhance the defect texture details in the mask area (value≈1). The middle layer signal (32x32 resolution), i.e., the intermediate signal, is converted into a structure scaling coefficient, and the feature map amplitude is adjusted by channel multiplication to control the defect geometry. The deep layer signal (16x16 resolution), i.e., the advanced signal, is converted into a semantic channel weight, and a channel dimension scaling vector is generated by a fully connected layer to adapt to the high-level semantics of the defect category.

[0080] In step S330, the primary control signal is injected into the shallow layer of the denoising network; the intermediate control signal is injected into the middle layer of the denoising network; and the advanced control signal is injected into the deep layer of the denoising network. The primary control signal is also a detail control signal, the intermediate control signal is also a structure control signal, and the advanced control signal is also a semantic control signal.

[0081] In each step of the diffusion model denoising process, the spatial condition control network encodes the specified defect type, location, and size, etc. spatial conditions into a feature tensor, and generates hierarchical control signals (detail, structure, semantic signals). When UNet performs forward propagation for noise prediction, these signals are respectively injected and fused into the feature maps of the shallow layer (3rd layer), middle layer (5th layer), and deep layer (7th layer) of the UNet decoder through adaptive instance normalization, channel multiplication, and cross attention mechanisms, to achieve spatial control of the generated image content at the pixel level.

[0082] That is, in each time step t of the above denoising process, the spatial condition control neural network will generate a set of hierarchical control signals according to the input spatial condition map. These signals are respectively injected into the 3rd, 5th, and 7th layer feature maps of the UNet network before it performs noise prediction calculation at the current step, thereby incorporating spatial control conditions into the noise prediction at this time step, and gradually guiding the latent variable to evolve towards an image that meets the predetermined spatial features.

[0083] The adjustment of the intermediate feature representation includes: injecting a detail control signal in the denoising network 3rd layer feature map, realizing the migration of the style (mean and variance) of the control signal to the features of the UNet through adaptive instance normalization (AdaIN), and regulating the edge / texture details. Inject a structure control signal in the denoising network 5th layer feature map, realize through channel multiplication, control the structure geometry, specifically, the channel-by-channel scaling (Scale), that is, the control signal is a vector with the same number of channels as the feature map, and the feature map is multiplied by channel by channel to enhance or suppress the activation strength of certain feature channels. Inject a semantic control signal in the denoising network 7th layer feature map, realize through a cross-attention mechanism, maintain semantic consistency, specifically, use the control signal as Key and Value, and use the UNet features as Query to perform attention calculation, so that semantic information can globally modulate the feature map.

[0084] Specifically, the control signal is injected into the stable diffusion model through the following way: hierarchical injection mechanism of the control signal: UNet 3rd layer (resolution 64x64): The primary control signal output by the zero convolution layer is connected through a residual connection Inject the shallow features of this layer, which act on the spatial detail dimension of the feature map. This level of injection mainly regulates the local high-frequency features of the generated image, which directly affects the fine structure of the edge texture.

[0085] UNet 5th layer (resolution 32x32): The intermediate control signal output by the zero convolution layer is injected into the intermediate layer features of the network The intermediate abstract representation space. This level is responsible for coordinating the combination relationship of local features and constraining the geometric shape of the structure within the medium receptive field.

[0086] UNet 7th layer (resolution 16x16): The high-level control signal output by the zero convolution layer is injected into the deep semantic feature space, which acts on the global semantic representation. This level regulates the overall semantic attributes to ensure that the generated content is consistent with the high-level concept.

[0087] The training of the spatial condition control neural network is based on the fine-tuned stable diffusion model, and the spatial condition control neural network training process includes: Step 30, for the real defect image of the training sample, randomly sampling time step t, adding noise to obtain a noisy latent variable z t = α t x + σ t ; wherein the training sample includes a real defect image, a label containing a special identifier, and a corresponding spatial condition. α t and σ t are diffusion scheduling parameters, respectively, is the sampling noise.

[0088] When training the spatial-condition control neural network, a fine-tuning dataset can be directly used as training samples, which already contains defect images and their corresponding defect types (which can be converted into spatial conditions), and the position and size information may need to be additionally labeled or automatically obtained through defect segmentation masks. Specifically, for example, the bounding box or minimum circumscribed rectangle of the defect area is calculated, and then a binary mask map is generated to represent the position and size; or additional training samples can be collected, including real defect images, labels containing special identifiers, and their corresponding spatial conditions, including defect types, positions, and sizes. The annotation of spatial conditions can use the aforementioned method, which is not described here.

[0089] Step 31, input the spatial condition c spatial to the spatial-condition control neural network to generate a hierarchical control signal; and input the text condition c to the CLIP text encoder in the frozen stable diffusion model to obtain a text embedding. The text condition c is a label containing a special identifier.

[0090] The weights of the zero convolution layer in the spatial-condition control neural network are initialized to 0 to control the output of the control signal to be 0 at the beginning of training.

[0091] Step 32, the hierarchical control signal is injected into the corresponding level feature map of the UNet of the frozen stable diffusion model; the UNet of the frozen stable diffusion model predicts the noise based on the noisy latent variable z t , the time step t, the text embedding, and the control signal.

[0092] Step 33, calculate the mean square error (MSE) between the predicted noise and the real noise as the loss function.

[0093] Step 34, update the parameters of the spatial-condition control network through backpropagation (the parameters of the frozen diffusion model are not updated).

[0094] At this stage, the training target becomes

[0095] where the noisy latent variable z t = α t x + σ t , and c is the text prompt condition, which is a specific task condition. θ represents a UNet noise prediction function fused with the control signal with the parameter θ, a spatial condition map c spatial In this stage, the c ="image with [category] defect", where [category] is the name of the defect category that the generated image is desired to have.

[0096] The spatial condition control neural network used in step S300 ensures that the generated image meets certain conditions, which belongs to the conditional generation network. The spatial condition control neural network is not directly used to generate images, but is used in combination with the fine-tuned generation model. It provides condition information according to the input condition information such as defect type, size and position, helps the generation model to "guide" the features of the image in the generation process, and generates an image that meets the control requirements. The introduction of this condition control makes the generated image not only maintain the high-quality features of the fine-tuned diffusion model, but also accurately reflect the requirements of the defects, enhancing the reality and diversity of the generated samples. The training of the spatial condition control neural network uses the fine-tuned stable diffusion model as a frozen backbone network, calculates the mean square error loss of the predicted noise and the real noise, and only updates the parameters of the spatial condition control neural network to achieve it.

[0097] Step S400, calculate the distribution difference value of the industrial defect generated sample and the original small sample data set, select the images with a distribution difference value less than a preset threshold value and add them to the original small sample data set to form an expanded training data set.

[0098] It should be noted that the embodiments of the present disclosure can calculate the Mahalanobis distance of each sample and the small sample data set to obtain the distribution difference value, or group the industrial defect generated samples and calculate the distribution difference value of each group and the original small sample data set. The calculation of the distribution difference value includes: Step S410, randomly divide the industrial defect generated samples into multiple groups, each group containing K samples; Step S420, use a pre-trained convolutional neural network to extract feature vectors of the original small sample data set and the industrial defect generated samples, and perform L2 normalization processing.

[0099] Step S430, calculate the Fréchet distance between each sample group and the feature distribution of the original small sample data set as the distribution difference value.

[0100] Finally, as shown in Figure 2 , select the images in the sample group with a distribution difference value less than a preset threshold value and add them to the original small sample data set to form an expanded training data set.

[0101] The distribution difference value of the industrial defect generated sample in the calculation step S400 and the original small sample data set is calculated, an image with a smaller distribution difference value is selected and added to the original small sample data set to form an expanded training data set, and the detection algorithm is trained using the expanded data set. The detection algorithm is yoloV5, as shown in Figure 2 to improve the generalization ability and performance of the model.

[0102] As shown in Figure 5 Another aspect of the embodiment of the present disclosure provides an industrial defect sample generation system based on small sample fine-tuning and controllable diffusion model, which comprises: The selection module 100 is configured to select a model fine-tuning data set from the surface defect image small sample data set by using a pre-constructed image distribution evaluation module; The fine-tuning module 200 is configured to fine-tune the pre-trained stable diffusion model according to the model fine-tuning data set, and a class priori loss function is used to jointly optimize the image reconstruction loss and the pre-training knowledge preservation loss during the fine-tuning process. The defect features are bound to the generation model weight through a special identifier. The fine-tuning data set comprises defect images and corresponding labels, and the labels comprise a preset identifier and a defect type description. The generation module 300 is configured to input the prompt word containing the preset identifier and the spatial condition of the specified defect type, size and position into the fine-tuned stable diffusion model and the trained spatial condition control neural network, generate a hierarchical control signal through the spatial condition control neural network, adjust the intermediate feature representation of the stable diffusion model in the denoising process, and generate an industrial defect sample image conforming to the specified defect features. The expansion module 400 is configured to calculate the distribution difference value of the industrial defect generated sample and the original small sample data set, select an image with a distribution difference value less than a preset threshold value, and add the image to the original small sample data set to form an expanded training data set.

[0103] The system of the embodiment of the present disclosure further comprises an acquisition module configured to acquire a small sample data set comprising surface defect images. The small sample data set can be an industrial product surface defect small sample data set. A small sample refers to a sample set with a sample capacity less than N. N can be 1500 or 2000, etc., and N is less than 10000.

[0104] Another aspect of the embodiment of the present disclosure provides a computer readable storage medium having a computer program stored thereon, wherein the computer program is executed by a processor to implement the steps of the method described above.

[0105] Another aspect of the embodiment of the present disclosure provides a computer program product comprising a computer program, wherein the computer program is executed by a processor to implement the steps of the method described above.

[0106] Example one, take the class 10 defect data set in the DAGM data set as an example, verify the effectiveness of the method of the embodiment of the present disclosure. The data set is divided into training set and test set, wherein the training set and the test set each contain 1150 pictures, of which 150 are pictures containing defects.

[0107] The generated pictures of the DAGM data set are generated using the personalized industrial defect model weight obtained in the embodiment of the present disclosure, and are mixed with the original data set, and finally tested using yoloV5. In order to test the effect of the data set, the original data set is first randomly reduced, and then the generated samples are added to obtain the expanded training data set. Reduce 15, 30 original NG images, wherein the data set of reducing 15 original NG images is or_15, the data set of reducing 15 original NG images and adding 15 generated defect images is add_15, and so on.

[0108] In order to exclude the interference of random factors, the experiment is repeated 5 times, and the parameter settings in the training of yolo are: epoch is 3000; image size is 512; initial weight is yolov5m pre-training weight. Finally, the average value of the test set results of the best model is taken as the basis, as shown in Table 1. From Table 1, it can be concluded that using the expanded training data set can effectively improve the performance of the model in precision, recall rate, mAP50 and mAP50-95, which is beneficial to the improvement of the performance of the target detection algorithm Table 1 is the comparison of YOLOv5 test set indicators

[0109] The above only describes the preferred embodiments of the present disclosure, and it should be pointed out that for ordinary skilled persons in the art, without departing from the principles of the present disclosure, a number of improvements and refinements can be made, which should be considered as the protection scope of the present disclosure.< / gjjm> < / defect> < / gjjm>

Claims

1. A method for generating industrial defect samples based on small-sample fine-tuning and a controllable diffusion model, characterized in that, The method includes: Based on a small sample dataset of surface defect images, a model fine-tuning dataset is selected using a pre-built image distribution evaluation module. The pre-trained stable diffusion model is fine-tuned based on the model fine-tuning dataset. During the fine-tuning process, a priori loss function is used to jointly optimize the image reconstruction loss and the pre-training knowledge retention loss. Defect features are bound to the weights of the generation model through special identifiers. The fine-tuning dataset includes: defect images and corresponding labels, and the labels include: preset identifiers and defect type descriptions. Based on the fine-tuned stable diffusion model and the trained spatial condition control neural network, the input includes prompts containing preset identifiers and spatial conditions specifying the defect type, size, and location. The spatial condition control neural network generates hierarchical control signals to adjust the intermediate feature representation of the stable diffusion model during the denoising process, thereby generating industrial defect sample images that conform to the specified defect characteristics.

2. The method according to claim 1, characterized in that, After generating industrial defect sample images that conform to the specified defect characteristics, the method further includes: Calculate the distribution difference between the industrial defect generated samples and the original small sample dataset, and select images with a distribution difference value less than a preset threshold to add to the original small sample dataset to form an expanded training dataset.

3. The method according to claim 1, characterized in that, The image distribution evaluation module includes: a convolutional neural network and the Fréchet function; The step of selecting the model fine-tuning dataset based on a small sample dataset of surface defect images using a pre-built image distribution evaluation module includes: A convolutional neural network (CNN) is used to extract feature vectors from each image in a small sample dataset. The eigenvectors are subjected to L2 normalization to obtain normalized eigenvectors. A preset number of samples are randomly sampled from the normalized feature vectors to form multiple groups. The Fréchet distance between the feature distribution of each group and the overall distribution of the small sample dataset is calculated to obtain the distribution difference value of each group. The group data with a distribution difference value less than the preset difference is selected as the model fine-tuning dataset.

4. The method according to any one of claims 1 to 3, characterized in that, The pre-trained stable diffusion model is fine-tuned based on the model fine-tuning dataset. During the fine-tuning process, a prior-like loss function is used to jointly optimize the image reconstruction loss and the pre-training knowledge retention loss. Binding defect features to generative model weights using special identifiers includes: Load the pre-trained stable diffusion model and the model fine-tuning dataset, and freeze the weights of the CLIP text encoder and VAE; the model fine-tuning dataset includes: defect images and their corresponding labels, and the labels include at least: a preset identifier and a defect type description; Noise is added through a diffusion process, the model is trained to predict the noise, the image reconstruction loss and class prior loss are jointly optimized, and the weights of some UNet layers are updated through backpropagation. By saving the weights, we obtain the fine-tuned stable diffusion model.

5. The method according to claim 4, characterized in that, The formulas for the joint optimization of image reconstruction loss and class prior loss are as follows: In the formula: α t , σ t ,and w t They are respectively Items that control the noise propagation timeline and sample quality at all times. α t x+σ t To combine real images with noise according to the diffusion timeline, c For text conditions, λ Scalar weights; [] represents the total loss in the real image. x Text conditions c Noise sampling , and time step t Find the expected value on the distribution; This represents the reconstruction loss of the user-input image; It is a type of prior loss; x An image that contains the target concept, namely, a real industrial defect; x pr For class prior images; c pr Category prior text; t and These represent time steps; 、 Let represent the sampled Gaussian noise, respectively, which follows the order of . ; α t and σ t These represent the weight parameters related to the diffusion process; and w t Representing time steps and The corresponding weighting coefficients; Represents the square of the L2 norm; θ The UNet noise prediction model is represented by parameter θ, where θ represents the set of parameters that need to be fine-tuned in the UNet noise prediction model.

6. The method according to any one of claims 1 to 5, characterized in that, The spatial condition control neural network includes: The conditional encoding module converts spatial conditions into structured tensors; The encoder consists of four levels of convolutional layers that downsample the conditional map to extract feature maps; Triple residual blocks are used to obtain feature flow from feature maps, which enhances gradient flow. Zero convolutional layers convert the feature stream output by the residual block into a unified high-dimensional tensor.

7. The method according to any one of claims 1 to 6, characterized in that, The method of generating hierarchical control signals through a spatial condition-controlled neural network to adjust the intermediate feature representation of the stable diffusion model during the denoising process includes: The spatial conditions are encoded into a spatial condition diagram, which includes: defect type, location, and size; The spatial condition map is downsampled by the encoder to obtain the feature map. The feature map is then passed through the residual module and the zero convolutional layer group to output the hierarchical control signal. Primary control signals are injected into the shallow layer of the denoising network; intermediate control signals are injected into the middle layer of the denoising network; and advanced control signals are injected into the deep layer of the denoising network.

8. The method according to claim 7, characterized in that, The primary control signal is injected into the shallow layer of the denoising network; Inject intermediate control signals into the middle layer of the denoising network; Injecting advanced control signals deep into the denoising network, including: The primary control signal is injected into the third layer of UNet, i.e., the shallow feature layer, through residual connections, and acts on the feature map. Spatial detail dimension; Intermediate control signals are injected into the 5th layer of UNet, i.e., intermediate layer features, and act on the intermediate abstract representation space; Advanced control signals are injected into UNet's 7th layer, the deep semantic feature space, to affect the global semantic representation.

9. An industrial defect sample generation system based on small-sample fine-tuning and a controllable diffusion model, characterized in that, The system includes: The selection module is used to select a model fine-tuning dataset based on a small sample dataset of surface defect images using a pre-built image distribution evaluation module. The fine-tuning module is used to fine-tune the weights of the pre-trained stable diffusion model based on the model fine-tuning dataset. During the fine-tuning process, a priori loss function is used to jointly optimize the image reconstruction loss and the pre-training knowledge retention loss. Defect features are bound to the weights of the generation model through special identifiers. The fine-tuning dataset includes: defect images and corresponding labels, and the labels include: preset identifiers and defect type descriptions. The generation module is used to generate hierarchical control signals based on the fine-tuned stable diffusion model and the trained spatial condition control neural network. The input includes prompt words containing preset identifiers and spatial conditions specifying the defect type, size, and location. The spatial condition control neural network generates hierarchical control signals to adjust the intermediate feature representation of the stable diffusion model in the denoising process and generate industrial defect sample images that conform to the specified defect characteristics.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that, When executed by a processor, the computer program implements the steps of the method according to any one of claims 1 to 8.

Images