A method for constructing a lightweight diffusion model based on a defect image and related devices

Abstract

The application discloses a construction method of a light-weight diffusion model based on a defect image and related devices, and the method comprises the following steps: acquiring an original defect-free image and a defect image, and fusing the original defect-free image and the defect image to obtain a fused image; performing multi-scale feature extraction on the fused image through a pre-trained network architecture to obtain scale features of multiple different levels; performing a pooling operation on the last-level scale feature in the scale features of the multiple different levels to obtain prior extraction features; based on a network search algorithm, optimizing a network structure of a preset light-weight network according to the fused image, the scale features of the multiple different levels and the prior extraction features, and outputting a primary light-weight network; performing a pruning operation on the primary light-weight network to obtain a target light-weight network; taking the target light-weight network as a backbone network to establish an initial light-weight diffusion model; acquiring a defect sample set, training the initial light-weight diffusion model through the defect sample set, and generating a target light-weight diffusion model.

Description

Technical Field

[0001] This invention relates to the field of defect detection technology, and in particular to a method and apparatus for constructing a lightweight diffusion model based on defect images. Background Technology

[0002] In industrial settings, acquiring and labeling defective workpieces is extremely difficult, posing a significant obstacle to defect detection. Generating a large number of defect samples from a small number of real defect samples greatly alleviates the problem of sample scarcity. However, traditional generative models typically require substantial computational resources and bandwidth to handle image denoising, missing image restoration, detail loss, and other defect image generation operations. Diffusion models are a powerful and novel deep learning generative model. They employ parameterized Markov chains to optimize the downvariant boundary of the likelihood function, learning the noise distribution and storing noise information through a progressive noise addition process. Subsequently, in the reverse denoising process, the model is trained to fit the previously obtained noise labels to obtain high-quality defect-generated images. Diffusion models have a rigorous mathematical derivation process and possess characteristics such as no aftereffect and homogeneity, resulting in more accurate target distributions in their generated samples compared to other generative models (such as generative adversarial networks). However, their progressive denoising process involves numerous iterative steps, often leading to significant computational resource consumption. This severely impacts the efficiency of defect image generation and computational overhead, making it difficult to meet the application requirements of lightweight and efficient models in industrial scenarios. Summary of the Invention

[0003] This invention provides a method and apparatus for constructing a lightweight diffusion model based on defect images, which solves the technical problem of low computational efficiency in generating high-quality defect images using existing diffusion models based on defect images.

[0004] This invention provides a method for constructing a lightweight diffusion model based on defect images, the method comprising:

[0005] Obtain the original defect-free image and the defective image, and then fuse the original defect-free image and the defective image to obtain the fused image;

[0006] Multi-scale feature extraction is performed on the fused image using a pre-trained network architecture to extract multiple scale features at different levels; pooling is then performed on the lowest-level scale feature among these multiple scale features to obtain prior extracted features.

[0007] Based on the network search algorithm, the network structure of the preset lightweight network is optimized according to the fused image, the scale features of multiple different levels and the prior extracted features, and the primary lightweight network is output.

[0008] The primary lightweight network is pruned to obtain the target lightweight network; an initial lightweight diffusion model is established using the target lightweight network as the backbone network.

[0009] Obtain a set of defective samples, train the initial lightweight diffusion model using the set of defective samples, and generate the target lightweight diffusion model.

[0010] Optionally, the pre-trained network architecture includes multiple residual blocks connected in sequence; the feature extraction operation of each residual block includes two convolution operations and one activation function activation operation;

[0011] The step of extracting multi-scale features from the fused image using a pre-trained network architecture to obtain multiple scale features at different levels, and then performing prior extraction on the lowest-level scale feature among these multiple scale features to obtain prior extracted features, includes:

[0012] In the pre-trained network architecture, feature extraction operations are performed on the fused image through multiple sequentially connected residual blocks, and the scale features corresponding to each residual block are output.

[0013] The final scale feature of the output of the last residual block in the pre-trained network architecture is selected, and average pooling, linear normalization and activation function integer operation are performed on the final scale feature in sequence to obtain the prior extracted feature.

[0014] Optionally, the preset lightweight network includes multiple standard modules, downsampling modules, and prior extraction modules at different levels; at the same level, the standard modules, downsampling modules, and prior extraction modules support direct connection and skip connection.

[0015] The standard module is used to maintain the feature size, and its convolution stride is set to 1; the downsampling module is used to downsample the features to obtain multidimensional information, and its convolution stride is set to 2; the prior extraction module is used to introduce the prior extracted features.

[0016] Optionally, the method further includes:

[0017] Establish an internal structure selection mechanism for the standard module and the downsampling module, and represent the internal structure of the standard module and the downsampling module as a directed acyclic graph;

[0018] Under the internal structure selection mechanism, the convolutional operations of the standard module or the downsampling module are selected and replaced by a preset internal selection set; wherein, the internal selection set includes skip connections, fully connected operations and max pooling.

[0019] Optionally, the step of optimizing the network structure of a preset lightweight network based on the network search algorithm, according to the fused image, the multiple different levels of scale features, and the prior extracted features, and outputting a primary lightweight network includes:

[0020] The fused image and the prior extracted features are input into a preset lightweight network to output multiple multi-scale fused features at different levels;

[0021] The scale loss value between the multi-scale fusion features at multiple different levels and the scale features at multiple different levels is calculated using a preset loss function.

[0022] The network search algorithm optimizes the network structure distribution of the lightweight network in the multi-scale search space based on the scale loss value to determine the optimal combination connection of the standard module, downsampling module and prior extraction module at the same scale, as well as the optimal internal selection set of the standard module and downsampling module, thereby outputting the primary lightweight network.

[0023] The calculation method for the multi-scale fusion features at different levels is as follows:

[0024]

[0025] In the formula: This represents the multi-scale fusion feature of the s-th layer of a lightweight network; This represents the output scale feature of the standard module at level s. This represents the output scale characteristics of the downsampling module at level s. This represents the output scale features of the prior extraction module; , and All of these are learnable weight parameters, used to control the contributions of the standard module, downsampling module, and prior extraction module, respectively.

[0026] The preset loss function is expressed as:

[0027]

[0028] In the formula: This is the scale loss value. To merge images; For the scale features of the s-th layer of the pre-trained network architecture, The multi-scale fusion features at the s-th layer of the pre-defined lightweight network; This represents the number of nodes in the directed acyclic graph represented by the standard module and the downsampling module.

[0029] Optionally, the step of pruning the primary lightweight network to obtain the target lightweight network includes:

[0030] Acquire a calibration image and input the calibration image into the primary lightweight network;

[0031] Traverse each convolutional kernel in the primary lightweight network and calculate the feature map corresponding to the output of each convolutional kernel; calculate the corresponding feature map based on the feature map corresponding to each convolutional kernel.

[0032] Based on a preset threshold, remove convolutional kernels whose thresholds are less than the preset threshold corresponding to the feature map, and update the primary lightweight network; determine whether the updated primary lightweight network meets a preset condition. If it does, proceed to the step of traversing each convolutional kernel in the primary lightweight network and calculating the feature map output by each convolutional kernel; if it does not meet the condition, use the primary lightweight network that previously met the preset condition as the target lightweight network.

[0033] Optionally, the process of acquiring the fused image includes:

[0034] The original defect-free image and the defective image are stitched together to obtain a stitched image; the stitched image is then subjected to a convolution downsampling operation to restore the image resolution, resulting in a fused image.

[0035] Optionally, the process of obtaining the defect sample set includes:

[0036] Based on the masking technology, a binary mask image is generated from the defect image, and the defect region in the defect image is identified and extracted based on the binary mask image.

[0037] The defective region is integrated, transformed, and enhanced using the original defect-free image to obtain a corresponding defect template.

[0038] Multiple new defect-free images are acquired, and each new defect-free image is subjected to defect fusion processing according to the defect template to obtain multiple defect samples; the defect sample set includes multiple defect samples.

[0039] The calculation method for the defect samples is expressed as follows:

[0040]

[0041] In the formula: Indicates a defective sample; Indicates a defect template. Defect template Defect mask; This indicates a new, defect-free image.

[0042] Optionally, the step of training the initial lightweight diffusion model using the defect sample set to generate the target lightweight diffusion model includes:

[0043] The defective samples and the original defect-free image from the defective sample set are input into the initial lightweight diffusion model;

[0044] Based on the defective sample and the corresponding original defect-free image, a feature prior extraction operation is performed to obtain the training prior extracted features; the guiding features are calculated based on the training prior extracted features and the corresponding original defect-free image.

[0045] The training prior extracted features are subjected to forward noise addition, and the target lightweight network is trained with the added Gaussian noise as the training label to obtain a noisy image, thereby realizing the forward noise addition training of the lightweight diffusion model.

[0046] Based on the noisy image, the guiding features are mapped to the intermediate representation layer of the target lightweight network through a cross-attention mechanism for iterative denoising prediction, thereby obtaining the predicted noisy image and realizing the reverse denoising training of the lightweight diffusion model.

[0047] The total loss value of the training is calculated by a preset overall loss function, and the model parameters of the initial lightweight diffusion model are optimized by the total loss value until the preset iteration conditions are met, thereby generating the target lightweight diffusion model.

[0048] The present invention also provides a computer system, including a memory, a processor, and a computer program stored in the memory, characterized in that the processor executes the computer program to implement the steps of the method for constructing a lightweight diffusion model based on defective images as described above.

[0049] As can be seen from the above technical solutions, the present invention has the following advantages:

[0050] This invention provides a method and related apparatus for constructing a lightweight diffusion model based on defective images. The method includes: acquiring an original defect-free image and a defective image; fusing the original defect-free image and the defective image to obtain a fused image; extracting multi-scale features from the fused image using a pre-trained network architecture to obtain multiple scale features at different levels; performing pooling operations on the lowest-level scale features among the multiple scale features at different levels to obtain prior extracted features; optimizing the network structure of a preset lightweight network based on a network search algorithm, according to the fused image, multiple scale features at different levels, and prior extracted features, to output a primary lightweight network; pruning the primary lightweight network to obtain a target lightweight network; establishing an initial lightweight diffusion model with the target lightweight network as the backbone network; acquiring a defect sample set, training the initial lightweight diffusion model using the defect sample set, and generating the target lightweight diffusion model.

[0051] This invention guides a lightweight network to optimize its architecture within the search space based on multi-scale features and prior extracted features of both original, defect-free images and defective images. Simultaneously, pruning techniques are used to further compress the structural size of the lightweight network. The optimized lightweight network serves as the backbone network of a lightweight diffusion model, effectively reducing computational resource consumption during model inference and avoiding redundant computational processes. Furthermore, by training the lightweight diffusion model with a set of defective samples, the lightweight diffusion model maintains high computational quality in generating high-quality defective images while improving computational efficiency. Attached Figure Description

[0052] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0053] Figure 1 A flowchart illustrating the steps of constructing a lightweight diffusion model based on defect images, as provided in an embodiment of the present invention;

[0054] Figure 2 This is a data processing flowchart of the pre-trained network architecture provided in an embodiment of the present invention;

[0055] Figure 3 A flowchart illustrating the construction process of a lightweight network provided in an embodiment of the present invention;

[0056] Figure 4 A schematic diagram illustrating the search strategy of a lightweight network in a multi-scale search space provided in an embodiment of the present invention;

[0057] Figure 5A schematic diagram illustrating the loss fitting between the pre-trained network architecture and the lightweight network provided in this embodiment of the invention;

[0058] Figure 6 This is a schematic diagram of the process for performing pruning operations on a lightweight network according to an embodiment of the present invention;

[0059] Figure 7 A schematic diagram illustrating the process of obtaining a defect sample set according to an embodiment of the present invention;

[0060] Figure 8 A flowchart illustrating the forward noise addition process of the lightweight diffusion model provided in this embodiment of the invention;

[0061] Figure 9 A flowchart illustrating the generation process of the lightweight diffusion model provided in this embodiment of the invention;

[0062] Figure 10 A flowchart illustrating the generation process of a high-quality defect image for a target lightweight diffusion model provided in this embodiment of the invention. Detailed Implementation

[0063] This invention provides a method and apparatus for constructing a lightweight diffusion model based on defect images, which addresses the technical problem of low computational efficiency in generating high-quality defect images using existing diffusion models based on defect images.

[0064] To make the objectives, features, and advantages of this invention more apparent and understandable, the technical solutions of the embodiments of this invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the embodiments described below are only some embodiments of this invention, and not all embodiments. Based on the embodiments of this invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this invention.

[0065] It should be noted that, in the optional embodiments of the present invention, the data related to object information, etc., requires the permission or consent of the object when the embodiments of the present invention are applied to specific products or technologies. Furthermore, the collection, use, and processing of the relevant data must comply with the relevant laws, regulations, and standards of the relevant countries and regions. In other words, if the embodiments of the present invention involve data related to an object, it needs to be obtained with the object's authorization and consent, the authorization and consent of relevant departments, and in accordance with the relevant laws, regulations, and standards of the country and region. If the embodiments involve personal information, the acquisition of all personal information requires the individual's consent. If sensitive information is involved, the separate consent of the information subject is required. The embodiments also need to be implemented with the object's authorization and consent.

[0066] Please see Figure 1This invention provides a method for constructing a lightweight diffusion model based on defect images, the method comprising:

[0067] Step 101: Obtain the original defect-free image and the defective image, and perform fusion processing on the original defect-free image and the defective image to obtain a fused image.

[0068] It should be noted that the original defect-free image Refers to high-quality, original, defect-free images, or defective images. This refers to an image that has various quality problems (such as dents, scratches, etc.) that correspond to the original defect-free image, such as a photograph of a circuit board with real scratches; the original defect-free image and the defective image must be paired, that is, the original defect-free image and the defective image must be taken of the same type of industrial product in the same industrial environment.

[0069] In this embodiment, after obtaining the paired original defect-free image and defective image, it is necessary to perform a fusion process on the original defect-free image and the defective image to extract their image feature information.

[0070] In one specific implementation method, the process of acquiring the fused image includes the following steps:

[0071] S11. Combine the original defect-free image and the defective image to obtain a stitched image;

[0072] S12. Perform convolution downsampling on the stitched image to restore the image resolution and obtain the fused image.

[0073] In the fusion process, the original defect-free image and the defective image are first concatenated. Then, a convolutional downsampling operation (PixelUnshuffle) is performed on the concatenated image to restore the image resolution. The fusion process can be represented as follows: Among them, the convolution downsampling operation can be performed using a 2*2 convolution kernel.

[0074] Step 102: Multi-scale feature extraction is performed on the fused image through a pre-trained network architecture to extract multiple scale features at different levels; pooling is performed on the lowest-level scale feature among the multiple scale features at different levels to obtain prior extracted features.

[0075] It should be noted that the pre-trained network architecture provided in this embodiment is used to extract multi-scale feature information from the fused image. The shallow scale features are used to capture details such as the texture of defect edges, while the higher-level scale features are used to capture semantic information. Then, pooling is performed using the lowest-level scale features to obtain the prior extracted features. The multiple different levels of scale features and prior extracted features provided in this embodiment provide a data foundation for lightweight training of the network.

[0076] In a specific implementation method, step 102 specifically includes the following steps:

[0077] S31. In the pre-trained network architecture, feature extraction operations are performed on the fused image through multiple sequentially connected residual blocks, and the scale features corresponding to each residual block are output.

[0078] S32. Select the last-level scale feature output by the last-level residual block in the pre-trained network architecture, and perform average pooling, linear normalization and activation function integer operation on the last-level scale feature in sequence to obtain the prior extracted features.

[0079] Please see Figure 2 The pre-trained network architecture provided in this specific embodiment includes multiple sequentially connected residual blocks; the feature extraction operation of each residual block includes two convolution operations and one activation function operation.

[0080] Specifically, the fused image is used as input to the pre-trained network architecture. Double sampling is performed in the residual block to extract scale features at multiple different levels from the fused image. The final scale features are then subjected to average pooling, linear normalization, and activation function integer operations to obtain prior extracted features. All of the above operations are the feature prior extraction (Fpe) operations of the pre-trained network architecture.

[0081] Among them, prior feature extraction The calculation process can be expressed as:

[0082]

[0083] Step 103: Based on the network search algorithm, optimize the network structure of the preset lightweight network according to the fused image, scale features of multiple different levels and prior extracted features, and output the primary lightweight network.

[0084] Understandably, web search algorithms have advantages such as fast search speed and low resource consumption. Please refer to [link / reference]. Figure 3 In this embodiment, a network search algorithm is used in a multi-scale search space. The algorithm takes the fused image, multiple different levels of scale features and prior extracted features as input, optimizes the structural parameters of the preset lightweight network, and determines the best combination of each module in the preset lightweight network, thereby obtaining the primary lightweight network.

[0085] This embodiment realizes the automatic design of lightweight network structure, avoiding the traditional design mode that relies on expert experience. The network obtained by architecture search significantly reduces parameter redundancy while maintaining the ability to extract features at multiple scales, laying a good foundation for subsequent convolution kernel pruning and optimizing computational efficiency and feature representation ability.

[0086] In one specific implementation method, step 103 includes the following steps:

[0087] S41. Input the fused image and prior extracted features into a preset lightweight network to output multiple scale secondary features at different levels;

[0088] S42. Calculate the loss values ​​between multiple scale sub-features at different levels and multiple scale features at different levels using a preset loss function.

[0089] S43. Based on the network search algorithm, optimize the network structure distribution of the lightweight network in the multi-scale search space according to the loss value, so as to determine the best combination connection of the standard module, downsampling module and prior extraction module under the same level, as well as the best internal selection set of the standard module and downsampling module, thereby outputting the primary lightweight network.

[0090] It should be noted that the preset lightweight network provided in this specific embodiment includes multiple standard modules, downsampling modules, and prior extraction modules of different scales; at the same level, the standard modules, downsampling modules, and prior extraction modules support direct connections and skip connections. Specifically, the standard modules are used to maintain the feature size, with a convolution stride of 1; the downsampling modules are used to downsample features to obtain multidimensional information, with a convolution stride of 2; and the prior extraction modules are used to introduce prior extracted features.

[0091] Please see Figure 4 Taking a total number of levels M=3 as an example, the standard module, downsampling module, and prior extraction module will be further explained:

[0092] (1) Standard Module:

[0093]

[0094] In the formula, This represents three different levels of scale. For the first The output scale characteristics of the standard modules at each level For the first Input characteristics of standard modules at each level; For mapping functions of standard modules, These are the learnable parameters for the standard module; the convolution stride of the standard module is set to 1, based on... and Extract and maintain the consistency of input and output feature sizes.

[0095] (2) Downsampling Module:

[0096]

[0097] In the formula, For the first The output scale characteristics of the hierarchical downsampling module. Downsampling operation, For the first The input features of the hierarchical downsampling module are set to 2, which enables downsampling of the feature map to obtain more dimensional information.

[0098] (3) Prior Extraction Module:

[0099]

[0100] In the formula, The output features of the prior extraction module, namely the prior extracted features, enable the lightweight network to jointly extract the prior features of the pre-trained network architecture, thereby enhancing the robustness and stability of the model.

[0101] In this specific embodiment, at the same level, the standard module, downsampling module, and prior extraction module support direct connections and skip connections; therefore, the multi-scale fusion feature at the s-th level is represented as:

[0102]

[0103] In the formula: This represents the multi-scale fusion feature of the s-th layer of a lightweight network; , and These are all learnable weight parameters, used to control the contributions of the standard module, downsampling module, and prior extraction module, respectively.

[0104] Thus, the entire multi-scale search space is complete. It can be represented as:

[0105]

[0106] In the above formula, the superscript 3 indicates the combination of feature inputs at three different scales.

[0107] To further streamline the structure of lightweight networks, the internal computational operations of the standard module and downsampling module can be further optimized. Specifically, this invention also includes: establishing an internal structure selection mechanism for the standard module and downsampling module and representing the internal structure of the standard module and downsampling module as a directed acyclic graph; under the internal structure selection mechanism, selecting and replacing the convolutional operations of the standard module or downsampling module through a preset internal selection set; wherein, the internal selection set includes skip connections, fully connected operations, and max pooling.

[0108] Understandingly, the internal structures of the standard and downsampling modules can be represented as a Directed Acyclic Graph (DAG). Each node in the DAG represents an operation (such as skip connections, fully connected operations, and max pooling), and each edge represents the flow of data between different operations, allowing for free combination of multiple inputs and outputs. In the DAG, data flows from one node to the next, progressively building and extracting feature information. Based on the overall requirements of the lightweight network, by analyzing the output requirements of each standard and downsampling module, the combination of their internal structures is determined, enabling the finding of a suitable optimal network architecture within a broad multi-scale search space.

[0109] Furthermore, considering the impact of noise on the feature maps obtained between nodes, a search strategy is designed by combining the scale features obtained from a supervised pre-trained network architecture with the multi-scale fusion features of a lightweight network. This aims to improve search efficiency while better acquiring multi-scale information, such as... Figure 5 As shown, the scale loss values ​​between multi-scale fusion features at multiple different levels and scale features at multiple different levels are calculated using a preset loss function. Based on the scale loss values, the network structure distribution of the lightweight network in the multi-scale search space is optimized to determine the optimal combination connection method of the standard module, downsampling module, and prior extraction module at the same scale, as well as the optimal internal selection set of the standard module and downsampling module, thereby outputting the primary lightweight network.

[0110] The preset loss function is expressed as:

[0111]

[0112] In the formula: This is the scale loss value. To merge images; For the scale features of the s-th layer of the pre-trained network architecture, The multi-scale fusion features at the s-th layer of the pre-defined lightweight network; Let be the number of nodes in the directed acyclic graph represented by the standard module and the downsampling module; s ranges from 2 to 1. .

[0113] Step 104: Prune the primary lightweight network to obtain the target lightweight network; establish an initial lightweight diffusion model with the target lightweight network as the backbone network.

[0114] Understandably, without affecting model performance, this embodiment further optimizes and compresses the initial lightweight network based on pruning techniques to improve the model's optimization efficiency. Using the target lightweight network obtained after pruning as the backbone network of the initial lightweight diffusion model significantly reduces memory usage and computational latency during inference while maintaining high-quality image generation capabilities, providing a feasible solution for deploying high-performance diffusion models in resource-constrained environments.

[0115] For a specific implementation method, please refer to Figure 6 The process of pruning a primary lightweight network to obtain the target lightweight network includes the following steps:

[0116] S51. Obtain the calibration image and input the calibration image into the primary lightweight network;

[0117] S52. Traverse each convolutional kernel in the primary lightweight network and calculate the feature map corresponding to the output of each convolutional kernel; calculate the corresponding feature map based on the feature map corresponding to the output of each convolutional kernel.

[0118] S53. Based on the preset egress threshold, remove convolutional kernels whose egress values ​​are less than the preset egress threshold corresponding to the feature map, and update the primary lightweight network; determine whether the updated primary lightweight network meets the preset conditions. If it does, jump to the step of traversing each convolutional kernel in the primary lightweight network and calculating the feature map output by each convolutional kernel; if it does not meet the conditions, use the primary lightweight network that previously met the preset conditions as the target lightweight network.

[0119] It should be noted that existing methods typically optimize network efficiency by adjusting hyperparameters of the network architecture, such as the number of feature channels, which can shorten search time. However, the optimization potential is very limited, and redundant convolutional kernels in the network still cause significant avoidable overhead and performance loss. Therefore, to reasonably reduce the number of convolutional kernels, this specific embodiment employs pruning techniques to optimize the structure of the primary lightweight network.

[0120] To ensure that the pruned lightweight network retains its core general feature extraction capabilities, this embodiment inputs a calibration image into the primary lightweight network. The importance of each convolutional kernel is measured by the alpha value corresponding to the feature map calculated by each convolutional kernel in the primary lightweight network. The calibration image is a high-quality, original, defect-free image.

[0121] The optimization objective of the pruning task in this specific embodiment is to selectively remove convolutional kernels while satisfying preset conditions (such as the performance metrics of the primary lightweight network), so that the sum of the tuples of the feature maps corresponding to the convolutional kernels of the pruned lightweight network is as small as possible; wherein, the optimization objective is expressed as:

[0122]

[0123] In the formula: This represents the total number of convolutional kernels in the initial lightweight network. Indicates the first The number of convolution kernels in each convolution operation. Indicates the first The first convolution operation One convolutional kernel, Indicates calibration image In the The convolution operation of the th _ Feature maps obtained from each channel This indicates the rank-finding operation; This indicates whether the convolution kernel should be pruned. A value of 1 indicates that the convolution kernel is pruned. A value of 0 indicates that the convolution kernel is retained.

[0124] In the pruning task, a preset threshold value S can be set. By continuously increasing S in an iterative operation, the above optimization objective is minimized, and the convolutional kernels with the smallest threshold value are iteratively removed to achieve the purpose of optimizing the lightweight network.

[0125] Step 105: Train the initial lightweight diffusion model using the defect sample set to generate the target lightweight diffusion model.

[0126] It should be noted that the defect sample set includes a large number of defect samples (defect images). The defect sample set is used to train the initial lightweight diffusion model, optimize the training parameters of the model, and generate the target lightweight diffusion model to achieve the generation of high-quality defect images.

[0127] For a specific implementation method, please refer to Figure 7 The process of obtaining the defect sample set includes the following steps:

[0128] S61. Generate a binary mask image based on the defect image using masking technology, and identify and extract the defect region in the defect image based on the binary mask image;

[0129] S62. By integrating and transforming the defective region and performing image enhancement processing on the original defect-free image, the corresponding defect template is obtained.

[0130] S63. Obtain multiple new defect-free images, perform defect fusion processing on each new defect-free image according to the defect template, and obtain multiple defect samples; the defect sample set includes multiple defect samples.

[0131] In practical image defect detection, the number of defective samples is usually far less than the number of normal samples, and the uneven distribution of positive and negative samples increases the risk of model overfitting. Therefore, this specific embodiment uses a defect template corresponding to the defect region of the defective image based on masking technology to integrate the original defect-free image with the defect template to generate a large number of defective samples, thereby increasing the number of defective samples and improving the accuracy of subsequent diffusion model training.

[0132] In this specific embodiment, the defective region of the defective image is masked using masking technology to form an original defective image with the defective region masked. Then, a new defect-free image similar to the defective background of the original defective image is obtained. It should be noted that the defective image, the original defect-free image, and the new defect-free image are all images actually captured by an industrial camera, not generated from a model, and all are images of the same model of product under the same environment.

[0133] The calculation method for defect samples is expressed as follows:

[0134]

[0135] In the formula: Indicates a defective sample; Indicates a defect template. Defect template Defect mask; This indicates a new, defect-free image. The function uses the defect template Binarization is performed to generate a binary mask image of the same size. The pixel value of the defect area in the binary mask image is 1, and the pixel value of the normal area is 0. This enables precise localization of the defect area and pixel-level fusion operation.

[0136] In a specific implementation method, step 105 specifically includes the following steps:

[0137] S71. Input the defective samples and the original defect-free image from the defective sample set into the initial lightweight diffusion model;

[0138] S72. Perform feature prior extraction based on defective samples and corresponding original defect-free images to obtain training prior extracted features; calculate guiding features based on training prior extracted features and corresponding original defect-free images.

[0139] S73. Add noise to the features extracted from the training prior, and use the added Gaussian noise as the training label to train the target lightweight network to obtain a noisy image, thereby realizing the forward noise training of the lightweight diffusion model.

[0140] S74. Based on noisy images, the guiding features are mapped to the intermediate representation layer of the target lightweight network through a cross-attention mechanism for iterative denoising prediction, thereby obtaining a predicted noisy image and realizing the reverse denoising training of the lightweight diffusion model.

[0141] S75. Calculate the total loss value of the training using the preset overall loss function, optimize the model parameters of the initial lightweight diffusion model using the total loss value, until the preset iteration conditions are met, and generate the target lightweight diffusion model.

[0142] It should be noted that the diffusion model has a rigorous mathematical derivation process. It first applies noise to the data step by step (forward noise addition process), and then learns how to gradually recover the data from the noise (reverse denoising process). In this specific embodiment, the initial lightweight diffusion model also includes an encoder and a decoder. The encoder and decoder use a convolutional neural network (CNN) as the backbone network to compress and retain important feature information.

[0143] First, please refer to Figure 8 via encoder Image fusion, multi-scale feature extraction, and pooling operations (i.e., feature prior extraction) are performed on defective samples and original defect-free images to obtain training prior extracted features. , ,in The image feature channel dimension; the above operation is represented as:

[0144]

[0145] Then, forward noise is added using a lightweight diffusion model. The forward noise addition process of the lightweight diffusion model can be expressed as:

[0146]

[0147] In the formula: t is the time step, , The number of iterations to add noise; The noise image is obtained after t rounds of noise iteration; = , is the cumulative product of image feature weights. In the early stages of adding noise to the model, this is done to maintain the correlation between the features and the previous time step as much as possible. The value is relatively large, and it increases as the iteration progresses. As the noise level gradually decreases, its proportion increases. It is a sample of random Gaussian noise that follows a standard normal distribution.

[0148] Save the Gaussian noise added each time during the noise addition process. The target lightweight network (i.e., U-Net network) with lightweight shared weights is trained using training labels, and the noise labels from the noise addition process are obtained through iterative steps.

[0149] For the inverse denoising process of the lightweight diffusion model, please refer to Figure 9 Extracting features using training priors and original defect-free images As a training guide, that is, extracting features from the training prior through the encoder. and original defect-free images Perform fusion and connection processing to obtain guiding features. , This will be used as one of the input terms in each iteration of the diffusion model. Represented as:

[0150]

[0151] A cross-attention mechanism based on guided features is constructed. Guided features are incorporated into each downsampling and upsampling layer of the U-Net network through a cross-attention layer mapping. This includes sequentially connecting the upsampling module and the downsampling module to form the U-Net network, with the upsampling module and the downsampling module being connected in a skip connection.

[0152] In cross-attention calculation, the query vector Derived from U-Net intermediate representations, and key-value vectors and By guiding features A simple linear transformation yields:

[0153]

[0154] in and These represent the linear projection layers that generate the key vector and value vector, respectively. Meanwhile, in the self-attention mechanism... This is a linear transformation derived from the intermediate representation layer of U-Net. The specific implementation process of the cross-attention mechanism in the U-Net network is shown below:

[0155]

[0156] in, This represents the value of self-attention. This represents the value of cross-attention. This represents the result of the fusion of cross-attention and self-attention. The query vector represents the intermediate representation layer of U-Net.

[0157] In lightweight diffusion models, features are extracted through training priors. Forward noise addition yields a noisy image. The U-Net network utilizes a cross-attention mechanism to... intermediate representation and By fusing these elements, the denoising process can reference prior knowledge and target background information at each step, outputting the predicted noise for the current step. This is used to calculate a clearer prediction of the noisy image at the next time step. ; and so on, until the final predicted noisy image is obtained. The cross-attention mechanism ensures that the process, starting from pure Gaussian noise, is stably guided towards generating a defect-free image. High-quality defect images with similar backgrounds and realistic defects.

[0158] To ensure the quality of defective image generation from the initial lightweight diffusion model, this specific embodiment combines image prediction loss and noise fitting loss to construct a preset overall loss function. The preset overall loss function is expressed as follows:

[0159]

[0160] In the formula: This is the total loss value. Image prediction loss, This represents the noise fitting loss.

[0161] Among them, image prediction loss The difference between the defective and non-defective images predicted by the diffusion model, i.e.:

[0162]

[0163] In the formula: It is an L1 norm.

[0164] Among them, noise fitting loss The noise loss used to fit forward noise addition and backward noise reduction is:

[0165]

[0166] In the formula: This represents the number of channels in the original feature map. Indicates the process Total number of fused feature channels after operation; and They represent the first The actual noise and predicted noise on each channel.

[0167] Minimizing the overall loss value mentioned above optimizes the image generation of the lightweight diffusion model using the supervised feature prior extraction proposed in this invention. While maintaining a lightweight network architecture and optimizing computational efficiency, it also ensures the quality of generated images with defects in the diffusion model as much as possible. Please refer to [link / reference]. Figure 10 After generating the target lightweight diffusion model, it can be used to perform the task of generating high-quality defect images. Specifically, a small number of original images (including defect-free and defective images) are input into a defect sample generation module based on traditional image fusion methods. This module outputs a diverse set of defect samples. The defect sample set is then input into the trained target lightweight diffusion model. Based on its powerful generation capabilities and defect knowledge learned from prior features, the target lightweight diffusion model optimizes, enhances, and diversifies the initial samples, ultimately automatically outputting higher-quality, more realistic, and more detailed defect images. This greatly reduces the cost of data collection and annotation, providing solid data support for defect detection technology.

[0168] The present invention also provides a computer system, including a memory, a processor, and a computer program stored in the memory, wherein the processor executes the computer program to implement the steps of any of the above methods for constructing a lightweight diffusion model based on a defect image.

[0169] Those skilled in the art will clearly understand that, for the sake of convenience and brevity, the specific working processes of the systems, devices, and units described above can be referred to the corresponding processes in the foregoing method embodiments, and will not be repeated here.

[0170] In the several embodiments provided in this application, it should be understood that the disclosed systems, apparatuses, and methods can be implemented in other ways. For example, the apparatus embodiments described above are merely illustrative; for instance, the division of units is only a logical functional division, and in actual implementation, there may be other division methods. For example, multiple units or components may be combined or integrated into another system, or some features may be ignored or not executed. Furthermore, the coupling or direct coupling or communication connection shown or discussed may be an indirect coupling or communication connection between apparatuses or units through some interfaces, and may be electrical, mechanical, or other forms.

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

[0172] Furthermore, the functional units in the various embodiments of the present invention can be integrated into one processing unit, or each unit can exist physically separately, or two or more units can be integrated into one unit. The integrated unit can be implemented in hardware or as a software functional unit.

[0173] If the integrated unit is implemented as a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, in essence, or the part that contributes to the prior art, or all or part of the technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods described in the various embodiments of the present invention. The aforementioned storage medium includes various media capable of storing program code, such as USB flash drives, portable hard drives, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0174] The above-described embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit it. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features. Such modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A method for constructing a lightweight diffusion model based on defective images, characterized in that, The method includes: Obtain the original defect-free image and the defective image, and then fuse the original defect-free image and the defective image to obtain the fused image; Multi-scale feature extraction is performed on the fused image using a pre-trained network architecture to extract multiple scale features at different levels; pooling is then performed on the lowest-level scale feature among these multiple scale features to obtain prior extracted features. Based on the network search algorithm, the network structure of the preset lightweight network is optimized according to the fused image, the scale features of multiple different levels and the prior extracted features, and the primary lightweight network is output. The primary lightweight network is pruned to obtain the target lightweight network; an initial lightweight diffusion model is established using the target lightweight network as the backbone network. Obtain a set of defective samples, train the initial lightweight diffusion model using the set of defective samples, and generate the target lightweight diffusion model.

2. The method for constructing a lightweight diffusion model based on defect images according to claim 1, characterized in that, The pre-trained network architecture includes multiple residual blocks connected in sequence; the feature extraction operation of each residual block includes two convolution operations and one activation function operation. The fused image is subjected to multi-scale feature extraction through a pre-trained network architecture, resulting in multiple scale features at different levels. The step of prior extraction of the lowest-level scale feature from the multiple different levels of scale features to obtain the prior extracted features includes: In the pre-trained network architecture, feature extraction operations are performed on the fused image through multiple sequentially connected residual blocks, and the scale features corresponding to each residual block are output. The final scale feature of the output of the last residual block in the pre-trained network architecture is selected, and average pooling, linear normalization and activation function integer operation are performed on the final scale feature in sequence to obtain the prior extracted feature.

3. The method for constructing a lightweight diffusion model based on defective images according to claim 1, characterized in that, The preset lightweight network includes multiple standard modules, downsampling modules, and prior extraction modules at different levels; at the same level, the standard modules, downsampling modules, and prior extraction modules support direct connection and skip connection. The standard module is used to maintain the feature size, and its convolution stride is set to 1; the downsampling module is used to downsample the features to obtain multidimensional information, and its convolution stride is set to 2; the prior extraction module is used to introduce the prior extracted features.

4. The method for constructing a lightweight diffusion model based on defect images according to claim 3, characterized in that, The method further includes: Establish an internal structure selection mechanism for the standard module and the downsampling module, and represent the internal structure of the standard module and the downsampling module as a directed acyclic graph; Under the internal structure selection mechanism, the convolutional operations of the standard module or the downsampling module are selected and replaced by a preset internal selection set; wherein, the internal selection set includes skip connections, fully connected operations and max pooling.

5. The method for constructing a lightweight diffusion model based on defect images according to claim 4, characterized in that, The step of optimizing the network structure of a preset lightweight network based on the network search algorithm, according to the fused image, the multiple different levels of scale features, and the prior extracted features, and outputting a primary lightweight network includes: The fused image and the prior extracted features are input into a preset lightweight network to output multiple multi-scale fused features at different levels; The scale loss value between the multi-scale fusion features at multiple different levels and the scale features at multiple different levels is calculated using a preset loss function. The network search algorithm optimizes the network structure distribution of the lightweight network in the multi-scale search space based on the scale loss value to determine the optimal combination connection of the standard module, downsampling module and prior extraction module at the same scale, as well as the optimal internal selection set of the standard module and downsampling module, thereby outputting the primary lightweight network. The calculation method for the multi-scale fusion features at different levels is as follows: In the formula: This represents the multi-scale fusion feature of the s-th layer of a lightweight network; This represents the output scale feature of the standard module at level s. This represents the output scale characteristics of the downsampling module at level s. This represents the output scale features of the prior extraction module; , and All of these are learnable weight parameters, used to control the contributions of the standard module, downsampling module, and prior extraction module, respectively. The preset loss function is expressed as: In the formula: This is the scale loss value. To merge images; For the scale features of the s-th layer of the pre-trained network architecture, The multi-scale fusion features at the s-th layer of the pre-defined lightweight network; This represents the number of nodes in the directed acyclic graph represented by the standard module and the downsampling module.

6. The method for constructing a lightweight diffusion model based on defect images according to claim 1, characterized in that, The step of pruning the primary lightweight network to obtain the target lightweight network includes: Acquire a calibration image and input the calibration image into the primary lightweight network; Iterate through each convolutional kernel in the primary lightweight network and calculate the feature map corresponding to the output of each convolutional kernel; calculate the corresponding feature map based on the feature map corresponding to each convolutional kernel. Based on a preset threshold, remove convolutional kernels whose thresholds are less than the preset threshold corresponding to the feature map, and update the primary lightweight network; determine whether the updated primary lightweight network meets a preset condition. If it does, proceed to the step of traversing each convolutional kernel in the primary lightweight network and calculating the feature map output by each convolutional kernel; if it does not meet the condition, use the primary lightweight network that previously met the preset condition as the target lightweight network.

7. The method for constructing a lightweight diffusion model based on defect images according to claim 1, characterized in that, The process of acquiring the fused image includes: The original defect-free image and the defective image are stitched together to obtain a stitched image; the stitched image is then subjected to a convolution downsampling operation to restore the image resolution, resulting in a fused image.

8. The method for constructing a lightweight diffusion model based on defect images according to claim 1, characterized in that, The process of obtaining the defect sample set includes: Based on the masking technology, a binary mask image is generated from the defect image, and the defect region in the defect image is identified and extracted based on the binary mask image. The defective region is integrated, transformed, and enhanced using the original defect-free image to obtain a corresponding defect template. Multiple new defect-free images are acquired, and each new defect-free image is subjected to defect fusion processing according to the defect template to obtain multiple defect samples; the defect sample set includes multiple defect samples. The calculation method for the defect samples is expressed as follows: In the formula: Indicates a defective sample; Indicates a defect template. Defect template Defect mask; This indicates a new, defect-free image.

9. The method for constructing a lightweight diffusion model based on defect images according to claim 8, characterized in that, The step of training the initial lightweight diffusion model using the defect sample set to generate the target lightweight diffusion model includes: The defective samples and the original defect-free image from the defective sample set are input into the initial lightweight diffusion model; Based on the defective sample and the corresponding original defect-free image, a feature prior extraction operation is performed to obtain the training prior extracted features; the guiding features are calculated based on the training prior extracted features and the corresponding original defect-free image. The training prior extracted features are subjected to forward noise addition, and the target lightweight network is trained with the added Gaussian noise as the training label to obtain a noisy image, thereby realizing the forward noise addition training of the lightweight diffusion model. Based on the noisy image, the guiding features are mapped to the intermediate representation layer of the target lightweight network through a cross-attention mechanism for iterative denoising prediction, thereby obtaining the predicted noisy image and realizing the reverse denoising training of the lightweight diffusion model. The total loss value of the training is calculated by a preset overall loss function, and the model parameters of the initial lightweight diffusion model are optimized by the total loss value until the preset iteration conditions are met, thereby generating the target lightweight diffusion model.

10. A computer system comprising a memory, a processor, and a computer program stored in the memory, characterized in that, The processor executes the computer program to implement the steps of the method for constructing a lightweight diffusion model based on defective images as described in any one of claims 1-9.

Images