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

Abstract

The invention provides a semi-supervised surface defect instance segmentation method based on pseudo label enhancement. The method comprises the following steps: acquiring a surface defect image data set; training an initial instance segmentation network by using a small amount of labeled defect image data; inputting the marked defect image data into a student instance segmentation network to train the student instance segmentation network; inputting the defect image data without labels into the teacher instance segmentation network and the student instance segmentation network, and respectively obtaining a mask and features of each defect instance; inputting the data into a cross-supervision comparative learning module and a general distribution fusion module; constructing a contrast learning loss function, obtaining a fused pseudo tag, and training a student instance segmentation network by using the pseudo tag; and updating the parameters of the teacher instance segmentation network by using the parameters of the student instance segmentation network. According to the method, challenges of small data set scale and defect boundary fuzziness can be handled, complex category feature distribution can be estimated through the general distribution fusion module, and the quality of mask false labels from a teacher model is improved.

Description

Technical Field

[0001] The present invention relates to the technical field of defect detection, and in particular to a semi-supervised surface defect instance segmentation method based on pseudo-label enhancement. Background Art

[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, many 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 that the defect data is uniformly distributed and the labeled data is sufficient. At the same time, some researchers are committed to estimating the fuzzy edges of defect areas and propose edge detection methods based on traditional filtering techniques, Gaussian functions, and improvements. However, obtaining labeled defect data is still a difficult, expensive, and time-consuming task, which makes vision-based industrial inspection challenging. In addition, directly locating the edge does not make full use of the information of the fuzzy area, which to some extent affects the performance of the model for defect detection.

[0003] Semi-supervised defect detection has attracted increasing attention. The existing technology has enhanced the defect segmentation ability of the model by using generative adversarial networks (GANs) to synthesize data, and improved the model performance by integrating multi-source information. In addition, the existing technology constructs a simple semi-supervised teacher-student model framework. The student model is trained using pseudo labels predicted by the teacher model, 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 performance of the framework is improved by reducing the difference in prediction distribution between the student model and the teacher model or the difference between different input scales. However, these studies mainly focus on classification, detection, and segmentation tasks, ignoring instance segmentation. Instance segmentation can not only accurately locate defects and provide detailed information, but also minimize the impact of background pixels on annotation. In addition, the method of generating pseudo labels based on thresholds is prone to overfitting of the model, 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, which can not only accurately locate defects and provide detailed information, but also minimize the impact of background pixels on annotation. In addition, the threshold-based pseudo-label generation method is prone to overfitting of the model, making it difficult to accurately identify defects in unlabeled images.

[0005] Secondly, the pseudo-labels generated by existing methods are inaccurate. Since the pseudo-labels are automatically generated by the model rather than based on manual annotation, their accuracy is difficult to guarantee. When the accuracy of the pseudo-labels is not high, the model will learn these inaccurate pseudo-labels as true labels in the subsequent iterative training process, gradually accumulating noise. As the training progresses, the accumulation of noise may cause the model's judgment of defects to deviate more and more from the actual situation, thereby reducing the overall performance of the model. This phenomenon is particularly evident in the long-term iteration process. Furthermore, the edges of defects are usually blurred, especially for smaller defects, which makes the existing semi-supervised surface defect detection methods easily mispredict the edges of defects, seriously affecting their performance. Summary of the invention

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

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

[0008] S1), obtaining a surface defect image dataset; and performing instance segmentation format annotation on some defect images;

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

[0010] S3), inputting the labeled defect image data into the student instance segmentation network, combining the output of the student instance segmentation network and the labeling of the defect image data to train the student instance segmentation network;

[0011] S4), inputting the unlabeled defect image data into the teacher instance segmentation network, and obtaining the mask of each defect instance output by the teacher instance segmentation network;

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

[0013] S5), inputting 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 a cross-supervised contrastive learning module; constructing a contrastive learning loss function;

[0014] S6), inputting 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, obtaining the fused pseudo labels, and training the student instance segmentation network with the pseudo labels;

[0015] S7) Use the parameters of the student instance segmentation network to update the parameters of the teacher instance segmentation network through momentum iterative update.

[0016] Preferably, in step S2), a small amount of labeled image data is used to train an initial instance segmentation network, which specifically includes the following steps:

[0017] S21), input the labeled defect image data into the instance segmentation network, and after the defect image passes through the feature extractor and decoder of the instance segmentation network, the prediction of each instance of the defect image is finally output, including the category prediction category of the instance And the mask prediction of the instance

[0018] S22), according to the category y of the real instance in the annotation of the defect image data i and mask m i , calculate the loss function of the predicted category and mask; among them, the predicted category loss function is the cross entropy loss function; the mask loss function is the mean square error loss function:

[0019] Cross entropy loss function L CE for:

[0020] Mean square error loss function L MSE for:

[0021] Preferably, in step S4), the unlabeled defect image is output to the teacher instance segmentation network. After the image passes through the feature extractor and decoder of the classroom network, the image finally outputs the prediction of each instance, including the category prediction of the instance and the mask prediction of the instance.

[0022] Preferably, 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, namely: (0-0.25] is the simple background area, [0.25-0.5] is the difficult background area, [0.5-0.75] is the difficult defect area, and [0.75-1] is the simple defect area.

[0023] Preferably, 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, w represent the channel dimension, length and width of feature C, respectively.

[0024] Preferably, in step S5), the four parts m of the mask are 1 -m 4Tensor multiplication is performed with the feature C output by the backbone network of the student instance segmentation network, and then global average pooling is performed to obtain the vector v of [c, 1, 1] 1 -v 4 , which is the feature center of the instance; the category of the feature center is the category predicted for the instance, and the vector v 1 -v 4 The calculation expression is:

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

[0026] Among them, Global is a global average pooling operation, which calculates the mean of the tensor [c,h,w] in the [h,w] dimension at the same time to obtain the vector v of [c,1,1] i .

[0027] Preferably, in step S5), for K defect categories and 4 mask area categories, the calculated feature center of each instance is collected into different memory banks, the instance centers of the memory banks of the same defect category and mask area category are positive samples, and the instance centers of the memory banks of other defect categories and other mask area categories are negative samples, and a contrastive learning loss function is constructed.

[0028] Preferably, in step S5), the expression of the contrastive learning loss function L constructed is:

[0029]

[0030] Among them, sim is the similarity calculation function, sim(z i ,z j )=z i T *z j , calculated as the inner product of two vectors; τ is the temperature coefficient; n is the number of positive samples, N is the total number of positive and negative samples; z j For z i Positive samples of the instance centers of the memory bank of the same defect category and mask region category, z k The instance center for all memories.

[0031] Preferably, 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 and feature C of each defect instance are multiplied pixel by pixel to obtain the feature [m, c, h, w] covered by the mask of each defect instance.

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

[0033]

[0034] Where v is the feature center; μ k is the mean of the kth instance; σ k The variance of the kth instance.

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

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

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

[0038] 1. The present invention can cope with the challenges of small data set size and blurred defect boundaries, and the general distribution fusion module of the present invention can estimate complex category feature distributions and improve the quality of pseudo-labels from teacher model masks;

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

[0040] Figure 1 It is a schematic diagram of the process of the present invention;

[0041] Figure 2 It is a flow chart of the method of the present invention; DETAILED DESCRIPTION

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

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

[0044] S1), obtaining a surface defect image dataset; and performing instance segmentation format annotation on some defect images;

[0045] S2), using a small amount of labeled defect image data to train an initial instance segmentation network; 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), input the labeled defect image data into the instance segmentation network, and after the defect image passes through the feature extractor and decoder of the instance segmentation network, the prediction of each instance of the defect image is finally output, including the category prediction category of the instance And the mask prediction of the instance

[0047] S22), according to the category y of the real instance in the annotation of the defect image data i and mask m i , calculate the loss function of the predicted category and mask; among them, the predicted category loss function is the cross entropy loss function; the mask loss function is the mean square error loss function:

[0048] Cross entropy loss function L CE for:

[0049] Mean square error loss function L MSE for:

[0050] S3), inputting the labeled defect image data into the student instance segmentation network, combining the output of the student instance segmentation network and the labeling of the defect image data to train the student instance segmentation network;

[0051] S4), input the unlabeled defect image data into the teacher instance segmentation network, and obtain the mask of each defect instance output by the teacher instance segmentation network; this embodiment outputs the unlabeled defect image to the teacher instance segmentation network, and after the image passes through the feature extractor and decoder of the classroom network, it finally outputs the prediction of each instance of the image, including the category prediction of the instance, and the mask prediction of the instance.

[0052] At the same time, the unlabeled defect image data is input into the student instance segmentation network to obtain the features 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 features output by the backbone network of the student instance segmentation network into a cross-supervised contrastive learning module; constructing a contrastive learning loss function;

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

[0055] In this 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, w represent the channel dimension, length and width of feature C respectively.

[0056] In this embodiment, the four parts m of the mask are 1 -m 4 Tensor multiplication is performed with the feature C output by the backbone network of the student instance segmentation network, and then global average pooling is performed to obtain the vector v of [c, 1, 1] 1 -v 4 , which is the feature center of the instance; the category of the feature center is the category predicted for the instance, and the vector v 1 -v 4 The calculation expression is:

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

[0058] Among them, Global is a global average pooling operation, which calculates the mean of the tensor [c,h,w] in the [h,w] dimension at the same time to obtain the vector v of [c,1,1] i .

[0059] In this embodiment, for K defect categories and 4 mask area categories, the feature center of each instance calculated is collected into different memory banks. In this embodiment, there are 4K memory banks in total.

[0060] The instance centers of the memory libraries of the same defect category and mask area category are positive samples, and the instance centers of the memory libraries of other defect categories and other mask area categories are negative samples. A contrastive learning loss function is constructed. The expression of the contrastive learning loss function L is:

[0061]

[0062] Among them, sim is the similarity calculation function, sim(z i ,z j )=z i T *z j, calculated as the inner product of two vectors; τ is the temperature coefficient; n is the number of positive samples, N is the total number of positive and negative samples; z j For z i The positive samples of the instance centers of the memory bank of the same defect category and the category of the mask region, z k The instance center for all memories.

[0063] S6), inputting 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, obtaining the fused pseudo labels, and training the student instance segmentation network with the pseudo labels;

[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 unlabeled image generated by the teacher instance segmentation network are input into the general distribution fusion module, and the mask and feature C of each defect instance are multiplied pixel by pixel to obtain the feature [m, c, h, w] covered by the mask of each defect instance.

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

[0066]

[0067] Where v is the feature center; μ k is the mean of the kth instance; σ k The variance of the kth instance.

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

[0069] The mean of the new mask and the old mask is calculated as the enhanced new mask, and the enhanced new mask is used to train the student instance segmentation network.

[0070] S7) Use the parameters of the student instance segmentation network to update the parameters of the teacher instance segmentation network through momentum iteration. In each iteration, the corresponding parameters p of the student instance segmentation network are combined s , the parameter p of the teacher instance segmentation network t Do the following to update:

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

[0072] The above embodiments and descriptions are only for illustrating the principles and best embodiments of the present invention. Without departing from the spirit and scope of the present invention, the present invention may be subject to various changes and improvements, all of which fall within the scope of the present invention to be protected.

Claims

1. A semi-supervised surface defect instance segmentation method based on pseudo-label enhancement, characterized in that: The steps include: S1), obtaining a surface defect image dataset; and performing instance segmentation format annotation on some defect images; S2), using a small amount of labeled defect image data to train an initial instance segmentation network; and using the parameters of the initial network to initialize the parameters of the student instance segmentation network and the teacher instance segmentation network; S3), inputting the labeled defect image data into the student instance segmentation network, combining the output of the student instance segmentation network and the labeling of the defect image data to train the student instance segmentation network; 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; at the same time, inputting 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; S5), inputting 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 a cross-supervised contrastive learning module; constructing a contrastive learning loss function; S6), inputting 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, obtaining the fused pseudo labels, and training the student instance segmentation network with the pseudo labels; S7) Use the parameters of the student instance segmentation network to update the parameters of the teacher instance segmentation network through momentum iterative update.

2. The method for semi-supervised surface defect instance segmentation based on pseudo-label enhancement according to claim 1, characterized in that: In step S2), an initial instance segmentation network is trained using a small amount of labeled image data, specifically including the following steps: S21), input the labeled defect image data into the instance segmentation network, and after the defect image passes through the feature extractor and decoder of the instance segmentation network, the prediction of each instance of the defect image is finally output, including the category prediction category of the instance And the mask prediction of the instance S22), according to the category y of the real instance in the annotation of the defect image data i and mask m i , calculate the loss function of the predicted category and mask; Among them, the prediction category loss function is the cross entropy loss function; the mask loss function is the mean square error loss function: Cross entropy loss function L CE for: Mean square error loss function L MSE for:

3. The method for semi-supervised surface defect instance segmentation based on pseudo-label enhancement according to claim 1, characterized in that: In step S4), the unlabeled defect image is output to the teacher instance segmentation network. After the defect image passes through the feature extractor and decoder of the classroom network, the image prediction for each instance is finally output, including the category prediction of the instance and the mask prediction of the instance.

4. The method for semi-supervised surface defect instance segmentation based on pseudo-label enhancement according to claim 1, characterized in that: 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, namely: (0-0.25] is the simple background area, [0.25-0.5] is the difficult background area, [0.5-0.75] is the difficult defect area, and [0.75-1] is the simple defect area.

5. The method for semi-supervised surface defect instance segmentation based on pseudo-label enhancement according to claim 4, characterized in that: 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], where c, h, w represent the channel dimension, length and width of feature C respectively.

6. The method for semi-supervised surface defect instance segmentation based on pseudo-label enhancement according to claim 5, characterized in that: In step S5), the four parts m1-m4 of the mask are tensor-multiplied with the feature C output by the backbone network of the student instance segmentation network, and then global average pooling is performed to obtain the vector v1-v4 of [c, 1, 1], which is the feature center of the instance; the category of the feature center is the category predicted for the instance, and the calculation expression of the vector v1-v4 is: v i =Global(m i *C),i=1,2,3,4; Among them, Global is a global average pooling operation, which calculates the mean of the tensor [c,h,w] in the [h,w] dimension at the same time to obtain the vector v of [c,1,1] i .

7. The method for semi-supervised surface defect instance segmentation based on pseudo-label enhancement according to claim 6, characterized in that: In step S5), for the K defect categories and the 4 mask area categories, the feature center of each instance calculated is collected into different memory banks, the instance center of the memory bank of the same defect category and mask area category is the positive sample, and the instance center of the memory bank of other defect categories and other mask area categories is the negative sample, and a contrastive learning loss function is constructed; the expression of the contrastive learning loss function L is: Among them, sim is the similarity calculation function, sim(z i ,z j )=z i T *z j , calculated as the inner product of two vectors; τ is the temperature coefficient; n is the number of positive samples, N is the total number of positive and negative samples; z j For z i The positive samples of the instance centers of the memory bank of the same defect category and the category of the mask region, z k The instance center for all memories.

8. The method for semi-supervised surface defect instance segmentation based on pseudo-label enhancement 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 and feature C of each defect instance are multiplied pixel by pixel to obtain the feature [m, c, h, w] covered by the mask of each defect instance.

9. The method for semi-supervised surface defect instance segmentation based on pseudo-label enhancement according to claim 8, characterized in that: In step S6), the mean μ and variance σ of each dimension of the feature center in each memory bank are calculated; then, based on the mean μ and variance σ in the memory bank, the probability γ that all features C covered by the defect instance mask belong to the Gaussian distribution of different memory banks is calculated, that is: Where v is the feature center; μ k is the mean of the kth instance; σ k The variance of the kth instance; Take the maximum value of the obtained probability γ[m,4*K,h,w] in the second dimension to obtain a tensor of [m,h,w] as the new mask, where m is the total number of memory banks; The mean of the new mask and the old mask is calculated as the enhanced new mask, and the enhanced new mask is used to train the student instance segmentation network.

10. The semi-supervised surface defect instance segmentation method based on pseudo-label enhancement according to claim 1, characterized in that: In step S7), the parameters of the teacher instance segmentation network are updated by using the parameters of the student instance segmentation network through momentum iterative update. In each iteration, the corresponding parameters p of the student instance segmentation network are combined. s , the parameter p of the teacher instance segmentation network t Do the following to update: p t =0.999p t +0.001p s 。

Images