Semi-supervised surface defect instance segmentation method based on pseudo label enhancement

Abstract

The application provides a semi-supervised surface defect instance segmentation method based on pseudo label enhancement, comprising the following steps: obtaining a surface defect image dataset; training an initial instance segmentation network with a small amount of labeled defect image data; inputting the labeled defect image data into a student instance segmentation network for training; inputting unlabeled defect image data into a teacher instance segmentation network and the student instance segmentation network to obtain the mask and features of each defect instance; then inputting the mask and features into a cross-supervised contrast learning module and a general distribution fusion module; constructing a contrast learning loss function, obtaining fused pseudo labels, training the student instance segmentation network with the pseudo labels, and updating the parameters of the teacher instance segmentation network using the parameters of the student instance segmentation network. The application can cope with the challenges of small data set size and fuzzy defect boundary, estimate complex class feature distribution through the general distribution fusion module, and improve the quality of the mask pseudo labels from the teacher model.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of defect detection, and particularly to a semi-supervised surface defect instance segmentation method based on pseudo-label enhancement. BACKGROUND

[0002] Surface defect detection ensures the quality of industrial products by accurately identifying and classifying surface defects. Thanks to the rapid development of artificial intelligence and convolutional neural networks, numerous defect detection methods have made significant progress in different fields. Many methods focus on conventional computer vision tasks such as classification, semantic segmentation, and object detection, assuming consistent defect data distribution and sufficient labeled data. At the same time, some researchers have also been committed to estimating the fuzzy edges of defect regions and proposing improved edge detection methods based on traditional filtering techniques and Gaussian functions. However, obtaining labeled defect data is still a difficult, expensive, and time-consuming task, which poses a challenge to vision-based industrial detection. In addition, directly locating the edges does not fully utilize the information of the fuzzy regions, which to some extent affects the performance of the model in defect detection.

[0003] Semi-supervised defect detection has attracted increasing attention, and existing technologies enhance the defect segmentation capability of the model by using generative adversarial networks (GANs) to synthesize data and improve the performance of the model by integrating multi-source information. And the existing technology builds a simple semi-supervised teacher-student model framework, the student model is trained using the pseudo-labels generated by the teacher model prediction, and the parameters of the teacher model are updated from the student model through the exponential moving average (EMA) strategy. At the same time, the existing technology reduces the difference between the prediction distribution of the student model and the teacher model or the difference between different input scales to improve the performance of the framework. However, these studies mainly focus on classification, detection, and segmentation tasks, ignoring instance segmentation. Instance segmentation not only accurately locates defects and provides detailed information, but also maximizes the reduction of background pixels on labeling. In addition, the method of generating pseudo-labels based on threshold values is prone to model overfitting, making it difficult to accurately identify defects in unlabeled images.

[0004] However, existing semi-supervised surface defect detection methods mainly focus on classification, detection, and segmentation tasks, ignoring instance segmentation. Instance segmentation not only accurately locates defects and provides detailed information, but also maximizes the reduction of background pixels on labeling. In addition, the method of generating pseudo-labels based on threshold values is prone to model overfitting, making it difficult to accurately identify defects in unlabeled images.

[0005] Secondly, the pseudo-label generated by the existing method is inaccurate. Since the pseudo-label is automatically generated by the model, rather than based on artificial labeling, its accuracy cannot be guaranteed. When the accuracy of the pseudo-label is not high, the model will learn these inaccurate pseudo-labels as real labels in the subsequent iterative training process, gradually accumulating noise. As the training progresses, the accumulation of noise may cause the model to judge defects more and more deviate from the actual situation, thereby reducing the overall performance of the model. This phenomenon is particularly evident in the long-term iterative process. Furthermore, the edges of defects are usually fuzzy, especially for smaller defects, which leads to the fact that the existing semi-supervised surface defect detection method is prone to mispredict the defect edges, seriously affecting its performance. SUMMARY

[0006] In view of the deficiencies of the prior art, the present application provides a semi-supervised surface defect instance segmentation method based on pseudo-label enhancement.

[0007] The technical scheme of the present application is: a semi-supervised surface defect instance segmentation method based on pseudo-label enhancement, comprising the following steps:

[0008] S1), obtain a surface defect image dataset; and perform instance segmentation format labeling on part of the defect images;

[0009] S2), train an initial instance segmentation network with a small amount of labeled defect image data; and use the parameters of the initial network to initialize the parameters of the student instance segmentation network and the teacher instance segmentation network;

[0010] S3), input the labeled defect image data into the student instance segmentation network, combine the output of the student instance segmentation network with the labeling of the defect image data, and train the student instance segmentation network;

[0011] S4), input the unlabeled defect image data into the teacher instance segmentation network to obtain the mask of each defect instance output by the teacher instance segmentation network;

[0012] Meanwhile, input the unlabeled defect image data into the student instance segmentation network to obtain the features output by the backbone network of the student instance segmentation network;

[0013] S5), input the mask of each defect instance output by the teacher instance segmentation network and the features output by the backbone network of the student instance segmentation network into the cross-supervised contrast learning module; construct a contrast learning loss function;

[0014] S6), input the mask of each defect instance output by the teacher instance segmentation network and the features output by the backbone network of the student instance segmentation network into the general distribution fusion module to obtain the fused pseudo-label, and use the pseudo-label to train the student instance segmentation network;

[0015] S7), updating the parameters of the teacher instance segmentation network in a manner of momentum iteration using the parameters of the student instance segmentation network.

[0016] As preferred, in step S2), the initial instance segmentation network is trained with a small amount of labeled image data, specifically comprising the following steps:

[0017] S21), inputting the labeled defect image data into the instance segmentation network, after the defect image passes through the feature extractor and the decoder of the instance segmentation network, finally outputting the prediction of each instance of the defect image, including the class prediction of the instance and the mask prediction of the instance

[0018] S22), calculating the loss function of the predicted class and the mask according to the class y i and the mask m i of the real instance in the annotation of the defect image data; wherein the predicted class loss function is a cross-entropy loss function; the loss function of the mask is a mean square error loss function:

[0019] The cross-entropy loss function L CE is:

[0020] The mean square error loss function L MSE is:

[0021] As preferred, in step S4), the unlabeled defect image is output to the teacher instance segmentation network, and after the image passes through the feature extractor and the decoder of the teacher network, finally outputting the prediction of each instance of the image, including the class prediction of the instance and the mask prediction of the instance.

[0022] As preferred, in step S5), the value range of the mask of the defect instance is [0, 1], and the mask of the defect instance is divided into four parts according to the value of the mask, which are: (0-0.25] part is simple background area, [0.25-0.5] is difficult background area, [0.5-0.75] part is difficult defect area, [0.75-1] is simple defect area.

[0023] As preferred, in step S5), for each instance of each image, the feature C output by the backbone network of the student instance segmentation network is a tensor of [c, h, w], wherein c, h and w represent the channel dimension, length and width of the feature C respectively.

[0024] As preferred, in step S5), the four parts m1-m4 of the mask and the features C output by the backbone network of the student instance segmentation network are subjected to tensor multiplication, and then global average pooling is performed to obtain vectors v1-v4 of [c, 1, 1], i.e., the feature centers of the instances; the class of the feature center is the predicted class for the instance, and the calculation expression of the vectors v1-v4 is:

[0025] v i = Global(m i *C), i = 1, 2, 3, 4.

[0026] wherein Global is a global average pooling operation, and the average value is calculated in the [h, w] dimension of the [c, h, w] tensor to obtain a vector v i .

[0027] As preferred, in step S5), the feature centers of each instance calculated for the K defect classes and the classes of the four mask regions are collected into different memory banks, and the instance centers in the memory bank of the same defect class and the class of the mask region are positive samples, and the instance centers in the memory bank of other defect classes and other mask region classes are negative samples, and a contrast learning loss function is constructed.

[0028] As preferred, in step S5), the expression of the constructed contrast learning loss function L is:

[0029]

[0030] wherein sim is a similarity calculation function, sim(z i ,z j ) = z i T *z j , and the calculation method is the inner product of two vectors; τ is a temperature coefficient; n is the number of positive samples, and N is the total number of positive samples and negative samples; z j is a positive sample of the instance center in the memory bank of the same defect class and the class of the mask region as z i , and z k is the instance center of all memory banks.

[0031] As preferred, in step S6), the features C output by the backbone network of the student instance segmentation network of each image and the mask of the defect instance generated by the teacher instance segmentation network for the unlabeled image are input into the general distribution fusion module, and the mask and the features C of each defect instance are multiplied pixel by pixel to obtain the features [m, c, h, w] covered by each defect instance mask.

[0032] As preferred, in step S6), the mean value μ and the variance σ of each feature center in the memory bank are calculated in each dimension; then the probability γ that all features C covered by the defect instance mask belong to the Gaussian distribution of different memory banks is calculated according to the mean value μ and the variance σ in the memory bank, that is:

[0033]

[0034] In the formula, v is the feature center; μ k is the mean value of the kth instance; σ k is the variance of the kth instance.

[0035] As preferred, in step S6), the probability [m, 4*K, h, w] obtained is taken as the maximum in the second dimension to obtain a tensor [m, h, w] as a new mask, wherein m is the total number of memory banks.

[0036] As preferred, in step S6), the mean value of the new mask and the old mask is calculated as an enhanced new mask, and the enhanced new mask is used to train the student instance segmentation network.

[0037] The beneficial effects of the present application are:

[0038] 1. The present application can cope with the challenges of small data set size and fuzzy defect boundary, and the general distribution fusion module of the present application can estimate complex class feature distribution, and improve the quality of pseudo-labels from the teacher model mask;

[0039] 2. The cross-supervised contrast learning module of the present application can distinguish easy / difficult areas and defect / background areas, and can successfully handle the fuzzy problem near the defect boundary. BRIEF DESCRIPTION OF DRAWINGS

[0040] Figure 1 is a flowchart of the method of the present application;

[0041] Figure 2 is a flowchart of the method of the present application; DETAILED DESCRIPTION

[0042] The specific embodiments of the present application will be further described below in conjunction with the accompanying drawings:

[0043] As Figure 1 shown, the semi-supervised surface defect instance segmentation method based on pseudo-label enhancement comprises:

[0044] S1), obtain a surface defect image data set; and perform instance segmentation format annotation on part of the defect images;

[0045] S2), training an initial instance segmentation network with a small amount of labeled defect image data; and using the parameters of the initial network to initialize the parameters of the student instance segmentation network and the teacher instance segmentation network; specifically comprising the following steps:

[0046] S21), inputting the labeled defect image data into the instance segmentation network, and after the defect image passes through the feature extractor and the decoder of the instance segmentation network, finally outputting the prediction of each instance of the defect image, including the class prediction of the instance and the mask prediction of the instance

[0047] S22), calculating the loss function of the predicted class and the mask according to the class y i and the mask m i of the real instance in the annotation of the defect image data; wherein the predicted class loss function is a cross-entropy loss function; and the loss function of the mask is a mean square error loss function:

[0048] The cross-entropy loss function L CE is:

[0049] The mean square error loss function L MSE is:

[0050] S3), inputting the labeled defect image data into the student instance segmentation network, and training the student instance segmentation network in combination with the output of the student instance segmentation network and the annotation of the defect image data;

[0051] S4), inputting the unlabeled defect image data into the teacher instance segmentation network to obtain the mask of each defect instance output by the teacher instance segmentation network; in this embodiment, the unlabeled defect image is output into the teacher instance segmentation network, and after the image passes through the feature extractor and the decoder of the teacher network, finally outputting the prediction of each instance of the image, including the class prediction of the instance, and the mask prediction of the instance.

[0052] Meanwhile, input the unlabeled defect image data into the student instance segmentation network to obtain the feature output by the backbone network of the student instance segmentation network;

[0053] S5), inputting the mask of each defect instance output by the teacher instance segmentation network and the feature output by the backbone network of the student instance segmentation network into the cross-supervised contrast learning module; and constructing a contrast learning loss function;

[0054] In the embodiment, the value range of the mask of the defect instance is [0, 1], and the mask of the defect instance is divided into four parts according to the value of the mask, which are respectively: the part of (0-0.25] is a simple background area, the part of [0.25-0.5] is a difficult background area, the part of [0.5-0.75] is a difficult defect area, and the part of [0.75-1] is a simple defect area.

[0055] In the embodiment, for each instance of each image, the feature C output by the backbone network of the student instance segmentation network is a tensor of [c, h, w], where c, h and w respectively represent the channel dimension, length and width of the feature C.

[0056] In the embodiment, the four parts m1-m4 of the mask and the feature C output by the backbone network of the student instance segmentation network are subjected to tensor multiplication, and then global average pooling is performed to obtain vectors v1-v4 of [c, 1, 1], that is, the feature centers of the instances; the category of the feature center is the category predicted for the instance, and the calculation expression of the vectors v1-v4 is:

[0057] v i =Global(m i *C),i=1,2,3,4;

[0058] wherein Global is a global average pooling operation, and the average value is calculated in the [h, w] dimension of the tensor of [c, h, w] to obtain a vector v i of [c, 1, 1].

[0059] In the embodiment, the feature centers of each instance are collected into different memory banks according to K defect categories and 4 mask region categories, and in the embodiment, the memory banks have a total of 4K;

[0060] The instance centers of the memory banks of the defect categories and the mask region categories are positive samples, and the instance centers of the memory banks of other defect categories and other mask region categories are negative samples, and a contrast learning loss function is constructed, and the expression of the contrast learning loss function L is:

[0061]

[0062] wherein sim is a similarity calculation function, sim(z i ,z j )=z i T *z j , and the calculation method is the inner product of two vectors; τ is a temperature coefficient; n is the number of positive samples, and N is the total number of positive samples and negative samples; z j is the positive sample, and z ipositive samples of the instance center of the memory bank of the same defect category and the category of the mask region, z k is the instance center of all memory banks.

[0063] S6), input the mask of each defect instance output by the teacher instance segmentation network and the feature output by the backbone network of the student instance segmentation network into the general distribution fusion module, obtain the fused pseudo label, and train the student instance segmentation network with the pseudo label;

[0064] In this embodiment, the feature C output by the backbone network of the student instance segmentation network of each image and the mask of the defect instance of the unannotated image generated by the teacher instance segmentation network are input into the general distribution fusion module, the mask of each defect instance and the feature C are multiplied pixel by pixel, and the feature [m, c, h, w] covered by each defect instance mask is obtained.

[0065] The mean μ and the variance σ of each dimension of the feature center in each memory bank are calculated, and then the probability γ that all features C covered by the defect instance mask belong to the Gaussian distribution of different memory banks is calculated according to the mean μ and the variance σ in the memory bank, that is:

[0066]

[0067] In the formula, v is the feature center; μ k is the mean of the kth instance; σ k is the variance of the kth instance.

[0068] The probability P[m, 4*K, h, w] obtained is taken as the maximum value in the second dimension to obtain a tensor of [m, h, w], which is used as a new mask, wherein m is the total number of memory banks.

[0069] The mean of the new mask and the old mask is calculated as an enhanced new mask, and the student instance segmentation network is trained using the enhanced new mask.

[0070] S7), the parameters of the teacher instance segmentation network are updated in a momentum iterative update manner using the parameters of the student instance segmentation network. In each iteration, the corresponding parameters p s of the student instance segmentation network are combined with the parameters p t of the teacher instance segmentation network to update the parameters p

[0071] p t = 0.999p t + 0.001p s .

[0072] The foregoing embodiments and description of the application only illustrate the principle and the best mode of the application, and various changes and modifications can be made to the application without departing from the spirit and scope of the application, and all these changes and modifications fall within the scope of the application.

Claims

1. A semi-supervised surface defect instance segmentation method based on pseudo label enhancement, characterized in that, Comprising the following steps: S1), obtaining surface defect image data set; and performing instance segmentation format annotation on part of the defect image; S2), training an initial instance segmentation network with a small amount of labeled defect image data; and using the parameters of the initial network to initialize the parameters of the student instance segmentation network and the teacher instance segmentation network; specifically comprising the following steps: S21), input the labeled defect image data into the instance segmentation network, after the defect image passes through the feature extractor and the decoder of the instance segmentation network, finally output the prediction of each instance of the defect image, including the class prediction class of the instance and the mask prediction of the instance ; S22)、according to the class of the true instance in the label of the defect image data and the mask , calculate the loss function of the predicted class and the mask; Wherein, the prediction class loss function is a cross-entropy loss function; the loss function of the mask is a mean square error loss function: cross-entropy loss function is: ; mean squared error loss function is: ; S3), input the labeled defect image data into the student instance segmentation network, combine the output of the student instance segmentation network and the annotation of the defect image data, and train the student instance segmentation network; S4), input the unlabeled defect image data into the teacher instance segmentation network to obtain the mask of each defect instance output by the teacher instance segmentation network; at the same time, input the unlabeled defect image data into the student instance segmentation network to obtain the feature output by the backbone network of the student instance segmentation network; S5), input the mask of each defect instance output by the teacher instance segmentation network and the feature output by the backbone network of the student instance segmentation network into the cross-supervision contrast learning module; construct a contrast learning loss function; The mask value of the defect instance ranges from 0 to 1, and the mask of the defect instance is divided into four parts according to the mask value The four parts are as follows: The part of the simple background area, The part of the difficult background area, The part of the difficult defect area, The part of the simple defect area; for each instance of each image, the feature output by the backbone network of the student instance segmentation network are tensors of shape ] where respectively represent the channel dimension, the length and the width of the feature ​ By dividing the mask into four parts m 1- m 4 and the backbone network output of the student instance segmentation network Tensor multiplication is performed, and global average pooling is performed to obtain a vector ,1,1] , that is, the feature center of the instance; the category of the feature center is the category predicted for the instance, and the vector The calculation expression of is: ； wherein, is a global average pooling operation, which computes the mean of the tensor ] over all dimensions ] to obtain a vector c ,1,1] ; S6), input the mask of each defect instance output by the teacher instance segmentation network and the feature output by the backbone network of the student instance segmentation network into the general distribution fusion module to obtain the fused pseudo label, and use the pseudo label to train the student instance segmentation network; S7), update the parameters of the teacher instance segmentation network by using the parameters of the student instance segmentation network through momentum iterative update.

2. The pseudo-labeling enhancement based semi-supervised surface defect instance segmentation method according to claim 1, characterized in that: In step S4), the unlabeled defect image is input into the teacher instance segmentation network, and after the defect image passes through the feature extractor and the decoder of the teacher instance segmentation network, the final output of each instance of the image is predicted, including the class prediction of the instance and the mask prediction of the instance.

3. The pseudo-labeling enhancement based semi-supervised surface defect instance segmentation method according to claim 1, characterized in that: In step S5), for K The system calculates feature centers for each instance across four defect categories and four mask regions, storing them in separate memory banks. Instance centers belonging to the same defect category and mask region category are considered positive samples, while instance centers belonging to other defect categories and other mask regions are considered negative samples. A contrastive learning loss function is then constructed. The expression is: wherein, is a similarity calculation function, = T is the inner product of two vectors; is a temperature coefficient; n is the number of positive samples, N is the total number of positive and negative samples; is a positive sample of the same defect category and the instance center of the memory bank of the mask area category, is a positive sample of the same defect category and the instance center of the memory bank of the mask area category, is the instance center of all memory banks.

4. The pseudo-labeling enhancement based semi-supervised surface defect instance segmentation method according to claim 1, characterized in that: In step S6, the feature C output by the backbone network of the student instance segmentation network of each image and the mask of the defect instance of the unlabeled image generated by the teacher instance segmentation network are input into the general distribution fusion module, and the mask of each defect instance and the feature C are multiplied pixel by pixel to obtain the feature covered by each defect instance mask [ ].

5. The pseudo-labeling enhancement based semi-supervised surface defect instance segmentation method according to claim 4, characterized in that: In step S6), the mean value of each dimension of the feature center in each memory bank is calculated and the variance ; then, according to the mean value and the variance of each memory bank, the probability that all features C covered by the defect instance belong to the Gaussian distribution of different memory banks is calculated , that is: ； wherein is the feature center; is the mean of the k th instance; is the variance of the k th instance; The resulting probability [ m ,4* K , h , w ] is maximized in the second dimension, resulting in a tensor of [ m , h , w ] as a new mask, where m is the total number of memory banks; Calculate the mean of the new mask and the old mask as the enhanced new mask, and use the enhanced new mask to train the student instance segmentation network.

6. The pseudo-labeling enhancement based semi-supervised surface defect instance segmentation method according to claim 1, characterized in that: In step S7, the parameters of the teacher instance segmentation network are updated using the parameters of the student instance segmentation network in a manner of momentum iterative update, and in each iteration, the corresponding parameters of the student instance segmentation network are combined , the parameters of the teacher instance segmentation network are updated using the parameters of the student instance segmentation network in a manner of momentum iterative update, and in each iteration, the corresponding parameters of the student instance segmentation network are combined The following operations are performed for updating: 。

Images