Small sample defect detection method based on a prototype learning guided discriminative segmentation network

Abstract

The application belongs to the technical field of small sample defect detection, and discloses a small sample defect detection method based on a prototype learning guided discriminative segmentation network, which comprises the following steps: obtaining a group of normal small sample images and defect small sample images; using a trained neural network to extract features in the normal small sample images and the defect small sample images to obtain initial features of the small samples; inputting the initial features of the small samples into a prototype generation network to obtain embedded feature blocks, wherein a loss function of the prototype generation network comprises a normal feature clustering loss function, a defect feature separation loss function and a classification loss function based on pixel-level differences; inputting a similarity map and the embedded feature blocks into a discriminative segmentation network after channel dimension splicing, and realizing defect segmentation and identification of the small samples. The application can realize more fine and accurate defect detection and positioning.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of small sample defect detection, and more particularly to a small sample defect detection method based on a prototype learning guided discriminative segmentation network. BACKGROUND

[0002] In order to improve the reliability and pass rate of product quality, surface defect detection in industrial production process becomes particularly important. Traditional surface defect detection mainly relies on manual experience, and quality inspection personnel need to visually inspect the product. However, due to visual fatigue, environmental factors and different personnel's judgment of defects, the detection efficiency and quality are difficult to guarantee. In order to solve these problems, machine vision technology for industrial defect detection has become a common technology. Machine vision is a technology that uses computer vision systems to simulate human visual function, and has wide application potential in industrial production.

[0003] The machine vision system converts the product surface image into a digital signal through optical equipment such as a camera, and then analyzes and recognizes the image through image processing algorithms. These algorithms can be trained and optimized through deep learning and artificial intelligence to identify defects and flaws on the surface of the product.

[0004] Image defect detection aims to distinguish between normal samples and defective samples. Unlike traditional object detection methods, there is no fixed defect category in defect detection, and any object that is different from the normal sample can be considered a defect. Existing object detection methods are not suitable for handling such a wide variety of defects, and the number of defect samples in actual production is limited, making it difficult to obtain a sufficient number of defect samples for accurate detection. Therefore, the current method relies heavily on normal samples for defect positioning. However, manually collecting a large number of normal samples in real-world scenarios requires a lot of time and cost. Therefore, research on small sample defect detection is urgently needed.

[0005] One potential research direction to solve the problem of small sample defect detection is prototype feature learning, which extracts the features of a support image as prototype features. However, comparing the features of a query image with prototype features, where similar parts are considered normal and locations with significant differences are identified as differences. Compared with directly comparing pixel features, this method of comparing with prototype features is more robust to noise. However, in the process of constructing prototype features, there is a problem that the lack of defect features will cause the normal feature space to become loose, which will cause the following problems:

[0006] (1) The recognition ability of the normal feature space decreases: When only normal samples are used for training in the network, it will be biased towards the main features observed in the normal samples. Therefore, it may not be able to capture the complexity of the defect samples, thereby ignoring the unique features that may exist therein. This can result in poor discrimination of the feature space, and the network is difficult to find the subtle differences between normal and defects.

[0007] (2) Unable to obtain fine-grained defect localization: When the defect area is very small or only a slight semantic difference exists in the normal prototype, directly comparing the extracted features with the prototype features may not accurately identify the defect area.

[0008] (3) Difficulty in identifying diversified defects: Given the heterogeneity of defect features, there are complex challenges in identifying defects, which can exhibit diverse and unique features. Without specific knowledge of their uniqueness, it is difficult to accurately detect defects.

[0009] Therefore, the current defect detection algorithm based on prototype learning is prone to problems such as loose normal feature space, rough defect localization results, and inability to detect multiple types of defects, and there is an urgent need to provide a method for small sample defect detection in industry. SUMMARY

[0010] In view of the above defects or improvement needs of the prior art, the present application provides a small sample defect detection method based on prototype learning guided discriminative segmentation network, which can realize more fine and accurate defect detection and localization.

[0011] To achieve the above object, according to one aspect of the present application, a small sample defect detection method based on a prototype learning guided discriminative segmentation network is provided, comprising: S1: obtaining a group of normal small sample images and defect small sample images; S2: extracting features in the normal small sample images and defect small sample images by using a trained neural network to obtain initial features of the small samples; S3: inputting the initial features of the small samples into a prototype generation network to obtain an embedded feature block, a loss function of the prototype generation network comprising a normal feature clustering loss function, a defect feature separation loss function and a pixel-level difference based classification loss function, wherein the pixel-level difference based classification loss function is created by the following way: calculating the average value of the normal features in the initial features, taking the average value as the initial center of the prototype features of the small samples, performing nearest neighbor search on the normal features with the initial center as the center to obtain k prototype features of the small samples, calculating the distance between the k prototype features and the embedded feature block to obtain a similarity map, obtaining the defect degree by averaging the similarity map at the pixel level, obtaining the distance between the defect degree and the small sample label, and then obtaining the pixel-level difference based classification loss function; S4: inputting the similarity map and the embedded feature block after channel dimension splicing into a discriminative segmentation network to realize defect segmentation and identification of the small samples; S5: training the prototype generation network and the discriminative segmentation network by using a training set to obtain a trained prototype generation network and a trained discriminative segmentation network, and inputting a to-be-tested small sample into the trained prototype generation network and the trained discriminative segmentation network to realize defect detection of the small sample.

[0012] Preferably, the prototype generation network comprises, in sequence, a convolutional layer, a regularization layer, a Relu activation function, an average pooling layer and two CoordConv2d layers.

[0013] Preferably, the normal feature clustering loss function L NFC is:

[0014]

[0015] wherein T = H x W, H is the height of the features extracted by the trained neural network, W is the width of the features extracted by the trained neural network, K is the total number of prototype features, F t is the embedded feature block, is the k prototype features closest to the initial center, d(.,.) is a predefined distance measure, and r is the range radius close to the initial center.

[0016] Preferably, the defect feature separation loss function L AFS is:

[0017] ​

[0018] where T = H x W, H is the height of the extracted feature of the trained neural network, W is the width of the extracted feature of the trained neural network, K is the total number of prototype features, F t is the embedding feature block, is the k prototype features closest to the initial center, d(.,.) is a predefined distance metric, r is the range radius close to the initial center, and a is used to adjust the relative importance of the normal feature clustering loss function and the defect feature separation loss function.

[0019] Preferably, the classification loss function L class based on pixel-level difference is:

[0020] L class = L(g s (F t ), y x )

[0021]

[0022]

[0023] where the prototype features are compared with the proposed normal features and defect features to obtain an image patch similarity map, P is the top P pixels of the image patch similarity map, is the similarity map, F t is the embedding feature block, K is the total number of prototype features, d(.,.) is a predefined distance metric, is the k prototype features closest to the initial center, y x represents the label assigned to the small sample x.

[0024] Preferably, the discriminative segmentation network is FCN, U-net or DSN.

[0025] Preferably, the segmentation loss function L seg of the discriminative segmentation network is:

[0026]

[0027] where M i and are the defect segmentation result output by the discriminative segmentation network and the true value label of the defect, i represents individual pixels in the entire pixel set I.

[0028] Preferably, in step S1, a set of normal small sample images and defect small sample images are obtained, specifically: a set of normal small sample images are randomly obtained from the training set, and defect samples are constructed based on the normal small sample images according to the prior knowledge of negative samples.

[0029] Overall, compared with the prior art, the small sample defect detection method based on prototype learning guided discriminative segmentation network provided by the present application mainly has the following beneficial effects:

[0030] 1. The normal feature clustering loss function can realize dense clustering of normal features, and the prototype feature set is obtained based on the clustered model, then the embedding feature block is embedded in the prototype feature set, the update of the prototype generation network is controlled through the distance between the embedding feature block and the prototype feature, and whether the embedding feature exceeds the pre-defined radius range is determined by clustering the normal features, thereby realizing the differentiation of normal and defect instances; the defect feature separation loss function can solve the problem of spatial mixing of normal features and defect features when there is large inter-class variation in normal samples. The loss function further separates the defect and the prototype feature by introducing synthetic defect samples, thereby further compressing the normal feature space. This enhancement method improves the discrimination ability of the integrated feature extractor, effectively separates normal and defect instances in the feature space, and reduces the possibility of misclassifying normal samples as defects; further, the present application provides a classification loss function based on pixel-level difference, which focuses on the pixels in the subtle defect area to maximize the dissimilarity between small size defects and normal samples.

[0031] 2. When using the discriminative segmentation network to realize fine-grained defect detection, in order to emphasize the area where defects may occur, the present application combines the possibility map with the obtained feature embedding. Then the spliced features are input into the discriminative segmentation network for effective defect segmentation. Through the training of the segmentation network, the network obtains the understanding ability of various forms of defect features, and the supervision signal provided by the label enables the segmentation network to realize more fine and accurate defect detection and positioning.

[0032] 3. The prototype generation network sequentially includes a convolution layer, a regularization layer, a Relu activation function, an average pooling layer, and two CoordConv2d layers. The convolution layer, the regularization layer, and the Relu activation function form a module to compress the channel number of the extracted features; then the average pooling operation is adopted to unify the size of the three feature layers extracted by the pre-trained network, so as to facilitate the alignment of the channel dimension splicing; finally, through the coordinate convolution CoordConv2d, the coordinate information is added to the input feature map, so that the network can better utilize the position information of the pixels. BRIEF DESCRIPTION OF DRAWINGS

[0033] Figure 1 is a step diagram of the small sample defect detection method based on prototype learning guided discriminative segmentation network;

[0034] Figure 2is a small sample defect detection method flowchart based on a prototype learning guided discriminative segmentation network.

[0035] Figure 3 is a prototype learning process schematic diagram of the present application. DETAILED DESCRIPTION

[0036] In order to make the objectives, technical solutions and advantages of the present application clearer, the present application is further described in detail below in combination with the drawings and examples. It should be understood that the specific examples described herein are only used to explain the present application and do not limit the present application. In addition, the technical features involved in each embodiment of the present application described below can be combined with each other as long as they do not conflict with each other.

[0037] The present application provides a small sample defect detection method based on a prototype learning guided discriminative segmentation network, as shown in Figure 1 and Figure 2 , the method comprises the following steps S1-S5.

[0038] S1: Obtain a group of normal small sample images and defect small sample images.

[0039] A group of normal small sample images x n is randomly selected from the training set. n According to the prior knowledge of negative samples, defect samples x a are constructed based on the normal small sample images x n .

[0040] S2: The initial features of the small samples are obtained by extracting the features in the normal small sample images and the defect small sample images using the trained neural network.

[0041] The present application first uses a pre-trained neural network E to extract relevant features, and the parameters of the network remain unchanged during the entire training process. The pre-trained neural network E can be ResNet 18, ResNet 50, ResNet101, VGGNet, MobileNet or Inception.

[0042] The initial features of the small samples are obtained by extracting the features in the normal small sample images and the defect small sample images using the trained neural network.

[0043] S3: The initial features of the small samples are input into the prototype generation network to obtain embedded feature blocks, and the loss function of the prototype generation network includes a normal feature clustering loss function, a defect feature separation loss function, and a classification loss function based on pixel-level differences.

[0044] The application further introduces a prototype generation network G with learnable parameters, the output of the pre-trained neural network is the input of the prototype generation network G, so the pre-trained neural network E and the prototype generation network G constitute an integrated feature extractor. The reason for using the pre-trained neural network E and the prototype generation network G at the same time also includes that the features extracted by the pre-trained neural network tend to be biased, by introducing another feature extraction network prototype generation network G, the features can be guided into the target domain.

[0045] The normal small sample image and the defect small sample image are input into the integrated feature extractor to obtain a target-oriented embedded feature block F t =G(E(x)), where t∈{1,2,…H×W}, H and W represent the height and width of the features extracted by E respectively.

[0046] The prototype generation network of the application sequentially includes a convolutional layer, a regularization layer, a Relu activation function, an average pooling layer, and two CoordConv2d layers.

[0047] A convolutional layer, a regularization layer, and a Relu activation function form a module to compress the channel number of the extracted features. Then, an average pooling operation is used to unify the size of the three feature layers extracted by the pre-trained network, so as to facilitate the alignment of the channel dimension splicing. Finally, the coordinate convolution CoordConv2d is used to introduce the coordinate information in the convolutional neural network. The traditional convolutional layer only considers the channel information of the input feature map, and ignores the spatial position information of the pixels. The CoordConv2d adds coordinate information to the input feature map, so that the network can better utilize the position information of the pixels.

[0048] The application is better for identifying defects in small samples, such as Figure 3 As shown in the figure, the normal feature clustering loss function, the defect feature separation loss function, and the pixel-level difference-based classification loss function are used to train the prototype generation network during the training process of the prototype generation network.

[0049] Specifically, the normal feature clustering loss function aims to find a hypersphere with the smallest possible radius to achieve dense clustering of normal features. First, the application calculates the average value of the normal features in the initial features to obtain an initial center as a prototype feature. Then, the nearest neighbor search is performed on the normal features to find k nearest prototype features, denoted as Finally, the embedded feature block F t is embedded in the range close to The range is determined by a predefined radius r, and the radius r is a threshold for determining whether the embedded feature block is considered as a defect. The normal feature clustering loss function L NFC of the process is:

[0050]

[0051] where T = H x W, H is the height of the extracted feature of the trained neural network, W is the width of the extracted feature of the trained neural network, K is the total number of prototype features, F t is the embedding feature block, is the k prototype features closest to the initial center, d(.,.) is a predefined distance metric, r is the range radius close to .

[0052] The significance of setting the defect feature separation loss function is that when there is a large inter-class difference in normal samples, there may be unoccupied space between normal features and prototype features, therefore, the normal feature space may be excessively expanded, so that some defect samples similar to normal instances can also exist therein, therefore, using the defect feature, by making them away from the prototype feature, the normal space is further compressed, and the discrimination ability of the integrated feature extractor is further enhanced, the defect feature separation loss function L AFS is defined as:

[0053]

[0054] where T = H x W, H is the height of the extracted feature of the trained neural network, W is the width of the extracted feature of the trained neural network, K is the total number of prototype features, F t is the embedding feature block, is the k prototype features closest to the initial center, d(.,.) is a predefined distance metric, r is the range radius close to .

[0055] The significance of setting the classification loss function based on the pixel-level difference lies in that some defect samples tend to exhibit very similar features to normal samples, for example, smaller defect regions or strong semantic similarity between defect regions and normal features. Traditional contrast learning methods often have difficulty in effectively distinguishing these samples. Therefore, the present application proposes a classification loss function based on pixel-level difference, which emphasizes training pixels with greater differences in defect regions compared to normal samples. By focusing on these specific pixels rather than overall features, the differences between these pixels and normal features are amplified. Specifically, the distance between k prototype features and the embedded feature block is calculated to obtain a similarity map, the similarity map is averaged at the pixel level to obtain a defect degree, the distance between the defect degree and the small sample label is obtained, and then the classification loss function based on the pixel-level difference is obtained. Specifically, the classification loss function L class is:

[0056] L class = L(g s (F t ), y x )

[0057]

[0058]

[0059] where the image patch similarity map P is obtained by comparing the prototype features with the proposed normal features and defect features, P is the top P pixels of the image patch similarity map, is the similarity map, F t is the embedded feature block, K is the total number of prototype features, d(.,.) is a predefined distance measure, is the k prototype features closest to the initial center, y x represents the label assigned to the small sample x, if x is a defect small sample, then y x = 1; if x is a normal sample, then y x = 0, p x , p y indicates the specific location of pixel p, g s (F t ) quantitatively evaluates the degree of defect present in the image by calculating the average of the selected pixels, L(·,·) represents the cross-entropy loss function.

[0060] S4: The similarity map and the embedded feature block are concatenated in the channel dimension and input into the discriminative segmentation network to perform defect segmentation of the small sample.

[0061] The discriminative segmentation network is defined as S, which uses a pyramid-like top-down structure to introduce multi-scale information. This information fusion allows the network to simultaneously utilize the captured spatial relationships and the rich semantic information in F t for accurate defect segmentation. The discriminative segmentation network can be FCN, U-net, DSN, etc.

[0062] The discriminative segmentation network mainly includes:

[0063] ResNet backbone network: the model uses a ResNet backbone network for feature extraction, which can choose a network of 50, 101 or 152 layers as the basis, and the network level is customized, including atrous convolution and modified down-sampling to improve segmentation performance.

[0064] Classification branch: composed of a series of convolutional layers, ReLU activation and dropout. It is used to predict the class label of the segmentation region.

[0065] Feature enhancement module: the model introduces a feature enhancement module (FEM) for context modeling. It includes pyramid pooling, correlation convolution, beta convolution and internal classification branch. FEM enhances the feature representation capability and captures multi-scale context information.

[0066] Alpha convolution: the model contains Alpha convolution, which is used to fuse features from different pyramid levels to improve the final feature representation.

[0067] ASPP (Atrous Spatial Pyramid Pooling): the model optionally contains an ASPP module for context modeling. It performs atrous convolution at multiple sampling rates to capture multi-scale context information.

[0068] Decoder: the model has a decoder module to improve feature representation and generate the final segmentation output, which may include additional convolutional layers.

[0069] The input of the discriminative segmentation network is created by concatenating the similarity map and the embedding feature F t in the channel dimension, which provides key information for accurate defect segmentation. Specifically, represents the degree of similarity between the embedding feature block and the corresponding prototype feature, which covers the spatial relationship. On the other hand, the embedding feature block F t contains rich semantic information related to the input image, which enables the segmentation network to capture different visual clues and patterns.

[0070] The discriminative segmentation network adopts a segmentation loss function L seg, for generating a pixel-level defect score map representing the likelihood of each pixel being classified as a defect. To enhance the network's understanding of defect features, the present application always provides pixel-level ground truth label guidance during the training phase, enabling the network to grasp the unique attributes of defects. The segmentation loss function is defined as follows:

[0071]

[0072] where M i and is the defect segmentation result output by the discriminative segmentation network, and the ground truth label of the defect, i represents individual pixels in the entire pixel set I.

[0073] Therefore, the total loss term of the prototype learning guided discriminative segmentation network based on the present application is:

[0074] L total = L NFC + L AFS + L class + L seg

[0075] S5: training the prototype generation network and the discriminative segmentation network using the training set, obtaining the trained prototype generation network and the discriminative segmentation network, and inputting the small sample to be tested into the trained prototype generation network and the discriminative segmentation network to realize small sample defect detection.

[0076] The prototype learning guided discriminative segmentation network is trained using the training set, and the trained prototype learning guided discriminative segmentation network is obtained. Inputting the small sample to be tested into the trained prototype learning guided discriminative segmentation network can realize small sample defect detection.

[0077] In summary, the small sample defect detection method of the present application mainly includes a prototype learning stage and a discriminative segmentation stage. In the prototype learning stage, some new strategies are introduced to enhance the ability of the network to generate discriminative feature representations. The present application defines three types of features: normal features, artificially synthesized defect features, and prototype features. Then, the contrast learning technique is used to minimize the distance between the prototype features and the normal features, while ensuring clear separation from the synthetic defect features. In addition, a classification loss function based on pixel-level differences is introduced, which focuses on the pixels within the subtle defect regions to maximize the dissimilarity between small size defects and normal samples. This process produces discriminative prototype features and a likelihood map obtained by comparing the feature embeddings with the prototype features. Subsequently, the present application proposes a discriminative segmentation network to achieve fine-grained defect detection. In order to emphasize the areas where defects may occur, the present application combines the likelihood map with the obtained feature embeddings, and then inputs the spliced features into the discriminative segmentation network for effective defect segmentation. Through the training of the segmentation network, the network acquires the ability to understand various forms of defect features, and the supervision signal provided by the labels enables the segmentation network to achieve more fine and accurate defect detection and positioning.

[0078] Those skilled in the art will easily understand that the above description is only the preferred embodiment of the present application, and is not intended to limit the present application. Any modifications, equivalent replacements and improvements made within the spirit and principles of the present application shall be included in the protection scope of the present application.

Claims

1. A small-sample defect detection method based on a prototype learning-guided discriminative segmentation network, characterized in that, include: S1: Obtain a set of normal small sample images and defective small sample images; S2: Use the trained neural network to extract features from the normal small sample images and defective small sample images to obtain the initial features of the small samples; S3: Input the initial features of the small sample into the prototype generation network to obtain embedded feature blocks. The loss function of the prototype generation network includes a normal feature clustering loss function, a defect feature separation loss function, and a classification loss function based on pixel-level differences. The classification loss function based on pixel-level differences is created as follows: calculate the average value of the normal features in the initial features, use the average value as the initial center of the prototype features of the small sample, perform nearest neighbor search on the normal features with the initial center as the center to obtain k prototype features of the small sample, calculate the distance between the k prototype features and the embedded feature blocks to obtain a similarity map, calculate the average of the similarity map at the pixel level to obtain the defect degree, obtain the distance between the defect degree and the small sample label, and then obtain the classification loss function based on pixel-level differences. S4: The similarity map and the embedded feature block are concatenated along the channel dimension and then input into the discriminative segmentation network to achieve defect segmentation and recognition for small samples; S5: The prototype generation network and the discriminant segmentation network are trained using the training set to obtain the trained prototype generation network and discriminant segmentation network. The small sample to be tested can be input into the trained prototype generation network and discriminant segmentation network to realize the defect detection of the small sample.

2. The small sample defect detection method according to claim 1, characterized in that, The prototype generation network includes, in sequence: It consists of one convolutional layer, one regularization layer, one ReLU activation function layer, one average pooling layer, and two CoordConv2d layers.

3. The small sample defect detection method according to claim 1 or 2, characterized in that, The normal feature clustering loss function L NFC for: Where T = H × W, H is the height of the features extracted by the trained neural network, W is the width of the features extracted by the trained neural network, K is the total number of prototype features, and F t For the embedded feature block, For the k prototype features closest to the initial center, d(.,.) is a predefined distance metric, and r is the distance between the k features and the initial center. The radius of the approximate range.

4. The small sample defect detection method according to claim 1 or 2, characterized in that, Defect Feature Separation Loss Function L AFS for: Where T = H × W, H is the height of the features extracted by the trained neural network, W is the width of the features extracted by the trained neural network, K is the total number of prototype features, and F t For the embedded feature block, For the k prototype features closest to the initial center, d(.,.) is a predefined distance metric, and r is the distance between the k features and the initial center. The approximate radius, α, is used to adjust the relative importance of the normal feature clustering loss function and the defect feature separation loss function.

5. The small sample defect detection method according to claim 1 or 2, characterized in that, The classification loss function L based on pixel-level differences class for: L class =L(g s (F t ),y x ) The image patch similarity map is obtained by comparing the prototype features with the proposed normal and defective features, where P represents the first P pixels of the image patch similarity map. For similarity graphs, F t Let K be the total number of prototype features, and d(.,.) be a predefined distance metric. For the k prototype features closest to the initial center, y x This represents the label assigned to a small sample x.

6. The small sample defect detection method according to claim 1, characterized in that, The discriminant segmentation network is FCN, U-net, or DSN.

7. The small sample defect detection method according to claim 1 or 6, characterized in that, The segmentation loss function L of the discriminative segmentation network seg for: Among them, M i and To determine the defect segmentation result and the ground truth label of the defect output by the segmentation network, i represents an individual pixel in the entire pixel set I.

8. The small sample defect detection method according to claim 1, characterized in that, In step S1, a set of normal small sample images and defective small sample images are obtained. Specifically, a set of normal small sample images is randomly obtained from the training set, and defective samples are constructed based on the prior knowledge of negative samples.

Images