Small sample defect segmentation method and system based on prompt learning

Abstract

The present invention provides a small sample defect segmentation method and system based on prompt learning, which relates to the field of image recognition technology, including obtaining a defect image to be tested, a known defect image and a defect mask map of the known defect image; extracting multi-dimensional features of the defect image to be tested and the known defect image respectively; obtaining global similarity features and local similarity features of the defect image to be tested and the known defect image according to the multi-dimensional features of the defect image to be tested and the known defect image, merging the global similarity features and the local similarity features to obtain aggregated similarity features; extracting dense prompt vectors and sparse prompt vectors respectively based on the aggregated similarity features; inputting the dense prompt vectors and the sparse prompt vectors and the multi-dimensional features of the defect image to be tested into a mask decoder to obtain a defect segmentation result of the defect image to be tested. The present invention can improve the accuracy of small sample defect segmentation.

Description

Technical Field

[0001] The present invention relates to the technical field of image recognition, and in particular to a small sample defect segmentation method and system based on prompt learning. Background Art

[0002] In industrial production, deep learning-based defect segmentation is the main defect segmentation solution currently used. Its basic idea is to collect a large number of defect images from industrial scenes, annotate them, and train a fully supervised model to achieve defect segmentation. However, in industrial scenes, a large amount of defect data is often difficult to collect and annotate, and traditional defect segmentation methods can only identify known defect categories and are difficult to achieve generalized segmentation of unknown defect categories.

[0003] Small sample segmentation can use a small number of labeled known defect images to compare features of defect images of the same category, and use regions with similar features as the final segmentation results, thereby achieving generalized defect segmentation under limited samples. However, since the feature representation forms are often complex and diverse and easily affected by noise, the accuracy of small sample segmentation is limited. Summary of the invention

[0004] In order to address the deficiencies of the prior art, the object of the present invention is to provide a small sample defect segmentation method and system based on prompt learning to achieve high-precision small sample defect segmentation.

[0005] To achieve the above object, according to some embodiments, a first aspect of the present invention provides a small sample defect segmentation method based on prompt learning, comprising:

[0006] Obtaining a defect image to be tested, a known defect image, and a defect mask image of the known defect image;

[0007] Extracting multi-dimensional features of the defect image to be tested and the known defect image respectively;

[0008] According to the multi-dimensional features of the defect image to be tested and the known defect image, the global similarity features and the local similarity features of the defect image to be tested and the known defect image are obtained, and the global similarity features and the local similarity features are merged to obtain the aggregated similarity features;

[0009] Based on the aggregated similarity features, dense hint vectors and sparse hint vectors are extracted respectively;

[0010] The dense hint vector, the sparse hint vector and the multi-dimensional features of the defect image to be tested are input into the mask decoder to obtain the defect segmentation result of the defect image to be tested.

[0011] A second aspect of the present invention provides a small sample defect segmentation system based on prompt learning, comprising:

[0012] An image acquisition module is configured to acquire a defect image to be tested, a known defect image, and a defect mask image of the known defect image;

[0013] A multi-dimensional feature extraction module is configured to extract multi-dimensional features of the defect image to be tested and the known defect image respectively;

[0014] A similarity feature generation module is configured to obtain global similarity features and local similarity features of the defect image to be tested and the known defect image according to the multi-dimensional features of the defect image to be tested and the known defect image, and merge the global similarity features and the local similarity features to obtain an aggregated similarity feature;

[0015] a hint vector generation module, configured to extract a dense hint vector and a sparse hint vector respectively based on the aggregated similarity features;

[0016] The defect segmentation module is configured to input the dense hint vector and the sparse hint vector as well as the multi-dimensional features of the defect image to be tested into the mask decoder to obtain the defect segmentation result of the defect image to be tested.

[0017] According to a third aspect of the present invention, an electronic device is provided, comprising a memory, a processor and a computer program stored in the memory, wherein the processor executes the computer program to complete the steps of the above-mentioned small sample defect segmentation method based on prompt learning.

[0018] According to a fourth aspect of the present invention, a computer-readable storage medium is provided for storing computer instructions, which, when executed by a processor, complete the steps of the above-mentioned small sample defect segmentation method based on prompt learning.

[0019] A fifth aspect of the present invention provides a computer program product, including a computer program / instruction, which, when executed by a processor, implements the steps of the above-mentioned small sample defect segmentation method based on prompt learning.

[0020] Compared with the prior art, the present invention has the following beneficial effects:

[0021] The present invention provides a small sample defect segmentation method and system based on prompt learning. With the powerful generalization segmentation ability of the SAM model, the automatic small sample industrial defect segmentation without manual prompts is realized, and the accurate segmentation of untrained defects can be realized through a very small number of known defect samples. By extracting the multi-dimensional features of the input image, based on these features, the semantic features of the known defects and the defects to be tested are compared to obtain the similarity of the defects, and the defect similarity is used to generate dense embedded prompt vectors and sparse embedded prompt vectors as the mask decoder input of SAM, guiding SAM to achieve autonomous small sample segmentation.

[0022] Advantages of additional aspects of the present invention will be given in part in the following description, and in part will become obvious from the following description, or will be learned through practice of the present invention. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] The accompanying drawings in the specification, which constitute a part of the present invention, are used to provide a further understanding of the present invention. The exemplary embodiments of the present invention and their descriptions are used to explain the present invention and do not constitute improper limitations on the present invention.

[0024] Figure 1 Schematic diagram of the method in Embodiment 1 of the present invention;

[0025] Figure 2 It is the algorithm flow chart of the present invention;

[0026] Figure 3 This is a schematic diagram of the SAM architecture;

[0027] Figure 4 Schematic diagram of multi-dimensional feature extraction;

[0028] Figure 5 It is a schematic diagram of the dynamic prototype learning module;

[0029] Figure 6 Schematic diagram of the cosine similarity attention module;

[0030] Figure 7 Schematic diagram of dense embedding generation module;

[0031] Figure 8 Schematic diagram of the sparse embedding generation module. DETAILED DESCRIPTION

[0032] The present invention will be further described below in conjunction with the accompanying drawings and embodiments.

[0033] Embodiment 1

[0034] The core goal of small sample industrial defect segmentation is to use a very small number of annotated known defect images to locate defects in the image to be tested at the pixel level. It should be noted that in the model inference stage, the defect categories detected have never been seen in the training stage. Considering that the basic model has the ability to obtain generalized recognition from a large amount of data, in the field of image segmentation, Segment Anything Model (SAM) can achieve general interactive segmentation performance with its powerful data engine. This rich prior knowledge can greatly make up for the difficulty of generalization in small sample segmentation. Figure 3The figure is a schematic diagram of the traditional SAM architecture. The SAM model is an interactive segmentation framework based on large-scale data training, which achieves target segmentation through prompts such as points, boxes, and masks. The model mainly consists of three parts: image encoder, prompt encoder, and mask decoder. The image encoder mainly uses the Vision Transformer structure in pre-training as the model backbone to extract semantic features from the input image; the prompt encoder encodes the interactive conditional prompts and provides them to the mask decoder; the mask decoder mainly uses a two-layer Transformer decoding unit to interact with the acquired image features and prompt encoding information to generate the final segmentation result.

[0035] However, the traditional SAM model must rely on artificial prompts such as points, boxes, and masks to achieve segmentation tasks, and it is difficult to perform fully automatic segmentation of industrial defects, which undoubtedly greatly limits its segmentation capabilities for industrial defects.

[0036] Based on this, the embodiment of the present invention makes use of the powerful generalization segmentation capability of the SAM model and improves the model to achieve automated small sample industrial defect segmentation.

[0037] like Figure 1-Figure 2 As shown, Embodiment 1 of the present invention provides a small sample defect segmentation method based on prompt learning, including:

[0038] S1. Obtain a defect image to be tested, a known defect image, and a defect mask image of the known defect image;

[0039] S2, extracting multi-dimensional features of the defect image to be tested and the known defect image respectively;

[0040] S3, according to the multi-dimensional features of the defect image to be tested and the known defect image, obtain the global similarity features and the local similarity features of the defect image to be tested and the known defect image, merge the global similarity features and the local similarity features, and obtain the aggregated similarity features;

[0041] S4, based on the aggregated similarity features, extract dense hint vectors and sparse hint vectors respectively;

[0042] S5. Input the dense hint vector, the sparse hint vector and the multi-dimensional features of the defect image to be tested into a mask decoder to obtain a defect segmentation result of the defect image to be tested.

[0043] In step S2, in order to fully exploit the features of the image encoder, a multi-scale feature extraction strategy is used to extract rich image features. Specifically, Figure 4As shown, the features output by different network positions of the image encoder, i.e., the features output by the shallow network layer, the middle network layer, and the deep network layer, are obtained, thereby obtaining a multi-scale feature set of the known defect image and the defect image to be tested. Then, the convolution downsampling operation is used to adjust the features of different scales to the same size. The multi-scale feature aggregation is performed using the convolution channel merging operation to generate the multi-scale features of the known defect image and the defect image to be tested. Finally, the position mask map of the known defect image is used to perform a dot multiplication operation with the multi-scale features of the known defect image to eliminate the interference of the background in the known defect features. Based on the above operations, the multi-dimensional features of the known defect image and the defect image to be tested can be obtained, wherein the multi-dimensional features of the known defect image are the features obtained by the dot multiplication operation of the position mask map of the known defect image and the multi-scale features of the known defect image, and the multi-dimensional features of the defect image to be tested are the multi-scale features of the defect image to be tested.

[0044] In step S3, after obtaining the multi-dimensional features of the known defect image and the defect image to be tested, the features of the two need to be compared, so as to use the target features in the known defect to match the parts with similar features in the defect to be tested, so as to locate the defect. In order to achieve this goal, the dynamic prototype learning module and the cosine similarity attention module are used to obtain the aggregated similarity features.

[0045] Dynamic prototype learning module: In order to eliminate the interference of small sample data noise and extract more robust global similarity, the dynamic prototype learning module uses distribution matching to compare the known defect features with the defect features to be tested, such as Figure 5 First, the multi-dimensional features of the known defect image are subjected to mask average pooling operation for feature compression to obtain the known defect prototype features. Then, the mean μ of the known defect prototype features is calculated. m and variance Probability distribution of known defect prototype features can be constructed In order to increase the noise resistance of the distribution, learnable parameters w1 and w2 are introduced as weights of the mean and variance of the probability distribution, respectively, to optimize the probability distribution of known defect prototype features. Through automatic learning, the probability distribution can be made more accurate, and then the global similarity feature is obtained by calculating the cosine similarity between the distribution and the multi-dimensional features of the defect image to be tested.

[0046] Cosine similarity attention module: It obtains local similarity features through dense feature matching, such as Figure 6First, two fully connected networks are used to map the multi-dimensional features of known defect images into the Key (K) vector and Value (V) vector of the attention mechanism, and another fully connected network maps the multi-dimensional features of the defect image to be tested into the Query (Q) vector of the attention mechanism. Based on the above settings, the cosine similarity attention module is defined as follows:

[0047]

[0048] Where d represents the dimension of the vector. Different from the common attention mechanism, cosine similarity attention achieves a denser similarity measurement between features by constructing the cosine similarity of Q and K to obtain local similar features.

[0049] Based on the above operations, global similarity features and local similarity features can be obtained, and then the channel merging operation can be used to obtain the aggregated similarity features. The aggregated similarity features combine the global similarity features obtained by the dynamic prototype learning module and the local similarity features obtained by the cosine similarity attention module, and use the aggregated similarity features as input for the subsequent generation of dense prompt vectors and sparse prompt vectors.

[0050] In step S4, based on the aggregated similarity features, the location of the target defect in the defect image to be tested is prompted, which can guide the model to achieve accurate defect segmentation. To achieve this goal, a dense embedding generation module and a sparse embedding generation module are constructed respectively to extract embedding prompt features of different dimensions to guide the model to perform targeted segmentation.

[0051] Dense embedding generation module: As can be seen from the above, the aggregate similarity feature is mainly obtained through global and local matching of known defect features and defect features to be tested. Since the aggregate similarity feature only merges two features through channel merging operation, but does not establish a connection between the two features, in order to achieve further fusion prompts, the dense embedding generation module is designed to make full use of these features to generate dense prompt vectors. Figure 7 As shown in the figure, a single convolution operation is first performed on the aggregated similarity features to achieve the fusion of global and local feature information by increasing the dimension. Then, a self-attention operation is used to further associate the feature context information. Finally, a single convolution operation is used again to reduce the feature dimension to generate the final dense prompt vector.

[0052] Sparse embedding generation module: This part uses a learnable query vector to obtain the implicit representation of the target to generate a sparse embedding hint vector to achieve target positioning. Figure 8As shown in the figure, first, a set of learnable feature vectors is randomly generated, which can be optimized by gradient descent. Then, the feature vector is cross-attended with the multi-dimensional features of the defect image to be tested to generate a preliminary feature vector to obtain the information of the image to be tested. Then, the preliminary feature vector is cross-attended with the global similarity features and the local similarity features to generate two sets of optimized feature vectors. Finally, the two sets of feature vectors are fused through a feedforward neural network to generate a sparse prompt vector.

[0053] In step S5, after obtaining the dense hint vector and the sparse hint vector, the dense hint vector and the sparse hint vector, as well as the multi-dimensional features of the defect image to be tested, are used as the input of the mask decoder to prompt the model to implement the segmentation task, thereby completing the entire small sample automatic segmentation process.

[0054] In order to achieve accurate small sample industrial defect segmentation, this embodiment provides a small sample defect segmentation method based on prompt learning. This method uses the powerful generalization segmentation ability of the SAM model to achieve automated small sample industrial defect segmentation and uses a very small number of known defect samples to achieve accurate segmentation of untrained defects. First, a multi-scale fusion strategy is designed to obtain the multi-scale semantic features of the input defect image from the encoder of SAM. Based on these features, the dynamic prototype learning module and the cosine similarity attention module are used to compare the semantic features of known defects and defects to be tested to obtain the similarity of defects. Then, the dense and sparse embedding generation modules are designed to use the defect similarity to generate dense embedding prompt vectors and sparse embedding prompt vectors as the mask decoder input of SAM to guide SAM to achieve autonomous small sample segmentation. This method makes full use of the general segmentation ability of the visual base model and shows better defect segmentation performance compared with existing segmentation methods.

[0055] Embodiment 2

[0056] This embodiment provides a small sample defect segmentation system based on prompt learning, including:

[0057] An image acquisition module is configured to acquire a defect image to be tested, a known defect image, and a defect mask image of the known defect image;

[0058] A multi-dimensional feature extraction module is configured to extract multi-dimensional features of the defect image to be tested and the known defect image respectively;

[0059] A similarity feature generation module is configured to obtain global similarity features and local similarity features of the defect image to be tested and the known defect image according to the multi-dimensional features of the defect image to be tested and the known defect image, and merge the global similarity features and the local similarity features to obtain an aggregated similarity feature;

[0060] a hint vector generation module, configured to extract a dense hint vector and a sparse hint vector respectively based on the aggregated similarity features;

[0061] The defect segmentation module is configured to input the dense hint vector and the sparse hint vector as well as the multi-dimensional features of the defect image to be tested into the mask decoder to obtain the defect segmentation result of the defect image to be tested.

[0062] It should be noted here that the various modules in this embodiment correspond one-to-one to the steps of the method in Example 1, and the specific implementation process is the same, which will not be repeated here.

[0063] Embodiment 3

[0064] This embodiment provides an electronic device, including a memory, a processor, and a computer program stored in the memory, wherein the processor executes the computer program to complete the steps of the method in the first embodiment.

[0065] Embodiment 4

[0066] This embodiment provides a computer-readable storage medium for storing computer instructions. When the computer instructions are executed by a processor, the steps of the method in the first embodiment are completed.

[0067] Embodiment 5

[0068] This embodiment provides a computer program product, including a computer program / instruction, characterized in that the computer program / instruction implements the steps of the method in the first embodiment when executed by a processor.

[0069] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. For those skilled in the art, the present invention may have various modifications and variations. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present invention shall be included in the protection scope of the present invention.

Claims

1. A small sample defect segmentation method based on hint learning, characterized in that: include: Obtaining a defect image to be tested, a known defect image, and a defect mask image of the known defect image; Extracting multi-dimensional features of the defect image to be tested and the known defect image respectively; According to the multi-dimensional features of the defect image to be tested and the known defect image, the global similarity features and the local similarity features of the defect image to be tested and the known defect image are obtained, and the global similarity features and the local similarity features are merged to obtain the aggregated similarity features; Based on the aggregated similarity features, dense hint vectors and sparse hint vectors are extracted respectively; The process of obtaining a dense hint vector includes: The aggregated similarity features are dimensionally increased through a single convolution operation; Associating feature context information through self-attention operation; The dense cue vector is obtained by reducing the dimension through a single convolution operation; The process of obtaining the sparse hint vector includes: Perform a cross-attention operation on the multi-dimensional features of the defect image to be tested and the randomly generated learnable feature vector to generate a preliminary feature vector; Perform cross-attention operations on the preliminary feature vector, the global similarity feature and the local similarity feature to obtain the first feature vector and the second feature vector; Fusing the first eigenvector and the second eigenvector to obtain a sparse hint vector; The dense hint vector, the sparse hint vector and the multi-dimensional features of the defect image to be tested are input into the mask decoder to obtain the defect segmentation result of the defect image to be tested.

2. The small sample defect segmentation method based on hint learning according to claim 1, characterized in that: The process of acquiring the multi-dimensional features of the known defect image includes: Extract multi-scale features of known defect images; The position mask map of the known defect image is multiplied by the multi-scale features of the known defect image to obtain the multi-dimensional features of the known defect image.

3. The small sample defect segmentation method based on prompt learning according to claim 1, characterized in that: Obtain the global similarity features of the defect image to be tested and the known defect image, including: Obtain known defect prototype features through mask average pooling operation; Construct a probability distribution of known defect prototype features; The cosine similarity between the probability distribution of known defect prototype features and the multi-dimensional features of the defect image to be tested is calculated to obtain the global similarity features.

4. The small sample defect segmentation method based on prompt learning according to claim 1, characterized in that: The local similarity features of the defect image to be tested and the known defect image are obtained by using cosine similarity attention, and the calculation formula is: Among them, the K vector and the V vector are obtained according to the multi-dimensional feature mapping of the known defect image, the Q vector is obtained according to the multi-dimensional feature mapping of the defect image to be tested, and d represents the dimension of the vector.

5. A small sample defect segmentation system based on prompt learning, characterized in that: include: An image acquisition module is configured to acquire an image of a defect to be tested and an image of a known defect; A multi-dimensional feature extraction module is configured to extract multi-dimensional features of the defect image to be tested and the known defect image respectively; A similarity feature generation module is configured to obtain global similarity features and local similarity features of the defect image to be tested and the known defect image according to the multi-dimensional features of the defect image to be tested and the known defect image, and merge the global similarity features and the local similarity features to obtain an aggregated similarity feature; a hint vector generation module, configured to extract a dense hint vector and a sparse hint vector respectively based on the aggregated similarity features; The process of obtaining a dense hint vector includes: The aggregated similarity features are dimensionally increased through a single convolution operation; Associating feature context information through self-attention operation; The dense cue vector is obtained by reducing the dimension through a single convolution operation; The process of obtaining the sparse hint vector includes: Perform a cross-attention operation on the multi-dimensional features of the defect image to be tested and the randomly generated learnable feature vector to generate a preliminary feature vector; Perform cross-attention operations on the preliminary feature vector, the global similarity feature and the local similarity feature to obtain the first feature vector and the second feature vector; Fusing the first eigenvector and the second eigenvector to obtain a sparse hint vector; The defect segmentation module is configured to input the dense hint vector and the sparse hint vector into the mask decoder to obtain a defect segmentation result of the defect image to be tested.

6. An electronic device, characterized in that: The invention comprises a memory, a processor and a computer program stored in the memory, wherein the processor executes the computer program to complete the steps of the method according to any one of claims 1 to 4.

7. A computer-readable storage medium, characterized in that: Used to store computer instructions, which, when executed by a processor, complete the steps of the method according to any one of claims 1 to 4.

8. A computer program product comprising a computer program / instructions, characterized in that When the computer program / instructions are executed by a processor, the steps of the method according to any one of claims 1 to 4 are implemented.

Images