A few-shot defect detection method based on prototype-cued fine-tuning of visual anchor model (SAM)

Abstract

The present invention belongs to the field of image recognition technology and discloses a few-shot defect detection method based on prototype-cue fine-tuning of a visual base model (SAM). The method extracts prototype prior information about the image to be segmented from a small number of example image-mask pairs, replacing human prior guidance. Based on defect data, representative foreground cue points and high-confidence background cue points (sparse cues), as well as a foreground mask (dense cues), are extracted from the prototype prior as prototype cues to guide the visual base model (SAM) to achieve more accurate segmentation. The method uses prototype cues and a small amount of defect data to fine-tune the visual base model (SAM), freeze the SAM's cue encoder and image encoder, and inject a LoRa bypass into the image encoder. By updating the parameters of the LoRA bypass and lightweight mask decoder, the SAM is made more adaptable to the defect data distribution.

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 fine-tuning of a visual base model (SAM). Background Art

[0002] Industrial defect segmentation plays a crucial role in ensuring the quality and safety of production processes. Visual defect segmentation is a contactless, non-destructive, automated inspection technology that effectively and precisely controls equipment. Deep learning has significantly improved the performance of this task, but it relies on large amounts of labeled data. In some industrial scenarios, defect data is scarce, making it difficult to form large datasets, which can lead to a sharp decline in model performance.

[0003] Few-shot learning methods can alleviate the scarcity 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, hindering their promotion in other industrial environments.

[0004] The Visual Foundation Model (SAM) is pre-trained on large-scale general datasets to learn rich visual features. It has strong performance and generalization capabilities and can quickly adapt to new data distributions and downstream tasks through fine-tuning. However, SAM is an interactive visual foundation 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 these hints and a small amount of defect data to fine-tune the visual base 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 of the visual basis model (SAM). Prior information about the image to be segmented is extracted through examples, representative cues are extracted from the prior information, and the visual basis model (SAM) is fine-tuned using the cues and a small amount of defect data to adapt it to the defect data distribution, thereby improving the performance and generalization of the algorithm in 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-cued fine-tuning of the visual base model (SAM), which extracts 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 the prototype prior information as prototype cues; 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 hints; the hint encoder and image encoder of SAM are frozen, and a 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] 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 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 by 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 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 dense cue; extract the rough mask Connected component A in 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 expressed 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 evenly 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 as follows:

[0035]

[0036] Among them, 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: the parameters of the SAM image encoder and prompt encoder are frozen, and the image encoder is embedded in the LoRA bypass. 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 then input into the mask decoder; the mask decoder obtains the predicted mask; and 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 Represent the linear change matrices of the query vector Q, key vector K and value vector V respectively, d and k represent their input dimension and output dimension respectively, and 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 base 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, and uses these prototypes to calculate similarities with query image features to obtain prototype-based foreground / background prior probability maps, 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 base 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 a LoRA bypass into the image encoder. By updating the parameters of the LoRA bypass and lightweight mask decoder, it can better adapt to the distribution of defect data. BRIEF DESCRIPTION OF THE DRAWINGS

[0047] Figure 1 A flowchart of a few-shot defect detection method based on prototype-cue fine-tuning of the 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 prototype prompt extraction provided in 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-cue 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 support 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] Among them F S Indicates the supported image features, Y S Denotes the support image mask, and c denotes 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] Where Avg represents the averaging function, 1 represents the foreground class, and 0 represents the background class. Subsequently, 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 using 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 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 and background cue points as sparse cues and a processed coarse mask as dense cues. First, we extract Connected component A in 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 cue points into Each connected component in is refined to a pixel width, and these refined 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 refinement 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 size, 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 high-confidence points as negative prompt points. The calculation is as follows:

[0078]

[0079] Among them, P negative Represents the set of background cue points. Foreground cue points and background cue points together constitute sparse cues.

[0080] Step 4: Freeze the parameters of the SAM image encoder and 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 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, ensuring that BA = 0, thereby maintaining consistency with the initial state.

[0083] This method was tested on MVTec-AD texture defects. To verify the generalization ability of the resulting model to unseen classes and unseen industrial scenarios, it was divided into MVTec_Seen and MVTec_Unseen models based on different data partitioning. MVTec_Seen was trained and tested on the same material, but with unseen defect classes in the seen material. With one supporting image-mask pair (one-shot), it achieved an average Intersection-Over-Union (IoU) of 63.07%, and with five supporting image-mask pairs (five-shot). MVTec_Unseen was trained and tested on different materials, with unseen defect classes in the unseen material. With one supporting image-mask pair (one-shot), it achieved an average IoU of 53.51%, and with five supporting image-mask pairs (five-shot). This method does not rely on human prior guidance and has excellent performance and generalization capabilities. 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 the visual base model (SAM), characterized by: Extracting prototype prior information about the image to be segmented from a small number of example image-mask pairs, which replaces human prior guidance; For defect image data, representative foreground cue points, high-confidence background cue points, and foreground masks are extracted from the prototype prior information as prototype cues; the foreground cue points and background cue points are sparse cues, and the foreground mask is a dense cue; Fine-tune the visual base model (SAM) through prototype hints; freeze the hint encoder and image encoder of SAM, inject LoRA bypass into the image encoder, and obtain the adjusted visual base model (SAM) for defect detection by updating the parameters of LoRA bypass and mask decoder; 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: 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 P B , 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 by 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 M F , The strategy for fine-tuning the visual base model (SAM) with prototype prompts is designed as follows: the parameters of the SAM image encoder and prompt encoder are frozen, and the image encoder is embedded in the LoRA bypass. 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 then input into the mask decoder; the mask decoder obtains the predicted mask; and 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.

2. 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 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 dense cue; extract the rough mask Connected component A in 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 expressed 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 size, P positive Represents a set of foreground cue points; 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 high-confidence points as negative prompt points. The calculation is as follows: Among them, P negative Represents a collection of background cue points; Foreground cue points and background cue points together constitute sparse cues.

3. 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 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 Represent the linear change matrices of the query vector Q, key vector K and value vector V respectively, d and k represent their input dimension and output dimension respectively, and 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