A Method for Generating Industrial Defect Samples Based on Small Sample Fine-tuning and Controlled Diffusion Model

Abstract

This invention proposes a method for generating industrial defect samples based on small-sample fine-tuning and a controllable diffusion model. The method uses a small-sample dataset of surface defect images. First, a pre-built image distribution evaluation module is used to select a fine-tuning dataset for the model. Then, the weights of a pre-trained stable diffusion model are fine-tuned. Next, a trained spatial condition control neural network, combined with preset identifiers as prompts, generates industrial defect samples that conform to specified defect characteristics. The distribution difference between the generated industrial defect samples and the original small-sample dataset is calculated, and images with a distribution difference value less than a preset threshold are added to the original small-sample dataset to form an expanded training dataset. This invention utilizes a small-sample dataset, fine-tunes the diffusion model weights, combines prompts containing preset identifiers, spatial conditions, and a spatial condition control neural network, expands the training samples, reduces computational costs, and has a wide range of applications.

Description

Technical Field

[0001] This disclosure relates to the field of computer vision technology, and in particular to a method, system, and computer-readable storage medium for generating industrial defect samples based on small-sample fine-tuning and a controllable diffusion model. Background Technology

[0002] In industrial production, manual inspection is the traditional method for detecting surface defects in 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. However, 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. However, insufficient data or class imbalance can lead to problems such as overfitting. Supervised deep learning-based industrial defect detection may face multiple obstacles, such as data quality, data scarcity, and difficulties in data annotation.

[0004] Industrial defect datasets commonly face problems such as small sample sizes and class imbalance. Furthermore, the type, location, and size of defects are often specific and unique. Existing image generation techniques based on stable diffusion models are generally applied to the natural and artistic fields. Research on defect image generation for the industrial domain is limited, and current stable diffusion models are difficult to control precisely. Additionally, training a large-capacity generation model from scratch requires extremely high computational costs, data requirements, and specialized knowledge. Summary of the Invention

[0005] The purpose of this disclosure is to provide a method, system, and computer-readable storage medium for generating industrial defect samples based on small-sample fine-tuning and a controllable diffusion model, thereby solving the aforementioned problems in the prior art.

[0006] To achieve the above objectives, the technical solutions adopted in the embodiments of this disclosure are as follows:

[0007] This disclosure proposes, in one aspect, a method for generating industrial defect samples based on a small-sample fine-tuning and controllable diffusion model, the method comprising:

[0008] 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.

[0009] 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.

[0010] 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.

[0011] Optionally, after generating industrial defect sample images that conform to the specified defect characteristics, the method further includes: calculating the distribution difference value between the generated industrial defect samples and the original small sample dataset, selecting images with a distribution difference value less than a preset threshold and adding them to the original small sample dataset to form an expanded training dataset.

[0012] Optionally, the image distribution evaluation module includes: a convolutional neural network and a Fréchet function;

[0013] 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:

[0014] A convolutional neural network (CNN) is used to extract feature vectors from each image in a small sample dataset.

[0015] The eigenvectors are subjected to L2 normalization to obtain normalized eigenvectors.

[0016] 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.

[0017] Optionally, the step of fine-tuning the weights of the pre-trained stable diffusion model based on the model fine-tuning dataset, wherein a prior-like loss function is used to jointly optimize the image reconstruction loss and the pre-training knowledge preservation loss during the fine-tuning process; binding defect features to the weights of the generated model through special identifiers includes:

[0018] 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;

[0019] 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.

[0020] By saving the weights, we obtain the fine-tuned stable diffusion model.

[0021] Optionally, the formulas for the joint optimization of image reconstruction loss and class prior loss are as follows:

[0022]

[0023] In the formula: α t , s t ,and w t They are respectively Items that control the noise propagation timeline and sample quality at all times. α t x+s t To combine real images with noise according to the diffusion timeline, c For text conditions, l 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 s 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.

[0024] Optionally, the spatial condition control neural network includes:

[0025] The conditional encoding module converts spatial conditions into structured tensors;

[0026] The encoder consists of four levels of convolutional layers that downsample the conditional map to extract feature maps;

[0027] Triple residual blocks are used to obtain feature flow from feature maps, which enhances gradient flow.

[0028] Zero convolutional layers convert the feature stream output by the residual block into a unified high-dimensional tensor.

[0029] Optionally, the step of generating hierarchical control signals through a spatial condition-controlled neural network to adjust the intermediate feature representations of the stable diffusion model during the denoising process includes:

[0030] The spatial conditions are encoded into a spatial condition diagram, which includes: defect type, location, and size;

[0031] 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.

[0032] 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.

[0033] Optionally, injecting primary control signals into the shallow layer of the denoising network; injecting intermediate control signals into the middle layer of the denoising network; and injecting high-level control signals into the deep layer of the denoising network includes:

[0034] 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.

[0035] Spatial detail dimension;

[0036] Intermediate control signals are injected into the 5th layer of UNet, i.e., intermediate layer features, and act on the intermediate abstract representation space;

[0037] Advanced control signals are injected into UNet's 7th layer, the deep semantic feature space, to affect the global semantic representation.

[0038] This disclosure proposes, in one aspect, an industrial defect sample generation system based on a small-sample fine-tuning and controllable diffusion model, characterized in that the system comprises:

[0039] 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.

[0040] 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.

[0041] 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.

[0042] Another aspect of this disclosure provides a computer-readable storage medium having a computer program stored thereon that, when executed by a processor, implements the steps of the method described above.

[0043] Another aspect of this disclosure provides a computer program product, including a computer program that, when executed by a processor, implements the steps of the method described above.

[0044] The beneficial effects of the embodiments disclosed herein are:

[0045] The method of this disclosure utilizes small sample data to fine-tune the weights of the diffusion model, and combines prompt words containing preset identifiers, spatial conditions, and spatial condition control neural networks to expand the training samples, reduce computational costs, and has a wide range of applications; the prior loss function supervises the consistency between the fine-tuned model and the pre-trained model's generated results. Attached Figure Description

[0046] Figure 1 This is a schematic diagram of the process for generating industrial defect samples based on a small-sample fine-tuning and controllable diffusion model, as proposed in an embodiment of this disclosure.

[0047] Figure 2 This is a schematic diagram of the entire process of an industrial defect sample generation method based on small sample fine-tuning and a controllable diffusion model proposed in this embodiment of the disclosure;

[0048] Figure 3 This is a schematic diagram illustrating the role of the class loss function in an industrial defect sample generation method based on small-sample fine-tuning and a controllable diffusion model proposed in this embodiment of the disclosure.

[0049] Figure 4 This is a flowchart illustrating the spatial control neural network model in an industrial defect sample generation method based on small-sample fine-tuning and a controllable diffusion model proposed in this embodiment of the present disclosure.

[0050] Figure 5 This is a schematic diagram of an industrial defect sample generation system based on a small-sample fine-tuning and controllable diffusion model proposed in an embodiment of this disclosure. Detailed Implementation

[0051] To make the objectives, technical solutions, and advantages of the embodiments of this disclosure clearer, the embodiments of this disclosure will be further described in detail below with reference to the accompanying drawings. It should be understood that the specific embodiments described herein are merely illustrative of the embodiments of this disclosure and are not intended to limit the embodiments of this disclosure.

[0052] This disclosure proposes a scheme for controllable sample generation via diffusion when faced with insufficient industrial product defect data and class imbalance. This disclosure relates to an industrial defect sample generation technology based on small-sample fine-tuning and a controllable diffusion model, belonging to the field of artificial intelligence-generated content, and particularly applied to the training process of industrial defect detection model algorithms.

[0053] like Figure 1 As shown, this disclosure proposes a method for generating industrial defect samples based on a small-sample fine-tuning and controllable diffusion model, the method comprising:

[0054] Step S100: Based on the pre-acquired small sample dataset of surface defect images, select the model fine-tuning dataset using the pre-built image distribution evaluation module.

[0055] Before step S100, the method of this embodiment further includes: acquiring a small sample dataset with surface defect images, which may be a small sample dataset of surface defects in industrial products. A small sample refers to a sample set with a sample size less than N. N may be 1500 or 2000, etc., and N is less than 10000.

[0056] The purpose of step S100 in this embodiment is to evaluate the original dataset of industrial defect small samples using the image distribution evaluation module based on the selected fine-tuning training samples, which is a small sample dataset of images with surface defects, to obtain the distribution difference value of each image, that is, the degree of deviation of each image from the overall distribution. The image samples with smaller distribution difference values ​​(i.e., the samples that are most consistent with the global distribution) are selected as the model fine-tuning dataset.

[0057] For example, the image distribution evaluation module includes a convolutional neural network and the Fréchet function.

[0058] The dataset selected for model fine-tuning using the pre-built image distribution evaluation module includes:

[0059] Step S110: Use a convolutional neural network (CNN) to extract feature vectors from each image in the small sample dataset.

[0060] 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 here.

[0061] Convolutional neural networks extract image features: v m = f cnn ( P m );

[0062] 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.

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

[0064] The normalization formula is as follows:

[0065]

[0066] Where v is the original feature vector, v i The eigenvector of the eigenvector i Each dimension value; v normThese 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.

[0067] 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.

[0068] Normalization results in standardized feature data, which aims to eliminate differences in the original features in terms of dimensions, range, or distribution, making different features comparable and thus improving the stability and effectiveness of subsequent calculations (such as distance metrics and model training).

[0069] For example, if the original eigenvector is [3,4], its L2 norm is... After normalization, it becomes [0.6, 0.8].

[0070] Step S130: Randomly sample a preset number of normalized feature vectors to obtain multiple groups, calculate the Fréchet distance between the feature distribution of each group and the overall distribution of the small sample dataset, thereby obtaining the distribution difference value of each group of data, and select the group data with the distribution difference value less than the preset difference as the model fine-tuning dataset.

[0071] For the small sample dataset of industrial defects, features are extracted using a convolutional neural network (CNN). After normalization, the dataset is randomly grouped, and the mean and covariance matrix of the feature vectors in each group are calculated. Then, the Fréchet distance between each group of feature vectors and the feature distribution of the small sample dataset is calculated to obtain the distribution difference value. The images of the group with the smaller distribution difference value are selected as the model fine-tuning dataset.

[0072] The convolutional neural network (CNN) used in step S100 is used to extract features. Specifically, convolutional layer modules are used for feature extraction. Since deep convolutional neural networks learn general features at the low levels, such as edges, textures, and corners, this model can extract useful low-level features.

[0073] The FID evaluation process first involves randomly sampling the feature vectors of the dataset into multiple groups, for example, 10-20 images per group. For each group, the Fréchet distance between its overall feature distribution (mean and covariance matrix) and the small sample original dataset is calculated. The group of images with higher FID scores is then selected as the model fine-tuning dataset.

[0074] Then, the extracted feature vectors are subjected to Fréchet distance calculation, which measures the distance between image distributions. This provides a quantitative indicator for evaluating the quality of the generative model; the smaller the value, the closer the generated samples are to the real samples, and the more the training samples conform to the overall sample distribution. The Fréchet distance distribution difference value used is... d FID ( p 1, p 2) The specific formula is expressed as follows:

[0075]

[0076] p 1, p 2 represents the Gaussian distribution of the randomly sampled group and the pre-acquired small sample dataset of surface defect images, respectively; The standard deviation of the data segment; Let ∑1 and ∑2 be the feature covariance matrices of the randomly sampled group and small sample datasets, respectively. These are the mean vectors of the feature vectors of the randomly sampled group and small sample datasets, respectively, which are the centers of the feature space; Represents the trace of a matrix.

[0077] Select d FID ( p 1, p 2) Group image data with differences less than a preset threshold are used as the model fine-tuning dataset. Group images with differences less than the threshold are selected and added to the original dataset to form an expanded dataset (step S500).

[0078] Step S200: 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. The labels include: preset identifiers and defect type descriptions.

[0079] like Figure 3 As shown, the specific process includes: freezing the weights of the CLIP text encoder and VAE decoder of the pre-trained stable diffusion model; jointly optimizing the following two items using a priori loss function: image reconstruction loss: constraining the model to learn defective image features; pre-training knowledge retention loss: constraining the output distribution to be consistent with the original model through priori image constraints; and binding defective features to UNet weights through special identifiers embedded in the text labels.

[0080] In this embodiment of the present disclosure, step S200 uses the model fine-tuning dataset from step S100 to fine-tune the weights of the existing pre-trained stable diffusion model. During fine-tuning, a priori loss function is used to preserve prior knowledge, and defect features are bound to special identifiers to generate a personalized industrial defect model. During image generation, prompts containing special identifiers enable the generation of images with specific defect features. These special identifiers are preset identifiers. The special identifiers are set before weight fine-tuning. Defect features are then bound to the special identifiers.

[0081] The fine-tuning dataset includes: defect images and corresponding labels. The labels include special identifiers and defect type descriptions. The special identifiers can be set according to the actual situation.

[0082] The fine-tuning dataset provides industrial defect samples, such as defect images and corresponding labels, for learning new features. It is used to calculate image reconstruction loss, ensuring the model learns to generate defects. The fine-tuning process is constrained by a class-prior loss function. A generic image, independent of the industrial domain, is generated using the pre-trained model with frozen weights as a distribution consistency benchmark, preventing the model from forgetting its original capabilities. The class-prior loss is calculated between the generic image and the results generated by the fine-tuned model, constraining the feature distributions of both to be consistent. The pre-trained model is a pre-trained stable diffusion model. By jointly optimizing the losses of these two types of inputs, the model adapts to industrial defect features while retaining pre-training knowledge. The special identifier is a unique text tag that is strongly associated with the defect category during fine-tuning. During the generation phase, inputting a prompt containing this identifier triggers the generation of a specific defect.

[0083] Pre-trained stable diffusion models (such as Stable Diffusion) are typically trained on large-scale general datasets (such as LAION). Fine-tuning refers to further adjusting the model weights using domain-specific data (such as images of industrial defects) to adapt it to new tasks. The pre-trained stable diffusion model can be considered as the initial stable diffusion model.

[0084] During training, defect features are bound to special identifiers to uniquely bind new defect concepts, avoiding contamination of pre-trained semantics. In the generation phase, inputting prompts containing special identifiers and defect type descriptions triggers the generation of specific defects.

[0085] For example, the step of fine-tuning the weights of the pre-trained stable diffusion model based on the model fine-tuning dataset involves jointly optimizing the image reconstruction loss and the pre-training knowledge preservation loss using a priori loss function during the fine-tuning process; and binding defect features to the generated model weights through special identifiers, including:

[0086] Step S210: Load the pre-trained stable diffusion model and the model fine-tuning dataset, and freeze the weights of CLIP text encoder and VAE; wherein, the model fine-tuning dataset includes: defect images and their corresponding labels, the labels include at least: preset identifiers and defect type descriptions, so that the preset identifiers are strongly associated with the defect categories.

[0087] The model fine-tuning dataset used in step S100 includes: defect images and their corresponding labels. In the labels of the model fine-tuning dataset, defect categories are associated with specific identifiers. Example: The label is "The product surface has..." <gjjm>"Cracks" refers to surface cracks on the product as a defect type description, with GJJM as a special identifier. Each defect type corresponds to a special identifier. The data needs to be preprocessed to a model input format, such as 512x512 resolution. Load the weights of an open-source pre-trained StableDiffusion model, such as stable-diffusion-v1-5, including its UNet, CLIP text encoder, and VAE parts. Fine-tune the weights of the U-Net denoising network.

[0088] Special identifiers (such as) <defect-type>`` is a placeholder used to trigger the model's ability to generate specific defect features during the generation phase. Its function is similar to the `[V]` identifier in DreamBooth, binding a new concept (scratch) to a unique identifier. Special identifiers (such as...) <defect>() is a key design element in fine-tuning, used to bind new concepts without interfering with existing knowledge.

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

[0090] 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.

[0091] 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.

[0092] 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... l· Prior loss function L pr .

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

[0094]

[0095] in, α t , s t ,and w t They are respectively Items that control the noise propagation timeline and sample quality at all times. α t x+s t To combine real images with noise according to the diffusion timeline, c For text conditions, λ·L pr , L pr For prior preservation terms, l 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 Find the expected value (average) of the distribution. The expression enclosed in brackets [] is a weighted MSE loss function for the user image and the class prior image, used to stabilize model training and prevent overfitting to the user image. The first term... This represents the reconstruction loss of the user input image; it calculates the model's response to real noise under condition c. Prediction error: By minimizing this error, the model learns how to reconstruct the target concept x from noise. The MSE (Mean Sequence Size) of the predicted noise versus the real noise is used to force the model to accurately predict the noise added to the real image, thus enabling the model to learn to reconstruct the unique features of the target image (user image) from the noise. (Second term) λ·L pr ,Right now It is a prior loss that constrains the output to be consistent with the distribution of the pre-trained model, and constrains the model not to disrupt the original class distribution when learning user concepts. x This represents an image that contains the target concept, namely, a real industrial defect; x pr Prior images (images representing general categories generated by pre-trained models) are typically generated by the original, untuned pre-trained model based on prior text representing categories. c pr The generated image; c Conditional text, used to describe user-provided concepts, such as "a photo of [V] dog"; c pr Category-prior text, typically general category text, such as "a photo of a dog"; t and These represent time steps (time step indices used in the diffusion process); and Let represent the sampled Gaussian noise, respectively, which follows the order of . ; α t and s t The weighting parameters, which are related to the diffusion process, determine the ratio of signal to noise at the current time step. and w t Representing time steps t´ and t The corresponding weighting coefficients (often related to the noise dispatcher); This represents the squared L2 norm, also known as the mean square error (MSE). θ The UNet noise prediction model has parameter θ, where θ represents the set of parameters that need to be fine-tuned in the UNet noise prediction model. θ The input is a noisy latent variable. α t x+s t Given time step t and condition c, the output is the predicted noise tensor. Before fine-tuning, training data needs to be prepared for the class prior loss. Using a pre-trained diffusion model with untuned, original weights and a sampler, training data is generated based on the class prior text. c pr Generate a batch of prior image data x pr = (z t1 , c pr The model contains random initial noise. And the condition vector. Where the second term... These are prior protection terms used to supervise the model using self-generated images. l Control the relative weight of this item. This represents the generative model, responsible for mapping latent variables to pixel space. t1 The latent variable represents time step t.

[0096] The purpose of the prior loss function is to address the catastrophic forgetting problem in diffusion models when learning new concepts, ensuring that the model learns the concept of industrial defects without forgetting existing knowledge. With limited data, it learns to generate defect features while retaining its original generative capabilities.

[0097] The mechanism of class-prior loss: By constraining the fine-tuning model, the noise prediction results of the class-prior image are consistent with the pre-trained model: Feature matching: Constraining the feature distribution of the fine-tuned model to be close to that of the pre-trained model (e.g., using KL divergence to measure the difference in features of intermediate layers of UNet). Output consistency: For the same input (text + noise), the model output before and after fine-tuning should remain similar (e.g., calculating the LPIPS distance of the generated image).

[0098] Before fine-tuning the pre-trained stable diffusion model in step S220, the core components of the pre-trained model are frozen, such as the CLIP text encoder of Stable Diffusion and the weights of the VAE. During the fine-tuning process, the weights of CLIP and VAE are frozen first.

[0099] Step S230: Save the weights to obtain the fine-tuned stable diffusion model.

[0100] The fine-tuned UNet or LoRA weights can be saved separately to obtain a personalized industrial defect model.

[0101] High-quality samples are generated using pre-trained models to assist in the fine-tuning process (similar to self-distillation).

[0102] Freezing the model: The CLIP, VAE, and other parts of the pre-trained model are frozen throughout the training process; only some layers of the UNet are trainable. Pre-sampling: In each training round, the current model (under fine-tuning) and the frozen pre-trained model are used for pre-sampling.

[0103] Do not generate images. Loss calculation: Task loss: supervises the difference between real defect data and model predictions. Prior loss: supervises the consistency between the results generated by the fine-tuned model and the pre-trained model. Parameter update: only update the weights of the unfrozen parts (such as the LoRA parameters in UNet).

[0104] The class prior loss used in step S200 constrains the generated image from deviating from common class attributes by comparing the difference between the feature distribution of the generated image and the distribution of the target class. This loss is often used in conjunction with other generation losses (such as image reconstruction loss, adversarial loss, etc.). The role of class prior loss includes: 1) helping the generative model maintain prior knowledge of a specific class during training, especially during fine-tuning. It avoids the model relying solely on features from a small number of samples, ensuring that the generated images have stronger class consistency and generality; 2) class prior loss encourages the model to generate images that conform to the general characteristics of the class, rather than overfitting based solely on a small number of input samples. This means that even with a very small number of input training samples, the model can still understand the overall distribution of the class, thereby generating diverse images that conform to the characteristics of the class.

[0105] Step S300: Based on the fine-tuned stable diffusion model and the trained spatial condition control neural network, input prompt words containing preset identifiers and spatial conditions specifying defect type, size and location. Generate hierarchical control signals through the spatial condition control neural network, 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. The spatial condition control neural network includes: a conditional encoding module, an encoder, a triple residual block and a zero convolutional layer group.

[0106] By using the fine-tuned stable diffusion model combined with the weights of the control network, and using the special identifier from step S200 as a prompt word, an industrial defect sample that conforms to the specified defect characteristics is generated. The special identifier is the preset identifier.

[0107] Based on a trained spatial conditional control neural network, the intermediate feature representation of a stable diffusion model during the diffusion process can be adjusted through hierarchical control signals according to input geometric conditions, such as defect type, location, and size, thus influencing the image generated by a large neural network, which refers to the stable diffusion model. In this embodiment, the network to be controlled is a stable diffusion network that has learned new concepts. Throughout the process, the model weights are first frozen to preserve the model's initial weights during training. Subsequently, the spatial conditional control neural network, the adapter network of ControlNet, is adjusted. It is a lightweight, trainable branch attached to the frozen master model and used as a conditional adapter to regulate the frozen copy. The spatial conditional control neural network is a structural copy of the master model UNet, but with independent and trainable weights. The structure of ControlNet replicates a portion of the layers of the master model (such as UNet), and these replicated layers are then connected to the corresponding layers of the master model, with control signals injected through a zero-convolution mechanism. The input conditions are an image displaying the rough geometry of the defect content, as well as the defect type and shape.

[0108] It should be noted that the spatial condition control neural network is trained independently. The spatial condition control neural network is trained to adjust the stepwise denoising process of the generated image according to the input conditions of the specified defect type, size and location. That is, to adjust the intermediate feature representation in the diffusion process, thereby controlling the stepwise denoising process of the image so that the generated image meets the predetermined conditions, thereby generating industrial defect samples that meet the specified defect characteristics.

[0109] Pixel-level spatial constraints are achieved by generating a trainable control path that runs parallel to a pre-trained diffusion model through external conditional input. Essentially, this involves constructing a conditional adapter network that injects precise control signals without compromising the original model's generative capabilities.

[0110] Step S200 outputs a finely tuned stable diffusion model, which serves only as the base model for generating defect images.

[0111] Spatial Condition Control Neural Network (SCNN) is a standalone lightweight network (such as ControlNet or a similar structure) responsible for dynamically adjusting the generation process of a stable diffusion model based on user-inputted conditions such as defect type, size, and location.

[0112] like Figure 4 As shown, the implementation of conditional control in the spatial conditional control neural network involves generating hierarchical control signals to adjust the intermediate feature representations of the stable diffusion model during the diffusion process. This includes:

[0113] Step S310: Encode the spatial conditions into a spatial condition diagram, wherein the spatial condition diagram includes: defect type, as well as location and size.

[0114] The conditional encoding module encodes user-specified spatial conditions (such as "scratch, 5px long, located in 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 heatmap to indicate the defect region, where the defect region is 1 and the background is 0. The defect location can be given by a bounding box or segmentation mask, and then converted into a fixed-resolution binary map.

[0115] Step S320: 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.

[0116] The spatial conditional control neural network takes the aforementioned conditional maps (mask map, category embedding) as input and first enters the encoder for downsampling. The encoder comprises four convolutional layers, which progressively compress features: the first layer extracts basic surface texture, the second layer captures contour structures, the third layer uses dilated convolution to expand the receptive field to 51×51 pixels to cover larger features, and the fourth layer compresses the feature map to 32×32 resolution and encapsulates global defect distribution information. This process achieves a physical mapping from pixel space to a high-dimensional feature space. The encoder's four convolutional layers achieve spatial resolution compression (from 512 to 32), and the receptive field expands progressively, serving as the UNet control signal for the fine-tuned stable diffusion model.

[0117] The feature maps output by the encoder are sequentially input into the residual module to obtain the feature flow, which enhances the gradient flow. The residual module includes a triple residual block.

[0118] The feature stream enters the zero-convolutional layer group, which transforms the feature stream output from the residual block into a unified high-dimensional tensor as the original control signal. The weights and biases of these 1×1 convolutions are initialized to zero. Initially, the output of this module is always a zero tensor, ensuring that it does not interfere with the knowledge system of the pre-trained UNet during the early stages of training. As training progresses, the zero-convolutional layers gradually learn to transform the control signal into a physical offset in the feature space. This tensor is then decomposed into shallow, middle, and deep control signals through a hierarchical projection module (composed of three parallel 1×1 convolutions), corresponding to layers 3, 5, and 7 of the UNet denoising network, respectively.

[0119] The shallow signal (64×64 resolution), i.e., the primary signal, is converted into a spatial attention mask. A binarized heatmap is generated using sigmoid activation, guiding the model to enhance defect texture details within the mask region (value ≈ 1). The mid-level signal (32×32 resolution), i.e., the intermediate signal, is converted into structural scaling factors. The feature map amplitude is adjusted through channel multiplication to control the defect geometry. The deep signal (16×16 resolution), i.e., the high-level signal, is converted into semantic channel weights. A channel-dimensional scaling vector is generated through a fully connected layer to adapt to the high-level semantics of the defect category.

[0120] Step S330: Inject the primary control signal into the shallow layer of the denoising network; inject the intermediate control signal into the middle layer of the denoising network; inject the high-level control signal 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 structural control signal, and the high-level control signal is also a semantic control signal.

[0121] In each step of the diffusion model's denoising process, the spatial condition control network encodes the specified spatial conditions, such as defect type, location, and size, into feature tensors and generates hierarchical control signals (detail, structural, and semantic signals). These signals are precisely injected into and fused into the shallow (layer 3), middle (layer 5), and deep (layer 7) feature maps of the UNet decoder through adaptive instance normalization, channel multiplication, and cross-attention mechanisms during UNet's forward propagation for noise prediction, respectively, to achieve pixel-level spatial control over the generated image content.

[0122] In other words, at each time step t in the above denoising process, the spatial conditional control neural network generates a set of hierarchical control signals based on the input spatial condition map. These signals are injected into the corresponding 3rd, 5th, and 7th layer feature maps before the UNet network performs the noise prediction calculation for the current step, thereby incorporating spatial control conditions into the noise prediction at that time step and gradually guiding the latent variables to evolve towards an image that conforms to the predetermined spatial features.

[0123] The adjustments to the intermediate feature representations include: injecting detail control signals into the feature map of the denoising network's third layer, implemented through adaptive instance normalization (AdaIN), transferring the style (mean and variance) of the control signals to the UNet features to modulate edge / texture details. Injecting structural control signals into the feature map of the denoising network's fifth layer, implemented through channel multiplication, controls the structural geometry; specifically, channel-wise scaling, where the control signal is a vector with the same number of channels as the feature map, multiplied channel-wise to enhance or suppress the activation intensity of certain feature channels. Injecting semantic control signals into the feature map of the denoising network's seventh layer, implemented through a cross-attention mechanism to maintain semantic consistency; specifically, using the control signals as the key and value, and the UNet features as the query, attention calculation is performed, allowing semantic information to globally modulate the feature map.

[0124] Specifically, the control signal is injected into the stable diffusion model in the following way: a hierarchical injection mechanism for the control signal:

[0125] UNet Layer 3 (64×64 resolution): The primary control signal output from the zero convolutional layer is passed through residual connections.

[0126] This shallow feature injection layer acts on the spatial detail dimension of the feature map. This layer of injection mainly modulates the local high-frequency features of the generated image, directly affecting the fine structure of edge textures.

[0127] UNet Layer 5 (32×32 resolution): The intermediate control signal output from the zero convolutional layer is injected into the network's intermediate layer features, acting on...

[0128] Intermediate-level abstract representation space. This level is responsible for coordinating the combination relationships of local features and constraining the geometry of the structure within a moderate receptive field.

[0129] UNet Layer 7 (16×16 resolution): High-level control signals from zero-convolutional layer outputs are injected into deep semantics.

[0130] The feature space operates on the global semantic representation. This level regulates the overall semantic attributes, ensuring that the generated content remains consistent with higher-level concepts.

[0131] The training of the spatial condition control neural network is carried out based on a fine-tuned stable diffusion model. The training process of the spatial condition control neural network includes:

[0132] Step 30: For real defect images of the training samples, randomly sample at time step t, add noise, and obtain noisy latent variables. z t = α t x + s t The training samples include real defect images, labels with special identifiers, and their corresponding spatial conditions. α t and s t These are the diffusion scheduling parameters. It is sampling noise.

[0133] When training a spatial condition control neural network, a fine-tuning dataset can be used directly as training samples. The dataset already contains defect images and their corresponding defect types (which can be converted into spatial conditions). The location and size information may need to be additionally labeled, or the location and size can be automatically obtained through defect segmentation masks. Specifically, for example, the bounding box or minimum bounding rectangle of the defect region can be calculated, and then a binary mask image can be generated to represent the location and size. Alternatively, additional training samples can be collected. The training samples include: real defect images, labels with special identifiers and their corresponding spatial conditions. The spatial conditions include: defect type, location and size. The labeling of spatial conditions can be done using the methods mentioned above, which will not be elaborated here.

[0134] Step 31, set spatial condition c spatial The input spatial conditional control neural network generates hierarchical control signals; the text condition c is input into the CLIP text encoder in the frozen stable diffusion model to obtain the text embedding. The text condition c is a label containing a special identifier.

[0135] The weights of the zero convolutional layer in the spatial condition control neural network are initialized to 0 so that the control signal output is 0 in the early stage of training.

[0136] Step 32: The hierarchical control signal is injected into the feature map of the corresponding layer in the UNet of the frozen stable diffusion model; the UNet of the frozen stable diffusion model is based on the noisy latent variables. z t The time step t, text embedding, and control signals are used to predict noise.

[0137] Step 33: Calculate the mean square error (MSE) between the predicted noise and the actual noise as the loss function.

[0138] Step 34: Update the parameters of the spatial condition control network through backpropagation (the parameters of the frozen diffusion model are not updated).

[0139] At this stage, the training objective becomes

[0140] Among them, noisy latent variables z t = α t x + s t ,and c These are text prompt conditions, or specific task conditions. θ This represents the UNet noise prediction function with parameter θ, which incorporates the control signal; spatial condition plot. c spatial At this stage, use c = "[V] image with [category] defect", where [category] is the name of the defect category that the image to be generated has.

[0141] The spatial conditional control neural network used in step S300 ensures that the generated image meets specific conditions; it belongs to the category of conditional generation networks. The spatial conditional control neural network does not directly generate images but is used in conjunction with a fine-tuned generation model. Based on input conditional information such as defect type, size, and location, it provides conditional information to help the generation model "guide" the image features during the generation process, thereby generating images that meet control requirements. This introduction of conditional control ensures that the generated images not only maintain the high-quality features of the fine-tuned diffusion model but also accurately reflect the defect requirements, enhancing the realism and diversity of the generated samples. The training of the spatial conditional control neural network uses the fine-tuned stable diffusion model as a frozen backbone network, calculates the mean squared error loss between predicted noise and real noise, and updates only the parameters of the spatial conditional control neural network.

[0142] Step S400: Calculate the distribution difference value between the industrial defect generated samples and the original small sample dataset, select images with a distribution difference value less than a preset threshold and add them to the original small sample dataset to form an expanded training dataset.

[0143] It should be noted that the embodiments of this disclosure can calculate the Mahalanobis distance between each sample and the small sample dataset to obtain the distribution difference value, and can also group industrial defect samples and calculate the distribution difference value between each group and the original small sample dataset. The calculation of the distribution difference value includes:

[0144] Step S410: Randomly divide the industrial defect generation samples into multiple groups, each group containing K samples;

[0145] Step S420: Use a pre-trained convolutional neural network to extract feature vectors from the original small sample dataset and industrial defect generation samples, and perform L2 normalization processing.

[0146] Step S430: Calculate the Fréchet distance between each sample group and the feature distribution of the original small sample dataset, as the distribution difference value.

[0147] Finally, as Figure 2 As shown, images from sample groups with distribution differences less than a preset threshold are selected and added to the original small sample dataset to form an expanded training dataset.

[0148] Calculate the distribution difference between the industrial defect generated samples in step S400 and the original small sample dataset. Select images with smaller distribution difference values ​​and add them to the original small sample dataset to form an expanded training dataset. Then, use the expanded dataset to train the detection algorithm. The detection algorithm is YOLOv5. Figure 2 As shown, this is done to improve the model's generalization ability and performance.

[0149] like Figure 5 As shown, another aspect of this disclosure proposes an industrial defect sample generation system based on a small-sample fine-tuning and controllable diffusion model, the system comprising:

[0150] Module 100 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.

[0151] The fine-tuning module 200 is used to fine-tune the weights of the pre-trained stable diffusion model according to 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.

[0152] The generation module 300 is used to generate a hierarchical control signal 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 a hierarchical control signal to adjust the intermediate feature representation of the stable diffusion model in the denoising process and generate an industrial defect sample image that conforms to the specified defect characteristics.

[0153] The expansion module 400 is used to calculate the distribution difference value between the industrial defect generation 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.

[0154] The system in this embodiment further includes: an acquisition module, used to acquire a small sample dataset containing surface defect images. The small sample dataset can be a small sample dataset of surface defects in industrial products. Small sample refers to a sample set with a sample size less than N. N can be 1500 or 2000, etc., and N is less than 10000.

[0155] Another aspect of this disclosure provides a computer-readable storage medium having a computer program stored thereon that, when executed by a processor, implements the steps of the method described above.

[0156] Another aspect of this disclosure provides a computer program product, including a computer program that, when executed by a processor, implements the steps of the method described above.

[0157] Example 1 uses the class 10 defect dataset from the DAGM dataset as a case study to verify the effectiveness of the method in this embodiment. This dataset is divided into a training set and a test set, each containing 1150 images, of which 150 are images containing defects.

[0158] The generated images for the DAGM dataset are generated using the personalized industrial defect model weights obtained in this embodiment of the disclosure, and then mixed with the original dataset. Finally, YOLOv5 is used for testing. To test the effect of this dataset, images in the original dataset are first randomly reduced, and then generated samples are added to obtain an expanded training dataset. 15 and 30 original NG images are reduced respectively, with the dataset having 15 original NG images reduced being designated as or_15, and the dataset having 15 original NG images reduced and then 15 generated defect images added being designated as add_15, and so on.

[0159] To eliminate the interference of random factors, the experiment was repeated 5 times. The YOLO training parameters were set as follows: epoch 3000; image size 512; initial weights were the YOLOv5M pre-trained weights. The final result was based on the average of the best model's performance on the test set, as shown in Table 1. Table 1 concludes that using the expanded training dataset can effectively improve the model's performance in precision, recall, mAP50, and mAP50-95, which is beneficial to improving the performance of the object detection algorithm.

[0160] Table 1 shows the comparison of YOLOv5 test set metrics.

[0161]

[0162] The above description is only a preferred embodiment of the present disclosure. It should be noted that for those skilled in the art, several improvements and modifications can be made without departing from the principle of the present disclosure, and these improvements and modifications should also be considered within 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, input prompts containing preset identifiers and spatial conditions specifying defect type, size and location are used. 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 generates industrial defect sample images that conform to the specified defect characteristics. 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. 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. Save the weights to obtain the fine-tuned stable diffusion model; 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.

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 any one of claims 2, 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.

4. The method according to any one of claims 3, 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.

5. The method according to claim 4, 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.

6. 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. 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. 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. Save the weights to obtain the fine-tuned stable diffusion model; 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.

7. 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 5.

Images