Surface defect classification method based on generalized category discovery

Abstract

The invention discloses a surface defect classification method based on generalized category discovery. According to the method, firstly, a whole image data set is divided into two parts, one part is a labeled data set containing a known type, and the other part is a non-labeled data set; secondly, a defect classification network model is constructed, and a composite loss function is constructed through set category perception contrast learning and a category space contraction strategy to guide the training process of the model; after the total loss is obtained, calculating the weight parameters of the network model through back propagation; and finally, applying the trained model to a label-free defect data set, and testing the classification accuracy and generalization ability of the model. According to the method, limited label data can be effectively utilized, the model performance is improved, the limitation that the category which does not contain the label data in the unlabeled training data cannot be effectively processed in semi-supervised learning is solved, the method has excellent precision under the condition of a small amount of label data, and the method is very suitable for actual industrial surface defect detection application scenes.

Description

Technical Field

[0001] The present invention belongs to the technical field of automatic defect detection and relates to a surface defect classification method discovered in a broad category. Background Art

[0002] In the manufacturing quality control system, accurate detection and classification of product surface defects is a key link in ensuring product consistency and reliability. Traditional defect detection methods mainly rely on manual visual inspection or automated image processing technology based on threshold segmentation. Such methods not only have low detection efficiency, but also have inherent defects such as strong subjectivity and high false detection rate. With the rapid development of artificial intelligence technology, especially the breakthrough progress of deep learning methods in the field of computer vision, automated surface defect detection technology based on machine vision has achieved a qualitative leap. Among them, semi-supervised learning technology has shown significant application value under practical constraints such as high labeling costs and sample imbalance in industrial scenarios due to its ability to effectively integrate a small amount of labeled data with a large amount of unlabeled data. It is gradually becoming a research hotspot in the field of industrial quality inspection.

[0003] Although semi-supervised learning methods have made significant progress in the field of surface defect detection, they still have several key limitations in actual industrial application scenarios. Specifically, traditional semi-supervised learning algorithms face three major challenges in surface defect classification tasks: First, the algorithm architecture lacks support for open set recognition and cannot effectively handle new defect categories that may exist in unlabeled data, which seriously restricts the generalization ability of the model in real industrial environments; second, most existing methods are based on the closed world assumption, that is, the preset test data only contains the defect categories that have been labeled in the training set. This idealized assumption deviates significantly from the complex and changeable production conditions of industrial sites; finally, conventional semi-supervised learning relies on high-threshold screening to generate pseudo-labels. When faced with unknown categories, pseudo-labels often have high uncertainty and low sample utilization. These limitations make it difficult for traditional semi-supervised learning methods to achieve the expected results in demanding industrial quality inspection scenarios. Summary of the Invention

[0004] In view of the deficiencies of the prior art, the present invention provides a surface defect detection method based on broad categories.

[0005] The method of the present invention comprises the following steps:

[0006] Dataset partitioning: The entire image dataset is divided into two parts: one is a labeled dataset containing known types, and the other is an unlabeled dataset, which contains data of types that are not in the labeled dataset;

[0007] Constructing a defect classification network model: The designed defect classification network model consists of multiple residual blocks, and the gated attention unit module aggregates the spatial context of information. An auxiliary classification head is designed for the final output to help the model find the category contraction subspace during the training phase.

[0008] Constructing a composite loss function: By combining category-aware contrastive learning with a category space contraction strategy, a composite loss function is constructed to guide the model training process.

[0009] Iteratively train the network model weight parameters: After obtaining the total loss, backpropagate to calculate the network model weight parameters. By continuously adjusting the model weight parameters, the loss function value is minimized, thereby achieving the purpose of optimizing model performance;

[0010] Defect detection application: Apply the trained model to an unlabeled defect dataset to test the model's classification accuracy and generalization ability.

[0011] The beneficial effects of the present invention include at least: not only can it effectively utilize limited labeled data to improve model performance, but it also solves the limitation of semi-supervised learning that it cannot effectively process categories that are not included in the labeled data in the unlabeled training data. It has excellent accuracy with a small amount of labeled data and is very suitable for actual industrial surface defect detection application scenarios; secondly, it introduces category-level contrastive learning, and performs mixed supervised contrastive learning by screening samples through confidence thresholds, so that the intra-class samples are compact and the inter-class samples are clearly separated; again, in order to alleviate the category confusion of unknown categories, the pseudo-label space of unlabeled data is compressed, and the original low-confidence samples are converted into high-confidence samples by adaptively removing the confused categories and re-weighting them for use; and a gating mechanism is added for sparsification, which reduces the number of network parameters and the amount of calculation, thereby improving the speed and robustness of classification. BRIEF DESCRIPTION OF THE DRAWINGS

[0012] FIG1 is a flow chart of the steps of the surface defect classification method for general category discovery according to the present invention;

[0013] FIG2 is a diagram showing the model network structure of the surface defect classification method for generalized category discovery according to the present invention;

[0014] FIG3 is a framework diagram of the surface defect classification method for generalized category discovery according to the present invention. DETAILED DESCRIPTION

[0015] In order to make the technical problems, technical solutions and beneficial effects to be solved by the present invention more clearly understood, the present invention is further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention.

[0016] See also Figure 1, which is a flowchart of the steps of the surface defect classification method for general category discovery provided in an embodiment of the present application, including the following steps:

[0017] S1 divides the entire defect dataset into two parts: one part is a labeled dataset containing known defect types; the other part is an unlabeled dataset, which contains defect type data that is not in the labeled dataset.

[0018] S2 builds a defect classification network model: The designed defect classification network model consists of multiple residual blocks, and aggregates the spatial context of information through the gated attention unit module; an auxiliary classification head is designed for the final output to help the model find the category contraction subspace during the training phase.

[0019] S3 constructs a composite loss function to guide the model training process by combining category-aware contrastive learning and category space contraction strategy. Including category-aware contrastive loss L c By using class-aware clustering for in-distribution data and image-level contrastive learning for out-of-distribution data, the quality of pseudo-labels is improved. shrink , by adaptively detecting and removing confusing categories from the top categories to shrink the original category space. Supervision loss L with uncertainty adaptive margin s The calculation of is designed to handle category uncertainty; the pairwise target loss L P The calculation of promotes the distinguishability of the internal features of the model; the calculation of the regularization loss R is used to balance the generalization performance of the model.

[0020] After calculating the composite loss function, S4 uses this loss function-guided strategy to iteratively train the network model. By continuously adjusting the model weight parameters, the loss function value is minimized, thereby achieving the goal of optimizing model performance.

[0021] S5 applies the trained model to an unlabeled defect dataset to test the model's classification accuracy and generalization capabilities. This step verifies the model's ability to effectively identify and classify new, unknown defect types, thereby evaluating its practical application value.

[0022] Preferably, S1 specifically includes: collecting an industrial dataset containing 12,380 image samples from an industrial copper-clad laminate machine vision inspection system. These images are divided into two categories: severe defects and non-severe defects, with 5,140 non-severe defects and 7,240 severe defects. Fifty percent of all non-severe defect samples are classified as a labeled dataset (non-severe defects are classified as a known class, and severe defects are classified as a new (unknown) class). The remaining 50% of non-severe defect samples and all severe defect samples are classified as an unlabeled dataset.

[0023] When loading all data (labeled and unlabeled datasets) into the network, the data is loaded in batches, with each batch size set to 512. Data is sampled in a batch according to the specified weights. For labeled data, their weight is set to 1, and for unlabeled data items, their weight is set to the length of the unlabeled dataset divided by the length of the labeled dataset to balance the sampling of labeled and unlabeled data items. Testing is performed on all unlabeled datasets.

[0024] Preferably, the S2 specifically includes:

[0025] S2.1. The first layer uses 64 convolutional kernels with a receptive field size of 3×3, a stride of 1, and an edge padding of 1. Batch normalization and ReLU activation functions are then used. The output is a 64-channel feature map of size 224×224.

[0026] S2.2. Layers 2 through 5 are residual layers. Each layer consists of two residual blocks, each of which contains two 3×3 convolutional layers, each followed by batch normalization and Reluctant Unit (ReLU) activation. The input of each residual block is dimensionally aligned through a 1×1 convolutional layer and then added to the output of the block, followed by Reluctant Unit (ReLU) activation. The second residual block outputs a 64-channel, 224×224 feature map. The third residual block outputs a 128-channel, 112×112 feature map. The fourth residual block outputs a 256-channel, 56×56 feature map. The fifth residual block outputs a 512-channel, 28×28 feature map.

[0027] S2.3. The sixth layer is a gated attention unit. This unit feeds the 512-channel feature map into 512 1×1 convolution kernels with a normalized thickness, outputting a 512-channel feature map. This feature map is then split into two 256-channel feature maps, each fed into two branches. The upper branch consists of 64 7×1 depthwise separable convolution kernels with a stride of 1, padded by 3, and grouped into 64 (ensuring each convolution is applied to a single channel). The 256-channel feature map is fed into this branch, outputting a 256-channel feature map. The lower branch retains the 256-channel feature map. The outputs of the two branches, i.e., the two 256-channel feature maps, are multiplied together and fed into 512 1×1 convolution kernels, outputting a 512-channel feature map. A residual connection is then added to output a 512-channel feature map.

[0028] S2.4. The seventh layer is global average pooling, which performs global average pooling on the feature map output by the previous layer. Each channel becomes a numerical value, and the output is a 512-dimensional vector.

[0029] S2.5. The eighth layer is a parallel classification head. One branch is a fully connected layer that connects the 512-dimensional feature vector to the dimension of the number of categories. The other branch is an auxiliary classification head consisting of three linear layers, each followed by a batch normalization layer and a Leaky ReLu activation function. The first fully connected layer reduces the dimension from 512 to 128; the second fully connected layer maintains the dimension; and the third fully connected layer connects the dimension from 128 to the dimension of the number of categories. The auxiliary classification head is only effective during the training phase. The classification head during the inference phase connects the 512-dimensional feature vector to the dimension of the number of categories, ultimately classifying the unlabeled dataset.

[0030] S2.6. The ninth layer is the Softmax classification layer, which uses the Softmax layer to convert the feature vector into the probability distribution of the category.

[0031] Preferably, the S3 specifically includes:

[0032] Preferably, the category-aware contrastive learning loss is calculated. The core of calculating this loss is to achieve feature aggregation of known categories and noise robustness separation of unknown categories through a confidence-adaptive hybrid contrastive learning mechanism. The category-aware contrastive learning loss calculation process is as follows:

[0033] S3.1.1. Feature vector extraction: First, in each data batch, two different strong enhancements are performed on the unlabeled image Ux to obtain U s_x1 、U s_x2 After the two enhancement results are spliced ​​together, their feature vector Z is extracted through the pre-trained neural network model. i 、Z j .

[0034] S3.1.2. Compute the similarity matrix S: Using the dot product Z i 、Z j To calculate the similarity between positive sample pairs, and through the temperature parameter Scaling is used to adjust the sensitivity of the dot product of the feature vectors. The larger the dot product, the more similar the two vectors are. The formula is as follows:

[0035]

[0036] S3.1.3. Generate category perception matrix: pass threshold T push To divide the unlabeled image, the category perception matrix is ​​defined as follows:

[0037] W clacon,ij =

[0038] S3.1.4. Generate target weight matrix: A foreground reweighting mechanism is introduced to adjust the weight of each sample according to the pseudo-label confidence score to obtain the target weight matrix, which is defined as follows:

[0039] Wtarget,ij=

[0040] S3.1.5. The final loss function is:

[0041] Lc =H(S,Wtarget)=

[0042] Where H is the cross entropy loss, Pi is the index set of other perspectives of the same category, P(i)+1 represents the number of positive samples, and L c,i The expression is as follows, where Z i * Represents an enhancement of the same view:

[0043]

[0044] The enhancement in step S3.1.1 is a strong enhancement, which includes two types: one is a combination of cropping, horizontal flipping, and random changes in image brightness; the other is a combination of cropping, horizontal flipping, and Gaussian blurring. τ is a temperature adjustment parameter used to adjust the dot product scaling between feature vectors.

[0045] S3.2 Calculate the pseudo-label space shrinkage loss L shrink ,This strategy shrinks the original category space by adaptively detecting and removing ,confused categories in the top categories. The model learns and predicts in a more clearly ,defined space, which reduces the noise caused by category confusion and helps ,improve the overall classification performance and the robustness of the ,model.

[0046] S3.2.1 First, in each data batch, the unlabeled image Ux is strongly enhanced and weakly enhanced to obtain U s_x3 、U w_x1 , input the two images into the network model, and finally the strongly enhanced image is output from the auxiliary classification head to obtain the logit; the weakly enhanced image is input from the main classification head to obtain the logit. And sort them in descending order to obtain (Weakly enhanced image output probability), (Strongly enhanced image output probability) corresponding category .

[0047] S3.2.2 Then, a set of low-confidence but most likely categories is determined by two constraints. , to ensure the determinism of the top-1 category in the contracted space and select the largest space among all possible spaces. The constraints are as follows:

[0048]

[0049]

[0050] in The expression is as follows:

[0051]

[0052] S3.2.3 Since the top-1 probabilities vary among different uncertain images and the predictions are noisy and random at the beginning of training, we introduce two reweighting principles to address this problem.

[0053] 1. Image-based variable reliability reweighting, the formula is as follows:

[0054]

[0055] 2. Reweighting based on the model's gradual improvement state, the formula is as follows:

[0056]

[0057] Where mg is initialized to 0 and updated in each training iteration. The update formula is:

[0058]

[0059]

[0060] S3.2.4 Combining the above two principles, the uncertainty loss after reweighting is:

[0061]

[0062] S3.3. Computing the supervised loss L with uncertainty adaptive bounds z , this loss function is based on the modified cross entropy formula and takes into account the difference between seen and unseen classes by normalizing the logical value. Normalization includes an uncertainty term, which is calculated from the confidence score of the unlabeled data, so that the margin can adapt to the learning process of the model and decrease as the model gains certainty. This loss function enables the model to learn new categories without being overly biased towards seen classes, by effectively balancing the gradient updates between seen and new classes. Using standard cross entropy loss on labeled data will cause an imbalance problem, so an uncertainty adaptive margin mechanism is introduced, the mathematical expression is:

[0063]

[0064] Among them, P iis the prediction score of the labeled sample, λ is the regularization factor that defines its strength, and s is an additional scaling parameter that controls the temperature of the cross entropy loss. Calculated by the following formula:

[0065]

[0066] Among them D u is the set of all unlabeled samples, P r It is the probability distribution of the unlabeled sample after passing through the model network and passing through the softmax function. To describe it concisely, the uncertainty of an unlabeled sample can be quantified as 1 minus the highest predicted probability of the sample. represents the average of this uncertainty computed over the entire set of unlabeled samples.

[0067] S3.4. Paired Target Loss L p The purpose of this item is to ensure that labeled samples can be grouped with samples of the same category, and that unlabeled samples are grouped with their nearest neighbor samples (which may be labeled or unlabeled). The formula is as follows:

[0068]

[0069] where |Bl| and |Bu| are the number of labeled and unlabeled samples respectively, σ represents the softmax function, is the output prediction score of the sample input model network, if It has labels. is with Random labeled samples of the same category, if There are no labels. It's distance The most recent example (may be labeled or unlabeled).

[0070] S3.5. Regularization loss R. Regularization avoids assigning all instances to the same class, ensuring that the predicted probability distribution is close to the prior probability distribution. However, in practice, this prior knowledge is often not available. In this case, maximum entropy regularization is used to make the predicted probability distribution close to a uniform distribution. The formula is as follows:

[0071]

[0072] Where P(i) is the predicted probability for class i. Without any prior knowledge, entropy is equal for each P(i), i.e., it is maximized under uniform distribution. Maximum entropy is the result that maximizes entropy under known constraints. However, there are no constraints here, so uniform distribution is the desired maximum entropy. When entropy is used as a loss term, R = −H should be used. The reason for this is that the greater the entropy, the smaller the loss. This results in a more even distribution of the classifier's probability predictions, allowing new classes to be discovered.

[0073] Preferably, the S4 iterative training network model weight parameters: after obtaining the total loss, back propagation is used to calculate the network model weight parameters. The total loss is as follows. Before back propagation, the gradient of the previous step needs to be cleared, because PyTorch accumulates gradients by default in order to perform gradient accumulation. The model is trained using SGD with a momentum of 0.9 and a weight decay of 1×10 -4 , the model was trained for 200 epochs; the initial learning rate was 0.1, and the decay was scheduled using cosine annealing.

[0074]

[0075] Among them L C The coefficient mask.mean() is the average value of the maximum probability of each sample of the predicted unlabeled data after threshold filtering.

[0076] Preferably, the S5 uses the trained model to perform defect detection: testing is performed on an unlabeled dataset.

[0077] Preferably, calculate accuracy at the end: train the model on all datasets (without access to the true labels in DU) and at test time, measure the model predictions Clustering accuracy between i and the true label yi:

[0078]

[0079] Here, M=|D u | and P( i ) all permutations of class labels in the unlabeled set. The main metric is ACC on 'All' instances, which represents the entire unlabeled set |D u | The image recognition accuracy on the 'Old' class subset (|D u |Instances of known classes) and the 'New' class subset (|D u| instances belonging to the new class). The maximum value of the set of permutations is computed using the Hungarian optimal allocation algorithm. Importantly, this application only computes the Hungarian allocation once, across all classes, and then measures classification accuracy on the 'Old' and 'New' subsets.

[0080] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and are not limiting. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that the technical solutions of the present invention may be modified or replaced by equivalents without departing from the spirit and scope of the technical solutions of the present invention.

Claims

1. A surface defect classification method based on generalized category discovery, characterized in that: The following steps are involved: Dataset partitioning: The entire image dataset is divided into two parts: one is a labeled dataset containing known types, and the other is an unlabeled dataset, which contains data of types that are not in the labeled dataset; Constructing a defect classification network model: The designed defect classification network model consists of multiple residual blocks, and the gated attention unit module aggregates the spatial context of information. An auxiliary classification head is designed for the final output to help the model find the category contraction subspace during the training phase. Constructing a composite loss function: By combining category-aware contrastive learning and category space contraction strategies, a composite loss function is constructed to guide the model training process. The composite loss function includes: Class-aware contrastive loss improves the quality of pseudo-labels by using class-aware clustering for in-distribution data and image-level contrastive learning for out-of-distribution data; Pseudo-label shrinkage loss, which shrinks the original category space by adaptively detecting and removing confusing categories from the top categories; A supervised loss with uncertainty-adaptive margin, based on a modified cross-entropy formula, that takes into account the difference between seen and unseen classes by normalizing the logical values; Paired target loss, which promotes the discriminability of internal features of the model; Regularization loss, used to balance the generalization performance of the model; Iteratively train the network model weight parameters: After obtaining the total loss, backpropagate to calculate the network model weight parameters. By continuously adjusting the model weight parameters, the loss function value is minimized, thereby achieving the purpose of optimizing model performance; Defect detection application: Apply the trained model to an unlabeled defect dataset to test the model's classification accuracy and generalization ability.

2. The surface defect classification method based on generalized category discovery according to claim 1, characterized in that: In step 1, the data set is divided into: An industrial dataset containing multiple image samples is collected from the industrial copper clad laminate machine vision inspection system and divided into two categories: severe defects and non-severe defects. A part of the non-severe defect samples are classified as the labeled dataset, and the remaining non-severe defect samples and all severe defect samples are classified as the unlabeled dataset.

3. The surface defect classification method based on generalized category discovery according to claim 1, characterized in that: In step 2, the defect classification network model structure constructed is: The first convolution layer uses 64 convolution kernels with a receptive field size of 3x3, a stride of 1, and an edge padding of 1, followed by batch normalization and ReLU activation function; The second to fifth layers are residual layers, each layer contains two residual blocks, each residual block contains two 3x3 convolutional layers, and each convolutional layer is followed by batch normalization and ReLU activation; The sixth layer is a gated attention unit, which separates the feature map into two branches, multiplies them, and adds a residual connection; The seventh layer is the global average pooling layer, which converts each channel of the feature map into a value; The eighth layer is a parallel-structured classification head, which includes a main classification head and an auxiliary classification head. The auxiliary classification head takes effect during the training phase. The ninth layer is the Softmax classification layer, which converts the feature vector into the probability distribution of the category.

4. The surface defect classification method based on generalized category discovery according to claim 3, characterized in that: In step 3, the calculation of the category-aware contrastive learning loss includes the following steps: Two different strong enhancements are performed on unlabeled images, and their feature vectors are extracted using a pre-trained neural network model; Use dot product to calculate the similarity between positive sample pairs and scale it by the temperature parameter; Divide the unlabeled image by threshold and generate a category perception matrix; Adjust the weight of each sample according to the pseudo-label confidence score to generate a target weight matrix; Calculate the final loss function and use cross entropy loss for optimization.

5. The surface defect classification method based on generalized category discovery according to claim 3, characterized in that: In step 3, the calculation of the pseudo-label space shrinkage loss includes: Perform strong enhancement and weak enhancement on the unlabeled image respectively, and input the network model to obtain logit; Two constraints are used to determine the most likely category with low confidence, ensuring the certainty of the top-1 category. A reweighting principle based on image reliability and model improvement status is introduced to adjust the uncertainty loss.

6. The surface defect classification method based on generalized category discovery according to claim 3, characterized in that: In step 3, the calculation of the supervised loss with uncertainty adaptive margin is based on the modified cross entropy formula, which considers the difference between seen and unseen classes by normalizing the logical value, where the uncertainty is calculated by the highest predicted probability of the unlabeled samples.

7. The surface defect classification method based on generalized category discovery according to claim 3, characterized in that: In step 3, the calculation of paired target loss ensures that labeled samples and samples of the same category are grouped together, and that unlabeled samples and nearest neighbor samples are grouped together. The similarity between samples is calculated using the softmax function.

8. The surface defect classification method based on generalized category discovery according to claim 3, characterized in that: In step 3, the regularization loss adopts maximum entropy regularization to make the predicted probability distribution close to a uniform distribution, thereby promoting the discovery of new categories.

9. A surface defect classification method based on generalized category discovery according to any one of claims 4 to 8, characterized in that: In step 4, the SGD optimizer is used for model training, with momentum set to 0.9 and weight decay to 1×10 -4 , train for 200 epochs with an initial learning rate of 0.1, and use cosine annealing scheduling to decay the learning rate.

10. The surface defect classification method based on generalized category discovery according to claim 9, characterized in that: In step 5, the maximum value of the permutation set is calculated by the Hungarian optimal allocation algorithm to evaluate the classification accuracy of the model on the unlabeled dataset. The main measurement indicators include the image recognition accuracy of the entire unlabeled set and the accuracy of the known class and new class subsets.

Images