A small sample defect detection method based on a similar perception prototype migration

Abstract

This application relates to a few-sample defect detection method based on class-aware prototype transfer. The method involves inputting labeled images into a student network containing a first CAM module, outputting a first prediction result and first features corresponding to various ground truth labels extracted by the first CAM module; initially optimizing the student network by calculating a first loss function based on the first prediction result, and updating the prototypes of various ground truth labels based on the first features; inputting unlabeled images into a teacher network containing a second CAM module, calculating a comprehensive credibility score for pseudo-labels based on the second features output by the teacher network, the second prediction result, and the updated prototypes, and then further optimizing the student network by calculating a second loss function based on the comprehensive credibility score; updating the network parameters of the teacher network based on the network parameters of the student network after the secondary optimization; and finally obtaining the final defect detection result by passing the image of the industrial product to be detected through the iteratively optimized teacher network.

Description

Technical Field

[0001] This application relates to the field of industrial product defect detection technology, and in particular to a small-sample defect detection method based on class-aware prototype transfer. Background Technology

[0002] Existing industrial product defect detection methods face the following challenges in detecting defects in small samples:

[0003] 1. The model has weak cross-category transfer ability, and semantic drift will occur when the model is transferred to heterogeneous materials or new workpieces;

[0004] 2. High rate of missed detection of minute defects in complex industrial settings;

[0005] 3. Traditional pseudo-labeling technology may lead to error propagation and overfitting problems due to improper threshold settings. Summary of the Invention

[0006] Therefore, it is necessary to provide a small-sample defect detection method based on class-aware prototype transfer, including:

[0007] S1: Obtain a set of industrial product images for training, which includes an unlabeled dataset and a labeled dataset;

[0008] S2: Input the labeled images in the labeled dataset into the student network containing the first CAM module, output the corresponding first prediction result and the first features corresponding to various ground truth labels extracted by the first CAM module; calculate the first loss function based on the ground truth labels of the labeled images and the first prediction result to preliminarily optimize the student network, and update the prototype of various ground truth labels based on the first features corresponding to various ground truth labels.

[0009] S3: Input the unlabeled images in the unlabeled dataset into the teacher network containing the second CAM module, and output the corresponding second prediction result and the second features of each unlabeled image extracted by the second CAM module; calculate the comprehensive credibility score of the pseudo-label based on the second features, the updated prototype and the second prediction result; calculate the second loss function based on the comprehensive credibility score of the pseudo-label to optimize the student network for a second time; update the network parameters of the teacher network based on the network parameters of the student network after the second optimization.

[0010] S4: Iterate through steps S2-S3 until the maximum preset number of rounds is reached to obtain the optimized teacher network; pass the image of the industrial product to be detected through the optimized teacher network to obtain the final defect detection result.

[0011] Preferably, a YOLOv8 model pre-trained on the ImageNet dataset is used as the backbone network of the teacher network or the student network, and the C2f module in the YOLOv8 model is replaced with a first CAM module or a second CAM module to obtain the teacher network or the student network.

[0012] Preferably, the first prediction result includes a bounding box, a first category confidence score, and an object confidence score; a classification loss is calculated based on the first category confidence score and the ground truth label, a regression loss is calculated based on the bounding box, and a confidence loss is calculated based on the object confidence score; the classification loss, the regression loss, and the confidence loss are added together to obtain the first loss function; based on the first loss function, the network parameters of the student network are initially updated using a backpropagation algorithm.

[0013] Preferably, the method further includes pre-training the teacher network based on labeled images in the labeled dataset. The pre-trained teacher network outputs the second features corresponding to each type of ground truth label based on the labeled images, and calculates the mean of the second features corresponding to each type of ground truth label as the initial value of the prototype of the corresponding ground truth label.

[0014] Preferably, the prototype of the corresponding truth label is updated based on the first feature corresponding to the truth label, and the calculation formula is:

[0015] ;

[0016] in, This represents the prototype of the truth label of class c at the t-th iteration. This represents the momentum coefficient for prototype updates. This represents the prototype of the truth label of class c in the (t-1)th iteration. Let represent the set of labeled images that belong to the c-th truth label class at the t-th iteration. Describing the L1 norm, Represents a set Image with label The corresponding first feature.

[0017] Preferably, the comprehensive credibility score of the pseudo-label is calculated based on the second feature, the updated prototype, and the second prediction result, including:

[0018] The second prediction result includes the prediction box and the second category confidence score for the pseudo-label;

[0019] The predicted bounding boxes are filtered. For any filtered predicted bounding box, the cosine similarity between the predicted bounding box and the prototype is calculated based on the second feature of the corresponding unlabeled image and the ground value label corresponding to the pseudo label.

[0020] The cosine similarity is normalized, and the second category confidence score is weighted and summed with the corresponding normalized cosine similarity to obtain the comprehensive credibility score of the pseudo-label.

[0021] Preferably, the second loss function is calculated based on the comprehensive credibility score of pseudo-labels to optimize the student network for a secondary purpose, including:

[0022] Calculate the first average confidence level of the corresponding truth label based on the confidence level of each first category corresponding to any one type of truth label;

[0023] Calculate the second average confidence level of the corresponding true value label based on the confidence levels of each second category corresponding to the same true value label;

[0024] The standard deviation is calculated based on the second average confidence level, and the confidence difference of the corresponding true value labels is calculated based on the first average confidence level and the second average confidence level.

[0025] The adaptive low threshold and adaptive high threshold are calculated based on the confidence difference and standard deviation of the corresponding truth labels, respectively, and the calculation formulas are as follows:

[0026] ;

[0027] ;

[0028] ;

[0029] in, This represents the adaptive high threshold for the truth label of class c. This represents the adaptive low threshold for the truth label of class c. This refers to the clip function. This represents the first average confidence level of the truth labels for class c. This represents the second average confidence level of the truth labels for class c. Indicates the first hyperparameter. Indicates the second hyperparameter. This represents the third hyperparameter. This represents the difference in confidence level between the true value labels of class c. This represents the standard deviation of the second average confidence level of the true value labels for class c. Indicates the lower boundary of the threshold. Indicates the upper boundary of the threshold;

[0030] The confidence interval of the pseudo-label is determined by comparing the overall confidence score with the adaptive low threshold and the adaptive high threshold.

[0031] The corresponding distillation loss is selected based on the confidence interval of each pseudo-label, and the second loss function is calculated based on each distillation loss and the first loss function.

[0032] Based on the second loss function, the network parameters of the initially optimized student network are updated a second time using the backpropagation algorithm.

[0033] Preferably, the process of constructing the second loss function includes:

[0034] When the overall credibility score of the pseudo-label is greater than or equal to the adaptive high threshold, the sum of the second classification loss, the second regression loss and the second confidence loss during the teacher network pre-training is used as the distillation loss.

[0035] When the overall credibility score of the pseudo-label is between the adaptive low threshold and the adaptive high threshold, the second regression loss or the second confidence loss during the pre-training of the teacher network is used as the distillation loss.

[0036] When the overall credibility score of the pseudo-label is less than or equal to the adaptive low threshold, a background penalty is constructed as the distillation loss.

[0037] Based on the distillation loss of the pseudo-labels corresponding to all unlabeled images, the semi-supervised loss is calculated as follows:

[0038] ;

[0039] in, Indicates semi-supervised loss. This represents an unlabeled dataset. Represents unlabeled images The corresponding partition weights, Represents unlabeled images Corresponding distillation losses;

[0040] The second loss function is calculated based on the semi-supervised loss and the first loss function, and the formula is as follows:

[0041] ;

[0042] in, This represents the second loss function. Denotes the first loss function. This represents the weight of the semi-supervised loss.

[0043] Preferably, updating the network parameters of the teacher network based on the network parameters after secondary optimization of the student network includes:

[0044] An exponential moving average mechanism is used to update the network parameters of the teacher network based on the network parameters of the student network after secondary optimization.

[0045] Preferably, the defect detection results include defect category, defect location, and confidence level.

[0046] Beneficial effects: This method inputs labeled images into a student network containing a first CAM module, outputting the corresponding first prediction result and the first features corresponding to various ground truth labels extracted by the first CAM module; it then calculates a first loss function based on the first prediction result to initially optimize the student network, and updates the prototypes of various ground truth labels based on the first features; it inputs unlabeled images into a teacher network containing a second CAM module, calculates the comprehensive credibility score of pseudo-labels based on the second features output by the teacher network, the second prediction result, and the updated prototypes, and then calculates a second loss function based on the comprehensive credibility score to further optimize the student network; it updates the network parameters of the teacher network based on the network parameters after the secondary optimization of the student network; and iteratively optimizes the teacher network through the teacher network to obtain the final defect detection result. This method can make full use of a small amount of labeled data and a large amount of unlabeled data, significantly reducing labeling costs while achieving high-precision, high-speed, and strong generalization ability industrial defect detection. Attached Figure Description

[0047] To more clearly illustrate the technical solutions in the embodiments of this application 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 this application. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0048] Figure 1 This is a flowchart of a small-sample defect detection method based on class-aware prototype transfer in an embodiment of this application. Detailed Implementation

[0049] To make the above-mentioned objectives, features, and advantages of this application more apparent and understandable, the specific embodiments of this application are described in detail below with reference to the accompanying drawings. Many specific details are set forth in the following description to provide a thorough understanding of this application. However, this application can be implemented in many other ways different from those described herein, and those skilled in the art can make similar modifications without departing from the spirit of this application. Therefore, this application is not limited to the specific embodiments disclosed below.

[0050] Furthermore, the terms "first" and "second" are used for descriptive purposes only and should not be construed as indicating or implying relative importance or implicitly specifying the number of technical features indicated. Thus, a feature defined as "first" or "second" may explicitly or implicitly include at least one of that feature. In the description of this application, "multiple" means at least two, such as two, three, etc., unless otherwise explicitly specified.

[0051] like Figure 1As shown, this embodiment provides a small-sample defect detection method based on class-aware prototype transfer, including:

[0052] S1: Obtain a set of industrial product images for training, which includes an unlabeled dataset and a labeled dataset.

[0053] The ratio of the labeled dataset to the unlabeled dataset is 1:1.

[0054] In this embodiment, the industrial product image set used for training is the PKU-Market-PCB dataset, which contains six typical PCB defects (short circuit, open circuit, excess copper, via, burr, and scratch). All images in the dataset are uniformly adjusted to 640×640 pixels before use.

[0055] S2: Input the labeled images in the labeled dataset into the student network containing the first CAM module, output the corresponding first prediction result and the first features corresponding to the ground truth labels extracted by the first CAM module; calculate the first loss function based on the ground truth labels of the labeled images and the first prediction result to initially optimize the student network, and update the prototype of each ground truth label based on the first features corresponding to each ground truth label.

[0056] Specifically, the first prediction result includes a bounding box, a first class confidence score, and an object confidence score. A classification loss is calculated based on the first class confidence score and the ground truth label; a regression loss is calculated based on the bounding box; and a confidence loss is calculated based on the object confidence score. The classification loss (e.g., Focal Loss, used to address class imbalance), the regression loss (e.g., CIoULoss, used to accurately locate the bounding box), and the confidence loss (e.g., binary cross-entropy loss, used to determine if there is an object within the box) are added together to obtain the first loss function. Based on the first loss function, the network parameters of the student network are initially updated using a backpropagation algorithm.

[0057] In this embodiment, the teacher network is pre-trained based on labeled images in the labeled dataset. The pre-trained teacher network outputs the second features corresponding to each type of ground truth label based on the labeled images, and calculates the mean of the second features corresponding to each type of ground truth label as the initial value of the prototype of the corresponding ground truth label.

[0058] Furthermore, the pre-training process of the teacher network is as follows: a small number of labeled images from the labeled dataset are used to conduct supervised training on the teacher network, and the initial network parameters of the teacher network are optimized by minimizing the pre-training loss function; the pre-training loss function includes: second classification loss (such as Focal Loss), second regression loss (such as CIoU Loss), and second confidence loss (such as binary cross-entropy loss).

[0059] Furthermore, the initial value of the truth label prototype is calculated as follows:

[0060] ;

[0061] in, This represents the initial value of the prototype of the truth label of class c. This represents the set of labeled images corresponding to the c-th class ground truth label in the labeled dataset. Represents a set The i-th labeled image The corresponding second feature, Indicates the modulus.

[0062] Optionally, the prototype of the corresponding truth label is updated based on the first feature corresponding to the truth label, calculated as follows:

[0063] ;

[0064] in, This represents the prototype of the truth label of class c at the t-th iteration. This represents the momentum coefficient for prototype updates. , This represents the prototype of the truth label of class c in the (t-1)th iteration. Let represent the set of labeled images that belong to the c-th truth label class at the t-th iteration. Describing the L1 norm, Represents a set Image with label The corresponding first feature. If the set in the current iteration When empty, The original remains unchanged.

[0065] S3: Input the unlabeled images in the unlabeled dataset into the teacher network containing the second CAM module, and output the corresponding second prediction result and the second features of each unlabeled image extracted by the second CAM module; calculate the comprehensive credibility score of the pseudo-label based on the second features, the updated prototype and the second prediction result; calculate the second loss function based on the comprehensive credibility score of the pseudo-label to optimize the student network for a second time; update the network parameters of the teacher network based on the network parameters of the student network after the second optimization.

[0066] In this embodiment, a YOLOv8 model pre-trained on the ImageNet dataset is used as the backbone network of the teacher network or the student network. The C2f module in the YOLOv8 model is replaced with the first CAM module or the second CAM module to obtain the teacher network or the student network.

[0067] Furthermore, the first CAM module and the second CAM module have the same structure and workflow. The input to the CAM module is a feature map. C represents the number of channels, H represents the height of the image, and W represents the width of the image; Output: Enhanced feature map ;

[0068] Its internal processing flow is as follows:

[0069] 1. Convolutional branch (the cornerstone of local feature extraction);

[0070] Objective: To capture local details and spatial information of features.

[0071] Operation: F conv =SiLU(BatchNorm(Conv k1×k1 (F)));

[0072] Among them, F conv This represents convolutional features, SiLU(·) represents the SiLU function, and BatchNorm(·) represents batch normalization. It is usually a 1x1 or 3x3 convolution.

[0073] 2. Attention feature branch (global dependency modeling);

[0074] Objective: To model the global channel and spatial dependencies of features. Channel attention must be performed first, followed by spatial attention.

[0075] a. Channel Attention Submodule:

[0076] ; ;

[0077] ;

[0078] ;

[0079] ;

[0080] in, Represents the global max pooling feature of the channel. This represents the global average pooling characteristic of the channel. Indicates global max pooling. Indicates global average pooling. Indicates attention, This represents a multilayer perceptron. This represents the sigmoid function. Indicates channel attention. This indicates element-wise multiplication. This indicates the characteristics of channel fusion.

[0081] b. Spatial Attention Submodule:

[0082] ; ;

[0083] ;

[0084] ;

[0085] ;

[0086] in, Represents the global max pooling feature in space. This represents the spatial global average pooling feature. Indicates splicing characteristics, , Represents a 7×7 convolution. This represents the sigmoid function. Indicates spatial attention. This indicates element-wise multiplication. Indicates spatial integration characteristics, This indicates that max pooling is performed along the channel dimension to highlight the most salient features of each spatial location across all channels. This indicates that average pooling is performed along the channel dimension to capture the average feature of each spatial location across all channels.

[0087] c. Multi-scale enhancement mechanism (linked with category information):

[0088] (1). Scale Attention Weight Matrix (SAAM):

[0089] Input: Original features F and class prediction vectors from the network's classifier head .

[0090] Calculate the weights of each scale s: ;

[0091] Multi-scale convolution: ;

[0092] in, The weights of scale s are represented. The weighting coefficients for scale s are represented by... The weighting coefficients representing scale s' are... Represents multi-scale convolutional features. Indicates the use of size Convolution operations are performed using convolution kernels, with different scales s corresponding to different... , The value can be 1, 3, 5, etc., to capture features at different scales.

[0093] (2). Dynamic Category Adaptation (DCA):

[0094] Core inputs: the original features F and the prototype vectors of various truth labels maintained by the class-aware prototype EMA module. .

[0095] calculate:

[0096] Similarity chart: ;

[0097] Category-enhanced fusion: ;

[0098] in, This represents the similarity between the feature map F and the prototype vector. This indicates element-wise multiplication. This indicates a category-enhanced feature.

[0099] (3). Sparse attentional inhibition (SIA):

[0100] Function: As a regularization loss, it facilitates attention graph processing. and The sparsity and uniformity of the data help prevent overfitting of background noise.

[0101] formula: ;in, Describing the L1 norm, Represents variance. Indicates the suppression weight coefficient;

[0102] integrated: As an additional term added to the total loss function, the weights The recommended value is 0.0001.

[0103] d. Final Fusion:

[0104] Combine the outputs of all the above branches using element-wise addition:

[0105] F CAM =F conv +F att +F scale +F DCA ;

[0106] It will serve as a feature input for subsequent processes, and its quality directly determines the performance of the entire system.

[0107] Specifically, the comprehensive credibility score of the pseudo-label is calculated based on the second feature, the updated prototype, and the second prediction result, including:

[0108] The second prediction result includes the prediction box and the second category confidence score for the pseudo-label;

[0109] The predicted bounding boxes are filtered, for example using Gaussian weighting, to provide a cleaner and more accurate candidate set; specifically: ;in, This represents the score of the i-th predicted box after the update. This represents the score of the i-th prediction box. Indicates the attenuation parameter. Represents the IoU function. The prediction box that represents the current highest score. This represents the i-th prediction box to be suppressed, when A value greater than or equal to the intersection-union ratio (IU) threshold indicates high overlap, with a small exponent value (close to 0), thus significantly reducing the score. When the value is less than the intersection-union ratio threshold, it indicates low overlap, the value of the exponent term is relatively large (close to 1), and the score remains basically unchanged.

[0110] Optionally, the screening process includes:

[0111] Step 1: Sort all prediction boxes in descending order of score;

[0112] Step 2: Select the prediction box with the highest score And iterate through other prediction boxes to be suppressed. ;

[0113] Step 3: For each Calculate its relationship with the IoU function The intersection-union ratio is calculated, and the corresponding score is updated based on the Gaussian weighted formula;

[0114] Step 4: Repeat steps 2-3 until all prediction boxes have been processed. Compare the scores of all processed prediction boxes with a preset score threshold (e.g., 0.01). Filter out prediction boxes with scores lower than the preset score threshold to obtain the filtered prediction boxes.

[0115] By continuously reducing the scores of overlapping boxes, soft suppression of dense detection boxes is achieved, while retaining some potentially effective predicted boxes with reduced scores, thus improving the recall rate of detection.

[0116] For any filtered predicted bounding box, based on the second feature of the corresponding unlabeled image and the ground truth label corresponding to the pseudo label, the cosine similarity between the predicted bounding box and the prototype is calculated as follows:

[0117] ;

[0118] in, Represents unlabeled labeled images The obtained prediction box and the first Cosine similarity between prototypes after truth label updates. Represents unlabeled images The corresponding second feature, Indicates the first The prototype after the truth value label is updated. Represents the L2 norm;

[0119] The cosine similarity is normalized, and the second category confidence score is weighted and summed with the corresponding normalized cosine similarity to obtain the comprehensive credibility score of the pseudo-label. The calculation formula is as follows:

[0120] ;

[0121] ;

[0122] in, The overall credibility score of pseudo-labels. , These represent the first weighting coefficient and the second weighting coefficient, respectively. Represents normalization , Indicates the first The second category confidence level of the truth label.

[0123] Furthermore, a second loss function is calculated based on the comprehensive credibility score of pseudo-labels to optimize the student network for a secondary purpose, including:

[0124] Calculate the first average confidence level of the corresponding truth label based on the confidence level of each first category corresponding to any one type of truth label;

[0125] Calculate the second average confidence level of the corresponding true value label based on the confidence levels of each second category corresponding to the same true value label;

[0126] The standard deviation is calculated based on the second average confidence level, and the confidence difference of the corresponding true value labels is calculated based on the first average confidence level and the second average confidence level.

[0127] The adaptive low threshold and adaptive high threshold are calculated based on the confidence difference and standard deviation of the corresponding truth labels, respectively, and the calculation formulas are as follows:

[0128] ;

[0129] ;

[0130] ;

[0131] in, This represents the adaptive high threshold for the truth label of class c. This represents the adaptive low threshold for the truth label of class c. This refers to the clip function. This represents the first average confidence level of the truth labels for class c. This represents the second average confidence level of the truth labels for class c. Indicates the first hyperparameter. Indicates the second hyperparameter. This represents the third hyperparameter. This represents the difference in confidence level between the true value labels of class c. This represents the standard deviation of the second average confidence level of the true value labels for class c. Indicates the lower boundary of the threshold. Indicates the upper boundary of the threshold;

[0132] The confidence interval of the pseudo-label is determined by comparing the overall confidence score with the adaptive low threshold and the adaptive high threshold.

[0133] The corresponding distillation loss is selected based on the confidence interval of each pseudo-label, and the second loss function is calculated based on each distillation loss and the first loss function.

[0134] Based on the second loss function, the network parameters of the initially optimized student network are updated a second time using the backpropagation algorithm.

[0135] Furthermore, the process of constructing the second loss function includes:

[0136] In the high confidence interval, when the overall confidence score of the pseudo-label is greater than or equal to the adaptive high threshold, the sum of the second classification loss, the second regression loss, and the second confidence loss during the teacher network pre-training is used as the distillation loss.

[0137] In the medium confidence interval, when the overall confidence score of the pseudo-label is between the adaptive low threshold and the adaptive high threshold, the second regression loss or the second confidence loss during the pre-training of the teacher network is used as the distillation loss.

[0138] In the low-confidence interval, when the overall credibility score of the pseudo-label is less than or equal to the adaptive low threshold, the pseudo-label is determined to be background, and a background penalty is constructed as the distillation loss; the expression for the distillation loss is:

[0139] ;

[0140] in, Represents unlabeled images The corresponding distillation loss, This represents the third category loss (e.g., Focal Loss). This indicates the student network's response to unlabeled images. The predicted probability of the category. Indicates the background category (usually corresponding to category index 0). This represents the distillation weighting coefficient. This represents the third confidence level loss (such as binary cross-entropy loss). This indicates the student network's response to unlabeled images. The prediction confidence level This represents the target value without objects; distillation loss can be achieved using only the third-class loss.

[0141] Based on the distillation loss of the pseudo-labels corresponding to all unlabeled images, the semi-supervised loss is calculated as follows:

[0142] ;

[0143] in, Indicates semi-supervised loss. This represents an unlabeled dataset. Represents unlabeled images Corresponding partition weights (high confidence intervals) Distillation loss is a positive loss term; medium confidence interval. Low confidence interval Distillation loss is not effective. When the value is negative, the distillation loss will update the network parameters in reverse. Represents unlabeled images The corresponding distillation loss, in this embodiment, tends to treat the pseudo-label as background and apply a positive loss term;

[0144] The second loss function is calculated based on the semi-supervised loss and the first loss function, and the formula is as follows:

[0145] ;

[0146] in, This represents the second loss function. Denotes the first loss function. This represents the weight of the semi-supervised loss.

[0147] In this embodiment, updating the network parameters of the teacher network based on the network parameters after secondary optimization of the student network includes:

[0148] The network parameters of the teacher network are updated using an exponential moving average mechanism based on the network parameters of the student network after secondary optimization. The update formula is as follows:

[0149] ;

[0150] in, This represents the network parameters of the teacher network at the (t+1)th iteration. This represents the network parameters of the teacher network at the t-th iteration. This represents the network parameters of the student network after secondary optimization in the t-th iteration. This represents the second momentum coefficient, which is usually close to 1 (e.g., 0.996, which can be adjusted from 0.996 to 1.0 using cosine tuning).

[0151] Optionally, for both labeled and unlabeled images, a dual-path data augmentation mechanism can be used to expand the decision boundary between normal and abnormal samples in the feature space.

[0152] Teacher network path (weak enhancement): Apply basic geometric transformations (such as random horizontal flip, p=0.5) and slight color jitter to the input image to maintain the stability of feature prediction.

[0153] Student network path (strong enhancement): applies a more drastic combination of transformations to the input image, including:

[0154] ColorJitter: Adjust brightness, contrast, saturation, and hue within a range of 0.4.

[0155] RandomGrayscale: Converts to grayscale with a probability of 0.2.

[0156] RandomAffine: Random rotation (±10 degrees), shearing (±5 degrees).

[0157] Cutout: Randomly occludes a rectangular area with a maximum side length of 20%.

[0158] S4: Iterate through steps S2-S3 until the maximum preset number of rounds is reached to obtain the optimized teacher network; pass the optimized teacher network through the industrial product image to be detected to obtain the final defect detection result, including defect category, defect location and confidence level.

[0159] Specifically, after training, only the optimized teacher network is used for defect detection inference.

[0160] 1. Input: The image to be detected is weakly enhanced (or directly scaled) to a uniform size (e.g., 640x640 pixels).

[0161] 2. Forward propagation: The image is input into the teacher network, passes through the backbone network containing the CAM module and the detection head, and the prediction result is obtained.

[0162] 3. Post-processing: Use standard NMS (not Soft-NMS) to filter the predicted bounding boxes and output the final defect category, location, and confidence level.

[0163] Deployment: Because it is based on the efficient YOLOv8 architecture and has been optimized, the model is small in size and fast, and can be directly deployed on edge computing devices or servers in the factory to achieve real-time quality inspection.

[0164] Optionally, to prevent the accumulation of erroneous labels and improve training stability, an iterative error label tracing process is performed, including:

[0165] For each sample that is judged as an "erroneous label" (i.e., the overall credibility score of the pseudo-label is less than or equal to the adaptive low threshold), its defect feature vector is extracted, stored in the "dynamic error feature library", and the iteration round is marked.

[0166] After generating pseudo-labels in the next round, calculate the cosine similarity between the second feature and the features in the database. If the similarity... Even if its confidence level is slightly higher than the adaptive low threshold, it is still forcibly classified as an "incorrect label" and discarded.

[0167] Clean up the "Dynamic Error Feature Library" regularly (e.g., every 100 rounds) and delete old features that have been marked for more than 50 rounds and have not reappeared.

[0168] Optionally, the method can also perform general object detection based on images from the COCO2017 data.

[0169] The small-sample defect detection method based on class-aware prototype transfer provided in this embodiment has the following beneficial effects:

[0170] 1. High detection accuracy and strong generalization: By stabilizing the feature space through prototypes of ground truth labels and refining the use of unlabeled data through adaptive thresholding, the model achieves leading detection accuracy (mAP) on datasets such as COCO2017 and PCB, especially performing well in cross-material and cross-workpiece scenarios.

[0171] 2. High training efficiency and stability: The teacher-student network is updated smoothly through EMA (exponential moving average) and combined with a dynamic threshold strategy to accelerate model convergence, suppress training oscillations, and make the training process more stable.

[0172] 3. Low annotation cost: This semi-supervised framework only requires a small amount of labeled data (such as 5% or 10%) to achieve performance close to that of fully supervised learning, which greatly reduces the annotation cost of industrial quality inspection.

[0173] 4. Real-time detection capability: Based on an optimized single-stage detector, the number of model parameters and computational load are effectively controlled, meeting the needs of real-time detection in industrial settings.

[0174] The technical features of the above embodiments can be combined in any way. For the sake of brevity, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.

[0175] The embodiments described above are merely illustrative of several implementation methods of this application, and while the descriptions are relatively specific and detailed, they should not be construed as limiting the scope of the patent application. It should be noted that those skilled in the art can make various modifications and improvements without departing from the concept of this application, and these all fall within the protection scope of this application. Therefore, the protection scope of this patent application should be determined by the appended claims.

Claims

1. A small-sample defect detection method based on class-aware prototype transfer, characterized in that, include: S1: Obtain a set of industrial product images for training, which includes an unlabeled dataset and a labeled dataset; S2: Input the labeled images in the labeled dataset into the student network containing the first CAM module, and output the corresponding first prediction result and the first features corresponding to the various ground truth labels extracted by the first CAM module; The student network is initially optimized by calculating the first loss function based on the ground truth labels of labeled images and the first prediction results, and the prototypes of various ground truth labels are updated based on the first features corresponding to various ground truth labels. S3: Input the unlabeled images in the unlabeled dataset into the teacher network containing the second CAM module, and output the corresponding second prediction result and the second features of each unlabeled image extracted by the second CAM module; calculate the comprehensive credibility score of the pseudo-label based on the second features, the updated prototype and the second prediction result; The student network is then optimized by calculating a second loss function based on the comprehensive credibility score of pseudo-labels, including: Calculate the first average confidence level of the corresponding truth label based on the confidence level of each first category corresponding to any one type of truth label; Calculate the second average confidence level of the corresponding true value label based on the confidence levels of each second category corresponding to the same true value label; The standard deviation is calculated based on the second average confidence level, and the confidence difference of the corresponding true value labels is calculated based on the first average confidence level and the second average confidence level. The adaptive low threshold and adaptive high threshold are calculated based on the confidence difference and standard deviation of the corresponding truth labels, respectively, and the calculation formulas are as follows: ； ； ； in, This represents the adaptive high threshold for the truth label of class c. This represents the adaptive low threshold for the truth label of class c. This refers to the clip function. This represents the first average confidence level of the truth labels for class c. This represents the second average confidence level of the truth labels for class c. Indicates the first hyperparameter. Indicates the second hyperparameter. This represents the third hyperparameter. This represents the difference in confidence level between the true value labels of class c. This represents the standard deviation of the second average confidence level of the true value labels for class c. Indicates the lower boundary of the threshold. Indicates the upper boundary of the threshold; The confidence interval of the pseudo-label is determined by comparing the overall confidence score with the adaptive low threshold and the adaptive high threshold. The corresponding distillation loss is selected based on the confidence interval of each pseudo-label, and the second loss function is calculated based on each distillation loss and the first loss function. Based on the second loss function, the network parameters of the initially optimized student network are updated a second time through the backpropagation algorithm. Update the network parameters of the teacher network based on the network parameters of the student network after secondary optimization. S4: Iterate through steps S2-S3 until the maximum preset number of rounds is reached to obtain the optimized teacher network; pass the image of the industrial product to be detected through the optimized teacher network to obtain the final defect detection result.

2. The small sample defect detection method according to claim 1, characterized in that, The YOLOv8 model pre-trained on the ImageNet dataset is used as the backbone network of the teacher network or the student network. The C2f module in the YOLOv8 model is replaced with the first CAM module or the second CAM module to obtain the teacher network or the student network.

3. The small sample defect detection method according to claim 1, characterized in that, The first prediction result includes a bounding box, a first category confidence score, and an object confidence score. A classification loss is calculated based on the first category confidence score and the ground truth label, a regression loss is calculated based on the bounding box, and a confidence loss is calculated based on the object confidence score. The classification loss, the regression loss, and the confidence loss are added together to obtain the first loss function. Based on the first loss function, the network parameters of the student network are initially updated using a backpropagation algorithm.

4. The small sample defect detection method according to claim 3, characterized in that, It also includes pre-training the teacher network based on labeled images in a labeled dataset. The pre-trained teacher network outputs the second features corresponding to each type of ground truth label based on the labeled images, and calculates the mean of the second features corresponding to each type of ground truth label as the initial value of the prototype of the corresponding ground truth label.

5. The small sample defect detection method according to claim 1, characterized in that, The prototype of the corresponding truth label is updated based on the first feature corresponding to the truth label. The calculation formula is as follows: ； in, This represents the prototype of the truth label of class c at the t-th iteration. This represents the momentum coefficient for prototype updates. This represents the prototype of the truth label of class c in the (t-1)th iteration. Let represent the set of labeled images that belong to the c-th truth label class at the t-th iteration. Describing the L1 norm, Represents a set Image with label The corresponding first feature.

6. The small sample defect detection method according to claim 4, characterized in that, The overall credibility score of the pseudo-label is calculated based on the second feature, the updated prototype, and the second prediction result, including: The second prediction result includes the prediction box and the second category confidence score for the pseudo-label; The predicted bounding boxes are filtered. For any filtered predicted bounding box, the cosine similarity between the predicted bounding box and the prototype is calculated based on the second feature of the corresponding unlabeled image and the ground value label corresponding to the pseudo label. The cosine similarity is normalized, and the second category confidence score is weighted and summed with the corresponding normalized cosine similarity to obtain the comprehensive credibility score of the pseudo-label.

7. The small sample defect detection method according to claim 1, characterized in that, The process of constructing the second loss function includes: When the overall credibility score of the pseudo-label is greater than or equal to the adaptive high threshold, the sum of the second classification loss, the second regression loss and the second confidence loss during the teacher network pre-training is used as the distillation loss. When the overall credibility score of the pseudo-label is between the adaptive low threshold and the adaptive high threshold, the second regression loss or the second confidence loss during the pre-training of the teacher network is used as the distillation loss. When the overall credibility score of the pseudo-label is less than or equal to the adaptive low threshold, a background penalty is constructed as the distillation loss. Based on the distillation loss of the pseudo-labels corresponding to all unlabeled images, the semi-supervised loss is calculated as follows: ； in, Indicates semi-supervised loss. This represents an unlabeled dataset. Represents unlabeled images The corresponding partition weights, Represents unlabeled images Corresponding distillation losses; The second loss function is calculated based on the semi-supervised loss and the first loss function, and the formula is as follows: ； in, This represents the second loss function. Denotes the first loss function. This represents the weight of the semi-supervised loss.

8. The small sample defect detection method according to claim 1, characterized in that, The network parameters of the teacher network are updated based on the network parameters after secondary optimization of the student network, including: An exponential moving average mechanism is used to update the network parameters of the teacher network based on the network parameters of the student network after secondary optimization.

9. The small sample defect detection method according to claim 1, characterized in that, The defect detection results include defect category, defect location, and confidence level.

Images