Few-sample defect detection method based on prototype prompt fine-tuning visual basic model SAM

Abstract

The invention belongs to the technical field of image recognition, and discloses a few-sample defect detection method based on a prototype prompt fine-tuning visual basic model SAM. Prototype prior information about an image to be segmented is extracted from a small number of example image-mask pairs, so that prior guidance of human beings is replaced; for defect data, a representative foreground prompt point, a high-confidence background prompt point (sparse prompt) and a foreground Mask (dense prompt) are extracted from prototype prior to serve as prototype prompt, and a visual basic model SAM is guided to achieve more accurate segmentation. A visual basic model SAM is finely adjusted by using prototype prompt and a small amount of defect data, a prompt encoder and an image encoder of the SAM are frozen, an LoRA bypass is injected into the image encoder, and the SAM is enabled to better adapt to defect data distribution by updating parameters of the LoRA bypass and a lightweight mask decoder.

Description

Technical Field

[0001] The present invention relates to the field of image recognition, and in particular to a few-sample defect detection method based on prototype-cue-based fine-tuning of a visual base model (SAM). Background Art

[0002] Industrial defect segmentation plays a vital role in ensuring the quality and safety of the production process. Visual defect segmentation is a non-contact, non-destructive automatic detection technology that can effectively control the equipment precisely. Deep learning has significantly improved the performance of this task, but it relies on a large amount of labeled data. In some industrial scenarios, there is less defect data, and it is difficult to form a large-scale data set, which will cause the performance of the model to drop sharply.

[0003] Few-shot learning methods can alleviate the scarcity problem of industrial defect data to a certain extent. However, since they are only trained on a small amount of defect data in a specific field, the limited and single training data leads to limited feature extraction capabilities of few-shot defect segmentation models, which hinders their promotion in other industrial environments.

[0004] The Visual Basic Model (SAM) has learned rich visual features through pre-training on large-scale general datasets, has strong performance and generalization ability, and can quickly adapt to new data distributions and downstream tasks through fine-tuning. However, SAM is an interactive visual basic model that relies on human prior knowledge to guide segmentation.

[0005] Therefore, we extract prior information about the image to be segmented through examples instead of human prior guidance. Focusing on defect data, we extract representative hints from the prior and use them together with a small amount of defect data to fine-tune the basic visual model SAM to adapt it to the distribution of defect data, thereby improving the performance and generalization of the algorithm. Summary of the invention

[0006] The purpose of this invention is to propose a few-shot defect detection method based on prototype-cue fine-tuning visual base model SAM, extracting prior information about the image to be segmented through examples, extracting representative cues from the prior information, and using the cues and a small amount of defect data to fine-tune the visual base model SAM to adapt it to the defect data distribution, thereby improving the performance and generalization of the algorithm on the few-shot defect segmentation task.

[0007] The technical solution of the present invention is as follows: a few-shot defect detection method based on prototype-cue fine-tuning visual base model SAM, extracting prototype prior information about the image to be segmented from a small number of example image-mask pairs, the prototype prior information replaces human prior guidance;

[0008] For defect image data, representative foreground cue points, high-confidence background cue points and foreground masks are extracted from prototype prior information as prototype cues; wherein the foreground cue points and background cue points are sparse cues, and the foreground mask is a dense cue;

[0009] The visual base model SAM is fine-tuned through prototype cues; the cue encoder and image encoder of SAM are frozen, and the LoRA bypass is injected into the image encoder. By updating the parameters of the LoRA bypass and mask decoder, the adjusted visual base model SAM is obtained for defect detection.

[0010] The method for extracting the prototype prior information is as follows: for a set of input support images and a query image, the corresponding support image features F are extracted through the CNN backbone network. S and query image features F Q , C represents the number of channels, H and W represent the height and width of the feature respectively; when there are multiple supporting images, F S Subdivided into The corresponding mask is subdivided into Subsequently, multiple prototypes are extracted from the supporting image features through mask average pooling MAP; the calculation of MAP is expressed as:

[0011]

[0012] where F S Indicates the supported image features, Y S represents the support image mask, c represents the category; after extracting multiple foreground prototypes and background prototypes from the support image features, the two are averaged to obtain the final foreground prototype P F and background prototype P B , Its calculation is expressed as:

[0013]

[0014]

[0015] Among them, Avg represents the average function, 1 represents the foreground category, and 0 represents the background category; calculate the background prototype P respectively B and the foreground prototype P F and query image features F Q The cosine similarity of Its calculation is expressed as:

[0016] S B =Cos(P B , F Q ) (4)

[0017] S F =Cos(P F ,F Q ) (5)

[0018] Cos represents the cosine similarity function;

[0019] These similarity graphs are concatenated and normalized through the exponential function to obtain the prototype prior probability graph. Its calculation is expressed as:

[0020] M BF =Softmax(Concat(S B , S F )) (6)

[0021] Among them, Softmax represents the normalized exponential function; M BF Decomposed into background probability map M B and the foreground probability map M F ,

[0022] The method for generating the prototype prompt is as follows: first, the prototype prior probability map M BF Upsample to the original size of the image; compare the predicted probability of the foreground and the predicted probability of the background at the same pixel position through Argmax, take the category with higher prediction probability as the predicted category, and generate a coarse mask Its calculation is expressed as:

[0023]

[0024] Among them, Upsample means upsampling;

[0025] Extract multiple foreground cue points and background cue points as sparse cues, and extract a processed rough mask as a dense cue; extract the rough mask Connected component A of the middle foreground area j , and calculate the area S of each connected component j , j∈{1,...,N}, N represents the number of connected components; at the same time, the maximum connected component area S is calculated max , its calculation is expressed as:

[0026] S max =Max{S1, S2, ..., S N} (8)

[0027] Among them, Max represents the maximum value function; set the threshold α to generate dense prompts Its calculation is expressed as:

[0028]

[0029] Through the thinning algorithm, dense hints Each connected component in is refined to a pixel width. These refined pixels are the skeleton of each connected component. The pixels on the skeleton are arranged in order and form a set This process is represented as:

[0030]

[0031] Skeletonize represents the Zhang-Suen thinning algorithm; when the total number of skeleton pixels is L, Traverse the collection for the step size And select 9 points as foreground prompt points, the process is as follows:

[0032]

[0033] Among them, step represents the step length, P positive Represents a set of foreground cue points;

[0034] From the background probability map M B 9 pixels are uniformly selected as candidate background cue points, and their set is defined as The threshold β is used to select points with high confidence as negative prompt points. The calculation is expressed as follows:

[0035]

[0036] Where P negative Represents a collection of background cue points;

[0037] Foreground cue points and background cue points together constitute sparse cues.

[0038] The strategy for fine-tuning the visual base model SAM with prototype prompts is designed as follows: freezing the parameters of the SAM image encoder and prompt encoder, and embedding the LoRA bypass for the image encoder. When fine-tuning the visual base model SAM, only the parameters of the LoRA bypass and the mask decoder are updated; the prototype prompt is encoded by the prompt encoder and input into the mask decoder; the query image is input into the image encoder embedded in the LoRA bypass to obtain the image embedding, which is input into the mask decoder; the mask decoder obtains the predicted mask; based on the comparison between the predicted mask and the true value, the parameters of the LoRA bypass and the mask decoder are further fine-tuned.

[0039] The image encoder contains a total of 12 Transformer blocks; the projection layer in the attention layer of each Transformer block A LoRA bypass consisting of a low-rank matrix is ​​added; W q , W k , W v Respectively represent the linear change matrices of the query vector Q, key vector K and value vector V, d and k represent their input dimension and output dimension respectively. The process is expressed as follows:

[0040] W0+ΔW=W0+BA (13)

[0041] Among them, the pre-training weight W0 represents any W q , W k or W v , ΔW represents the incremental parameter matrix, and Represents a low-rank matrix with rank r<<min(d, k); at the beginning of training, A is initialized with a random Gaussian value distribution and B is initialized with a zero matrix to ensure that BA=0, thereby maintaining consistency with the initial state.

[0042] Beneficial effects of the present invention:

[0043] (1) The present invention uses limited defect data to fine-tune the visual basic model SAM and applies it to few-sample defect segmentation, which improves the segmentation performance of the model and its generalization ability in unknown industrial scenarios.

[0044] (2) The present invention extracts foreground prototypes and background prototypes from a small number of example image-mask pairs, respectively, and uses these prototypes to calculate similarities with query image features, thereby obtaining a prototype-based foreground / background prior probability map, thereby replacing human prior guidance.

[0045] (3) Aiming at defect data, the present invention proposes a prototype hint generation method, which extracts representative hint points and masks from the prototype prior as prototype hints to guide the visual basic model SAM to achieve more accurate segmentation.

[0046] (4) The present invention designs a fine-tuning strategy for the visual base model SAM, which uses prototype hints and a small amount of defect data to fine-tune SAM, freezes the hint encoder and image encoder of SAM, and injects LoRA bypass into the image encoder. By updating the parameters of LoRA bypass and lightweight mask decoder, they can better adapt to the distribution of defect data. BRIEF DESCRIPTION OF THE DRAWINGS

[0047] Figure 1 A flowchart of a few-sample defect detection method based on prototype-cue fine-tuning visual base model SAM provided in this embodiment;

[0048] Figure 2 The overall network framework diagram provided for this embodiment;

[0049] Figure 3 A schematic diagram of the prototype prompt extraction provided for this embodiment.

[0050] Figure 4 This is a schematic diagram of the LoRA bypass provided in this embodiment. DETAILED DESCRIPTION

[0051] This example provides a few-shot defect detection method based on prototype-cued fine-tuning of the visual base model SAM. Figure 1 and Figure 2 As shown, the following steps are included:

[0052] Step 1: Input a query image and a set of support image-mask pairs;

[0053] Step 2: For a set of input support images F S and a query image F Q , use the CNN backbone network to extract the corresponding support image features F S and query image features F Q , Where C represents the number of channels, H and W represent the height and width of the feature map respectively. When there are multiple supporting images, F S Can be divided into The corresponding mask is also subdivided into Subsequently, multiple prototypes are extracted from the supporting image features respectively through mask average pooling (MAP). The calculation of MAP is expressed as:

[0054]

[0055] where F S Indicates the supported image features, Y S represents the support image mask, and c represents the category. After extracting multiple foreground prototypes and background prototypes from the support image features, we average them to obtain the final foreground prototype and background prototype Its calculation is expressed as:

[0056]

[0057]

[0058] Avg represents the averaging function, 1 represents the foreground class, and 0 represents the background class. Then, we calculate the cosine similarity between the background prototype and the foreground prototype and the query image features, and obtain the corresponding similarity graph Its calculation is expressed as:

[0059] S B =Cos(P B,F Q ) (4)

[0060] S F =Cos(P F ,F Q ) (5)

[0061] Where Cos represents the cosine similarity function.

[0062] These similarity graphs are then concatenated and normalized through the exponential function to obtain the prototype prior probability graph. Its calculation is expressed as:

[0063] M BF =Softmax(Concat(S B , S F )) (6)

[0064] Where Softmax represents the normalized exponential function. BF It can be further decomposed into the background probability map M B and the foreground probability map MF,

[0065] Step 3: Convert the prototype prior probability map M BF Upsample to the original size of the image. Then, by using Argmax to compare the predicted probability of the foreground and the predicted probability of the background at the same pixel position, the higher probability is selected as the predicted category and a coarse mask is generated. Its calculation is expressed as:

[0066]

[0067] Among them, Upsample means upsampling.

[0068] We then extract multiple foreground cue points and background cue points as sparse cues and extract a processed coarse mask as dense cue. First, we extract Connected component A of the middle foreground area j , and calculate the area of ​​each connected component as S j . Where j∈{1,...,N}, N represents the number of connected components. At the same time, the maximum connected component area S is calculated max , its calculation is expressed as:

[0069] S max =Max{S1, S2, ..., S N} (8)

[0070] Among them, Max represents the maximum value function. Then, the threshold α is used to eliminate small segmentation fragments to generate dense prompts. Its calculation is expressed as:

[0071]

[0072] In order to extract representative cue points, we use a thinning algorithm to transform dense cues into Each connected component in is thinned to a pixel width, and these thinned pixels are the skeleton of each connected component. These pixels on the skeleton are arranged in order and form a set This process can be expressed as:

[0073]

[0074] Skeletonize represents the Zhang-Suen thinning algorithm. When the total number of skeleton pixels is L, we use Traverse the collection for the step size And select 9 points as foreground prompt points, the process is as follows:

[0075]

[0076] Among them, step represents the step length, P positive Represents a collection of foreground cue points.

[0077] Then, from the background probability map M B 9 pixels are evenly selected as candidate background prompt points, and their set is defined as The threshold β is used to select points with high confidence as negative prompt points. The calculation is expressed as follows:

[0078]

[0079] Where P negative Represents the set of background cue points. Foreground cue points and background cue points together form sparse cues.

[0080] Step 4: Freeze the parameters of the SAM image encoder and the hint encoder, and design a LoRA bypass for the image encoder. When fine-tuning the visual base model SAM on the defective image data, only the parameters of the LoRA bypass and the lightweight mask decoder are updated. The image encoder contains a total of 12 Transformer blocks. We use the projection layer in the attention layer in each Transformer block. A LoRA bypass consisting of a low-rank matrix is ​​designed. q , W k , W v Represent the linear change matrices of query vector Q, key vector K and value vector V respectively, d and k represent their input dimension and output dimension respectively. The process is expressed as follows:

[0081] W0+ΔW=W0+BA (13)

[0082] Among them, the pre-training weight W0 represents any W q , W k or W v , ΔW represents the incremental parameter matrix, and Denotes a low-rank matrix with rank r < min(d, k). At the beginning of training, we initialize A with a random Gaussian value distribution and B with a zero matrix to ensure that BA = 0, thereby maintaining consistency with the initial state.

[0083] This method is tested on MVTec-AD texture defects. In order to verify the generalization ability of the model obtained by this method on unseen classes and unseen industrial scenes, it is divided into MVTec_Seen and MVTec_Unseen according to different data partitions. MVTec_Seen is trained and tested on the same materials, but uses unseen defect categories in seen materials during the test process. When one pair of supporting image-mask pairs (One-Shot) is provided, it can achieve an average intersection ratio of 63.07%, and when five pairs of supporting image-mask pairs (Five-Shot) are provided, it can achieve an average intersection ratio of 66.48%. MVTec_Unseen is trained and tested on different materials. When one pair of supporting image-mask pairs is provided, it can achieve an average intersection ratio of 53.51%, and when five pairs of supporting image-mask pairs are provided, it can achieve an average intersection ratio of 59.34%. This method does not rely on human prior guidance and has excellent performance and generalization ability. It can not only clearly segment the boundaries of various defects, but also adapt to different industrial scenarios.

[0084] Table 1. Performance comparison with different advanced methods on MVTec_Seen

[0085]

[0086] Table 2. Performance comparison with different advanced methods on MVTec_Unseen

[0087]

Claims

1. A few-shot defect detection method based on prototype-cued fine-tuning of visual base model SAM, characterized in that: Extracting prototype prior information about the image to be segmented from a small number of example image-mask pairs, the prototype prior information replaces human prior guidance; For defect image data, representative foreground cue points, high-confidence background cue points and foreground masks are extracted from prototype prior information as prototype cues; wherein the foreground cue points and background cue points are sparse cues, and the foreground mask is a dense cue; The visual base model SAM is fine-tuned through prototype cues; the cue encoder and image encoder of SAM are frozen, and the LoRA bypass is injected into the image encoder. By updating the parameters of the LoRA bypass and mask decoder, the adjusted visual base model SAM is obtained for defect detection.

2. The few-shot defect detection method based on prototype-cued fine-tuning visual base model (SAM) according to claim 1, characterized in that: The method for extracting the prototype prior information is as follows: for a set of input support images and a query image, the corresponding support image features F are extracted through the CNN backbone network. S and query image features C represents the number of channels, H and W represent the height and width of the feature respectively; when there are multiple supporting images, F S Subdivided into The corresponding mask is subdivided into Subsequently, multiple prototypes are extracted from the supporting image features through mask average pooling MAP; the calculation of MAP is expressed as: Among them, F S Indicates the supported image features, Y S represents the support image mask, c represents the category; after extracting multiple foreground prototypes and background prototypes from the support image features, the two are averaged to obtain the final foreground prototype P F and background prototype Its calculation is expressed as: Among them, Avg represents the average function, 1 represents the foreground category, and 0 represents the background category; calculate the background prototype P respectively B and the foreground prototype P F and query image features F Q The cosine similarity of Its calculation is expressed as: S B =Cos(P B ,F Q ) (4) S F =Cos(P F ,F Q ) (5) Cos represents the cosine similarity function; These similarity graphs are concatenated and normalized through the exponential function to obtain the prototype prior probability graph. Its calculation is expressed as: M BF =Softmax(Concat(S B ,S F )) (6) Among them, Softmax represents the normalized exponential function; M BF Decomposed into background probability map M B and foreground probability map 3. The method for few-sample defect detection based on prototype-cued fine-tuning visual base model SAM according to claim 2, characterized in that: The method for generating the prototype prompt is as follows: first, the prototype prior probability map M BF Upsample to the original size of the image; compare the predicted probability of the foreground and the predicted probability of the background at the same pixel position through Argmax, take the category with higher prediction probability as the predicted category, and generate a coarse mask Its calculation is expressed as: Among them, Upsample means upsampling; Extract multiple foreground cue points and background cue points as sparse cues, and extract a processed rough mask as a dense cue; extract the rough mask Connected component A of the middle foreground area j , and calculate the area S of each connected component j , j∈{1,...,N}, N represents the number of connected components; at the same time, the maximum connected component area S is calculated max , its calculation is expressed as: S max =Max{S1,S2,...,S N } (8) Among them, Max represents the maximum value function; set the threshold α to generate dense prompts Its calculation is expressed as: Through the thinning algorithm, dense hints Each connected component in is refined to a pixel width. These refined pixels are the skeleton of each connected component. The pixels on the skeleton are arranged in order and form a set This process is represented as: Skeletonize represents the Zhang-Suen thinning algorithm; when the total number of skeleton pixels is L, Traverse the collection for the step size And select 9 points as foreground prompt points, the process is as follows: Among them, step represents the step length, P positive Represents a set of foreground cue points; From the background probability map M B 9 pixels are uniformly selected as candidate background prompt points, and their set is defined as The threshold β is used to select high-confidence points as negative prompt points. The calculation is expressed as follows: Where P negative Represents a collection of background cue points; Foreground cue points and background cue points together constitute sparse cues.

4. The method for few-shot defect detection based on prototype-cued fine-tuned visual base model (SAM) according to claim 1, characterized in that: The strategy of fine-tuning the visual base model SAM with prototype prompts is designed as follows: freeze the parameters of the SAM image encoder and the prompt encoder, and embed the LoRA bypass into the image encoder. When fine-tuning the visual base model SAM, only the parameters of the LoRA bypass and the mask decoder are updated; the prototype prompt is encoded by the prompt encoder and input into the mask decoder; the query image is input into the image encoder embedded in the LoRA bypass to obtain the image embedding, which is input into the mask decoder; The mask decoder obtains a prediction mask; Based on the comparison between the predicted mask and the true value, continue to fine-tune the parameters of the LoRA bypass and mask decoder.

5. The method for few-shot defect detection based on prototype-cued fine-tuning visual base model (SAM) according to claim 4, characterized in that: The image encoder contains a total of 12 Transformer blocks; the projection layer in the attention layer of each Transformer block A LoRA bypass consisting of a low-rank matrix is ​​added; W q , W k , W v Respectively represent the linear change matrices of the query vector Q, key vector K and value vector V, d and k represent their input dimension and output dimension respectively. The process is expressed as follows: W0+ΔW=W0+BA (13) Among them, the pre-training weight W0 represents any W q , W k or W v , ΔW represents the incremental parameter matrix, and Represents a low-rank matrix with rank r<<min(d, k); at the beginning of training, A is initialized with a random Gaussian value distribution and B is initialized with a zero matrix to ensure that BA=0, thereby maintaining consistency with the initial state.

Images