Interactive fine tuning SAM for segmenting industrial surface defect picture and training method thereof

Abstract

The invention provides an interactive fine-tuning SAM for segmenting an industrial surface defect picture and a training method thereof, and belongs to the technical field of image segmentation. Fine tuning SAM includes: an image encoder, a hint encoder, a mask decoder embedded with an adapter structure including a downward projection that compresses a given embedding to a lower dimension using multiple perceptron layers, a ReLU activation function, and an upward projection that compresses a given embedding to a lower dimension using multiple perceptron layers. The ReLU activation function is responsible for nonlinear transformation, while the upward projection uses another multi-layer perceptron layer to extend the compressed embedding back to its original dimension. The interactive habits of the user are fully considered in the fine adjustment process, experiments show that the fine adjustment technology has the best experiment result in ten comparative fine adjustment methods, and an effective solution is provided for interactive image segmentation of the user on an industrial surface defect scene.

Description

Technical Field

[0001] The present invention belongs to the technical field of image segmentation, and specifically refers to a user in an image segmentation system, using SAM to interactively segment defects in an industrial surface defect image, and more particularly relates to an interactive fine-tuning SAM for segmenting industrial surface defect images, and a training method, training device, computing device and readable storage medium thereof. Background Art

[0002] Segment Anything Model (SAM) is a general-purpose image segmentation model released by Meta AI in April 2023. For almost any image, SAM can accurately segment the objects in it, even images and objects that have never been encountered during training. SAM can segment them, and has good versatility. In addition, unlike previous image segmentation models, SAM innovatively combines image segmentation with user prompts. Users can segment target objects by using prompt boxes and prompt points in the image, and can also perform interactive segmentation, continuously improving SAM's segmentation results, and well realizing the "cutout" function.

[0003] The structural diagram of SAM is shown in Figure 1As shown, it includes three components: image encoder, prompt encoder, and mask decoder. The image encoder is responsible for extracting image features and obtaining image embedding vectors. SAM's image encoder uses ViT (Vision Transformer) pre-trained by MAE (Masked Autoencoders) as the core architecture. ViT is a Transformer-based image processing model that can divide an image into multiple image blocks (patches) and convert the pixel information of each image block into a vector to learn image features. There are three different versions of ViT used by SAM, namely ViT-B, ViT-L and ViT-H. The main difference between them is the different scale of the model. The experiment of the present invention uses ViT-B. The main function of the prompt encoder is to convert the prompt information input by the user into a prompt vector. The prompt information can be a prompt point, a prompt box, text, etc. The mask decoder takes the output of the image encoder and the prompt encoder as input to generate the final segmentation mask and the corresponding IoU score. It is worth noting that for a picture, SAM only needs to use the image encoder once to extract the image feature vector. For subsequent interactions, the same image feature vector is used to combine the prompt encoder and mask decoder to obtain the interactively improved segmentation result. In addition, each time you segment, you can choose to return one mask or three masks by setting the multimask_output option. If true, three masks will be returned. If false, only one mask will be returned.

[0004] In the industrial production process, due to process technology and material problems, it is difficult to achieve zero defects in the products produced. Many industrial products may have production defects on their surfaces, such as metal surface defects, tile surface defects, cloth surface defects, etc. On the one hand, we hope to quickly and accurately mark the defect locations; on the other hand, these marked images can be used to train industrial visual target detection models for automatic detection and identification of defects in subsequent products. Therefore, the annotation of industrial surface defect images is a very valuable task. The interactive image segmentation function of SAM is very suitable for image annotation. Users only need to draw prompt boxes and prompt points, and SAM can segment the defects that users want to segment without the need for users to manually draw defect boundaries. In addition, the good versatility of SAM enables images in different fields to achieve good segmentation effects without retraining the model.

[0005] However, SAM is trained using large-scale natural scene images, which are rich in color and have clear boundaries of target objects, while industrial surface defect images such as Figure 2As shown in the figure, the color is relatively dim and single, and the boundary of the defect is not obvious enough, so the characteristics of natural scene pictures and industrial surface defect pictures are completely different, which restricts the segmentation performance of SAM on industrial surface defect pictures. In addition, considering the interactive habits of users during annotation, it is found through experiments that the selection area of ​​subsequent prompt points seriously affects the segmentation performance. This is because the training process of SAM does not fully consider the interactive habits of users. Therefore, based on the above two deficiencies, the present invention improves the interactive segmentation performance of SAM on industrial surface defect pictures. Summary of the invention

[0006] The purpose of the present invention is to solve the problem that SAM uses large-scale natural scene images and does not fully consider the user's interaction habits for training, resulting in unsatisfactory interactive segmentation performance in industrial surface defect images. The present invention uses adapter technology and considers the user's interaction habits to perform interactive and efficient parameter fine-tuning on SAM in industrial surface defect image scenes, thereby improving the interactive segmentation performance of SAM in industrial surface defect scenes.

[0007] According to a first aspect of the present invention, an interactive fine-tuning SAM for segmenting industrial surface defect images is provided, comprising: an image encoder, a hint encoder, and a mask decoder, wherein the image encoder and the mask decoder are embedded with an adapter structure, wherein the adapter structure comprises a down-projection, a ReLU activation function, and an up-projection, wherein the down-projection uses a multi-layer perceptron layer to compress a given embedding to a lower dimension, the ReLU activation function is responsible for performing a nonlinear transformation, and the up-projection uses another multi-layer perceptron layer to expand the compressed embedding back to its original dimension.

[0008] According to a second aspect of the present invention, there is provided a training method for interactively fine-tuning a SAM for segmenting industrial surface defect images, comprising:

[0009] Step 1: loading a pre-trained SAM and initializing it, including loading the pre-trained SAM, adding a first adapter layer and a second adapter layer to the image encoder in the pre-trained SAM, adding a third adapter layer, a fourth adapter layer and a fifth adapter layer to the mask decoder, and setting parameters for updating only the adapter layers in the image encoder and the mask decoder;

[0010] Step 2: Load the training data, which is industrial surface defect images, and preprocess the industrial surface defect images and masks;

[0011] Step 3: Start training. In each round of training, simulate nine interactions for each preprocessed image. For the first interaction, randomly select a foreground point from the true value region gt of the target object as the initial prompt point with equal probability, or add noise to the true value box as the initial prompt box, and the initial prompt box is larger than the true value box. For subsequent interactions, randomly select a subsequent prompt point from the error area of ​​the true value region gt of the target object and the segmentation region pred returned by the fine-tuning SAM. For the case where the true value region gt and the segmentation region pred intersect with each other, only select the foreground point from the false negative region, where the false negative region refers to the error area in the target area but not in the segmentation area.

[0012] In each interaction, the prompt encoder and mask decoder are used, and only the first interaction uses the image encoder to calculate the image embedding of the preprocessed image; for the prompt encoder, in addition to the prompt box and prompt point, the low-resolution mask map returned by the previous interaction is used as the input of the mask_input option, and sparse embedding and dense embedding are returned; for the mask decoder, the image embedding, sparse embedding and dense embedding are used as input, the multimask_ouput option is set to False, a low-resolution mask map and the corresponding IoU score are returned, and then the low-resolution mask map is scaled to obtain the mask map, and the three are returned. The loss function is calculated based on the three and the gradient is updated;

[0013] Step 4: After each round of training, save the model and record the average loss.

[0014] According to a third aspect of the present invention, there is provided an interactive fine-tuning SAM training device for segmenting industrial surface defect images, which is used to perform an interactive fine-tuning SAM training method for segmenting industrial surface defect images, comprising:

[0015] A first loading module is used to load a pre-trained SAM and initialize it, including loading the pre-trained SAM, adding a first adapter layer and a second adapter layer to the image encoder in the pre-trained SAM, adding a third adapter layer, a fourth adapter layer and a fifth adapter layer to the mask decoder, and setting parameters for updating only the adapter layers in the image encoder and the mask decoder;

[0016] The second loading module is used to load training data, the training data is industrial surface defect pictures, and preprocess the industrial surface defect pictures and masks;

[0017] Training module, used to start training;

[0018] The saving module is used to save the model and record the average loss after each round of training.

[0019] According to a fourth aspect of the present invention, there is provided a computing device comprising: at least one processor and a memory storing program instructions; when the program instructions are read and executed by the processor, the computing device executes an interactive fine-tuning SAM training method for segmenting industrial surface defect images.

[0020] According to a fifth aspect of the present invention, a readable storage medium storing program instructions is provided, and when the program instructions are read and executed by a computing device, the computing device executes an interactive fine-tuning SAM training method for segmenting industrial surface defect images.

[0021] Beneficial effects:

[0022] The present invention takes into account the user's interactive habits and uses industrial surface defect images for interactive fine-tuning, attempting to improve the interactive segmentation effect of the SAM model on industrial surface defect images. For the three component structures of SAM, the present invention inserts adapter layers in both the image encoder and the mask decoder. During fine-tuning, only the parameters in the adapter layer are updated, while other parameters in the image encoder and the mask decoder are kept frozen. The user's interactive habits are fully considered during the fine-tuning process. Experiments have shown that this fine-tuning technology has the best experimental results among ten compared fine-tuning methods, providing an effective solution for users' interactive image segmentation on industrial surface defect scenes. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] Figure 1 It is a schematic diagram of the structure of SAM;

[0024] Figure 2 This is a sample picture of industrial surface defects;

[0025] Figure 3 A flowchart for interactive image segmentation using SAM;

[0026] Figure 4 Design diagrams for interactive use of SAM;

[0027] Figure 5 This is an example diagram of a segmentation of SAM;

[0028] Figure 6 It is the relationship diagram between the segmented area pred returned by SAM and the true value area gt of the target object;

[0029] Figure 7 A schematic diagram of an adapter structure used in the interactive fine-tuning SAM for segmenting industrial surface defect images of the present invention;

[0030] Figure 8 A schematic diagram of the structure of an improved image encoder in the interactive fine-tuning SAM for segmenting industrial surface defect images of the present invention;

[0031] Fig. 9 It is a schematic diagram of the structure of the mask decoder of SAM;

[0032] Fig.10 A schematic diagram of the structure of the improved bidirectional attention block in the interactive fine-tuning SAM for segmenting industrial surface defect images of the present invention;

[0033] Fig.11 A training flow chart of the interactive fine-tuning SAM for segmenting industrial surface defect images of the present invention;

[0034] Fig.12 A comparison chart of the relationships between various fine-tuning methods;

[0035] Fig.13 A block diagram of the interactive fine-tuning SAM training device for segmenting industrial surface defect images according to the present invention. DETAILED DESCRIPTION

[0036] The technical problem to be solved by the present invention is: in the industrial surface defect image scene, considering the user's interaction habits, interactively fine-tuning SAM is performed to improve the interactive segmentation performance of SAM in the industrial surface defect scene, and to improve the speed and accuracy of marking defects in industrial surface defect images.

[0037] The design idea of ​​the present invention is:

[0038] SAM's fine-tuning techniques can be divided into two categories: Full Fine-Tuning and Parameter Efficient Fine-Tuning (PEFT). Full fine-tuning refers to fine-tuning all parameters based on the pre-trained model, while efficient parameter fine-tuning only updates some parameters in the model and freezes most parameters, which can significantly reduce training time and cost and achieve efficient learning through faster updates. In addition, existing studies have shown that PEFT is better than full fine-tuning because it avoids catastrophic forgetting and can be better generalized to out-of-domain scenarios, especially when there is only a small amount of fine-tuning data. There are two classic fine-tuning techniques for PEFT: Adapter and Low Rank Adaptation (LoRA). Adapter technology is achieved by inserting adapters in certain layers of the original model. These adapters are specially designed small sub-networks. The parameters of the adapters are trainable, while the other parameters of the original model remain frozen. The LoRA technology fine-tunes the weights of the original model by introducing a low-rank matrix. LoRA does not add a new network layer. The low-rank matrix it introduces is a special parameter module attached to the original layer. It has no independent input, output, and complete forward calculation process. Its function is to efficiently adjust the parameters of the existing model by low-rank decomposition of the weight matrix. During fine-tuning, only the introduced low-rank matrix is ​​trained, and the original weight matrix remains frozen.

[0039] Among the three components of SAM, the prompt encoder is a very simple module, so its structure is generally not changed during fine-tuning. For the image encoder and mask decoder, you can fine-tune only the image encoder, only the mask decoder, or both; if you fine-tune the image masker or the mask decoder, you can fine-tune the entire amount or fine-tune the parameters efficiently; for efficient parameter fine-tuning, you can use Adapter technology or LoRA technology. Therefore, there are many ways to fine-tune. In the experiment of the present invention, the Adapter structure and the LoRA structure were not modified, but the existing Adapter structure and LoRA structure used for fine-tuning SAM were used, because this technology is relatively mature, but the present invention combines and designs whether these structures are used and the insertion position, and finally found that the fine-tuning effect of inserting the Adapter layer in both the image encoder and the mask decoder for efficient parameter fine-tuning is the best.

[0040] First, the steps of interactive image segmentation using SAM are introduced.

[0041] For interactive image segmentation using SAM, the two functions provided by SAM are mainly used. The following is a simple example code:

[0042] sam=sam_model_registry["vit_b"](checkpoint="weights / sam_vit_b.pth").to(device="cuda")

[0043] predictor = SamPredictor(sam)

[0044] image = cv2.imread("image.jpg")

[0045] image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

[0046] predictor.set_image(image)

[0047] masks, iou_predictions, low_res_masks = predictor.predict(

[0048] point_coords=np.array([[234, 248]]),

[0049] point_labels = np.array([1]),

[0050] box=np.array([100, 100, 400, 400]),

[0051] multimask_output=False )

[0053] masks = masks.astype(np.uint8)

[0054] _, mask = cv2.threshold(masks[0], 0, 255, cv2.THRESH_BINARY)

[0055] cv2.imwrite('mask.jpg', mask)

[0056] like Figure 3 As shown, this code includes the following steps (Note: The element type in the following formula represents an integer, represents real numbers):

[0057] Step 1: Load the SAM model.

[0058] Step 2: For the new image, read the image and convert it to RGB format. .

[0059] Step 3: Call set_image() to calculate the image embedding of the image. The predictor object uses a variable to store the image embedding. For subsequent interactions, if the image has not changed, the image embedding will not be recalculated; if the image has changed, the image needs to be read again and the image embedding needs to be calculated.

[0060] Step 4: For subsequent interactions, call predict() to return the segmentation mask. SAM has three return values, namely, the segmentation mask , IoU score , low-resolution mask For the formula , when the input multimask_output is set to False, then , when the input multimask_output is set to True, then .

[0061] For the prediction function predict(), there are a total of six input options, namely:

[0062] : Prompt points. You can enter multiple prompt points for each prediction, or you can enter none.

[0063] : The type of the cue point. If it is 1, it means the foreground point. If it is 0, it means the background point (the foreground point and background point will be explained in detail later). The type of the cue point needs to correspond to the cue point one by one.

[0064] : Prompt box, the four values ​​represent the x-coordinate of the upper left point, the y-coordinate of the upper left point, the width of the prompt box, and the height of the prompt box. Only one prompt box is allowed to be entered for each prediction, or none is allowed. If a prompt point is entered, a prompt box can be entered or not, and vice versa;

[0065] : Low-resolution mask, usually from the return value of the previous interactive prediction;

[0066] : Whether to return multiple masks. If set to False, only one mask will be returned; if set to True, three masks will be returned. The default value is True. For interactive image segmentation, since only one mask needs to be returned, a post-processing step will be performed, and the mask corresponding to the highest IoU score will be returned as the segmentation result.

[0067] : Whether to return the binary mask. If set to False, it means it has been binary processed; if set to True, it means it has not been binary processed. The default value is False, and it is generally OK to use the default value.

[0068] Step 5: Save the segmentation mask.

[0069] It is worth noting that step 1 takes x seconds to load the SAM model, step 3 takes 0.x seconds to compute the image embedding, and step 4 takes only 0.0x seconds to return the segmentation mask.

[0070] By using multiple industrial surface defect datasets to simulate interactive experiments on these input options of the prediction function predict(), the interactive use method of SAM is designed as follows: Figure 4 As shown:

[0071] When starting the system, the SAM model is loaded;

[0072] When the user loads or switches an image, call set_image() to calculate the image embedding of the image;

[0073] Each time the user draws a prompt box or prompt point on the image, call predict() to return the segmentation mask and display the corresponding target area on the image;

[0074] The user can double-click the target area corresponding to the segmentation mask to save it;

[0075] The prompt box can only be drawn at the beginning, and then multiple prompt points can be drawn. The prompt point can be either a foreground point or a background point (in the prompt point mode, the left mouse button indicates the foreground point, and the right mouse button indicates the background point). If you draw the prompt box after drawing the prompt point, all the previously drawn prompt points will be cleared;

[0076] If a prompt box has been drawn, you can also draw a new prompt box, and the new prompt box will replace the existing prompt box. If prompt points have been drawn, the new prompt points drawn will be used as input to predict() together with the existing prompt points. If the drawn prompt points need to be deleted, you can only use the clear button to delete all prompts on the picture;

[0077] For subsequent interactions on the same image, use the low-resolution mask returned by the previous interaction as input to the mask_input option of predict();

[0078] Each time predict() is called, multimask_output is set to False, because it is found through experiments that setting it to False produces better segmentation results than setting it to True.

[0079] Then analyze the user's interaction habits.

[0080] like Figure 5 As shown in the figure, this is a segmentation example. Suppose the user wants to segment out a large black area in the image, and the area within the white boundary is the segmentation result returned by SAM after the user draws a hint point. It is far from the target area to be segmented, so the user intends to improve the segmentation result through interactive segmentation. There are two types of error areas between the segmentation area returned by SAM last time and the target area the user wants. The first is a false negative (in the target area but not in the segmentation area), and the second is a false positive (in the segmentation area but not in the target area). For the next interaction, the user will draw subsequent hint points in the error area, drawing foreground points in false negatives and background points in false positives. In this case, the user is likely to draw foreground points in false negatives and not background points in false positives.

[0081] In fact, there are three situations in which the relationship between the segmented region pred returned by SAM and the true value region gt of the target object exists, such as Figure 6 As shown:

[0082] The first case: gt completely contains pred, the error area only has false negatives, and the user will only select a foreground point as the subsequent prompt point.

[0083] The second case: pred completely contains gt, there are only false positives in the error area, and the user will only select a background point as the subsequent prompt point.

[0084] The third case: gt and pred intersect each other, and there are both false negatives and false positives in the error area. In this case, the user is more inclined to select a foreground point as the subsequent prompt point rather than a background point as the subsequent prompt point. However, the training process of SAM randomly selects subsequent prompt points in the error area, which may be foreground points or background points, and the user's interaction habit is basically to only select foreground points. It is also found through experiments that the average segmentation effect of only selecting foreground points is much worse than randomly selecting subsequent prompt points from foreground and background points.

[0085] Table 1

[0086]

[0087] Table 1 shows the experimental results of simulating interactive segmentation on eight industrial surface defect datasets. The leftmost 0bbox indicates that no prompt box is used at the beginning, 0.64 bbox indicates that a prompt box with an area of ​​0.64 times the area of ​​the true value box is used at the beginning (the width and height of the prompt box are both 0.8 times the corresponding values ​​of the true value box. In order to control the variables for comparison, the width or height of the prompt box is shrunk to the inside by the same length relative to the true value box). The processing of 0.81 bbox, 1 bbox, etc. is similar. 0point indicates that no prompt point is used subsequently, 1 point indicates that a subsequent prompt point is selected in the last error area. If the prompt box has not been used before, 1 point indicates that a foreground point is selected in the true value area, 3 point indicates that three prompt point interactions are performed, and each time a subsequent prompt point is selected in the last error area, and 5 point is similar.

[0088] For the selection of subsequent prompt points, two situations are simulated. The first situation is Figure 6 In the three cases shown, subsequent prompt points are randomly selected, so for the third case, the foreground point may be selected, or the background point may be selected; the second case is that for the third case, only the foreground point is selected from the false negatives. The evaluation index is the average of the dice coefficients of the last interaction of all defects in the data set. The experiment is repeated three times and the average of the three experimental results is taken.

[0089] ,

[0090] In the above formula, Indicates the image number. Indicates the defect number in the image. represents the area of ​​the corresponding region, represents the true value region of the defect, Indicates the segmented area returned by the last interaction SAM. represents the number of defects, represents the dice coefficient of the last interaction of the defect, represents the result of one experiment, and result represents the average of three experimental results.

[0091] The number in front of each result cell in Table 1 represents the simulation of the first situation, and the number behind represents the simulation of the second situation. It can be seen that in most result cells, the number in front is larger than the number behind, and the gap is very large, especially when only cue points are used. The maximum gap is 0.272. The case where the number behind is larger than the number in front is displayed in bold, and the maximum gap is only 0.013. In addition, for the second situation, as the number of cue points increases, the improvement in the segmentation effect is not obvious, and may even regress. For the first situation, as the number of cue points increases, the improvement in the segmentation effect is more obvious. This is because for example, Figure 6 In the third case shown, the training process of SAM randomly selects subsequent prompt points in the wrong area, which may be foreground points or background points, and does not fully consider the user's interaction habit of basically only selecting foreground points.

[0092] In addition, we can see that the segmentation effect when the prompt box is slightly smaller than the true value box is also good, and is not worse than the segmentation effect when the prompt box is slightly larger than the true value box. This is because when using the prompt box in the training process of SAM, noise is randomly added to the true value box, so the prompt box may be larger or smaller than the true value box. Users tend to draw prompt boxes larger than the true value box.

[0093] Therefore, the present invention takes these two interactive habits of users into consideration and uses the industrial surface defect dataset to interactively fine-tune SAM, attempting to improve the interactive segmentation effect of SAM in industrial surface defect scenarios.

[0094] The design of the network structure of the present invention is described below.

[0095] Figure 7 Schematic diagram of the adapter structure used in the present invention, including a downward projection, a ReLU activation function and an upward projection. The downward projection uses a simple multi-layer perceptron to compress a given embedding to a lower dimension, the ReLU activation function is responsible for nonlinear transformation, and the upward projection uses another multi-layer perceptron to expand the compressed embedding back to its original dimension.

[0096] The following describes the interactive fine-tuning SAM structure of the present invention for segmenting industrial surface defect images taking into account user interaction habits. Figure 8This is a schematic diagram of the structure of the improved image encoder, which contains multiple transformer blocks. The transformer blocks include a normalization layer, a multi-head attention, a first adapter layer, a normalization layer, a multi-layer perceptron, and a second adapter layer in parallel with the multi-layer perceptron in order. For each transformer block, adapters are used in two places. The first adapter layer is located after the multi-head attention and uses skip connections. It emphasizes lightweight adjustments while retaining the original features and focuses on fine-tuning in the spatial dimension. The second adapter layer is used for parallel operations with the multi-layer perceptron. It does not use skip connections and is more inclined to fine-tune at the feature level to increase the adaptability of the model to new tasks. During fine-tuning, only the parameters in the adapter are trained, and the parameters of other trained network layers are frozen.

[0097] In addition to using the adapter to fine-tune the mask decoder, the adapter is also used to fine-tune the mask decoder. Fig. 9 As shown in the figure, in the mask decoder, the two-way transformer is the core module, which is mainly used to realize multi-modal feature interaction and propagation. The two-way transformer accepts image embedding and hint embedding provided by the user, and realizes the interaction of these two features through the transformer layer, thereby aligning the information of different modalities in the latent space, so that the model can flexibly and accurately complete a variety of segmentation tasks.

[0098] Fig. 9 In the above figure, the image embedding is the output of the image encoder, the cue tag is the output of the cue masker, the output tag is a learnable code randomly initialized according to the output features, and two output tensors are obtained after two bidirectional attention blocks. Then, the predicted low-resolution mask map is obtained after a convolution transformation and a multi-layer perceptron, and the predicted IoU score is obtained after another multi-layer perceptron.

[0099] The bidirectional transformer mainly consists of two bidirectional attention blocks (TwoWayAttentionBlock). The structural diagram of the improved bidirectional attention block is as follows: Fig.10As shown in the figure, the input q and k represent sparse embedding and dense input respectively, while q_pe and k_pe are the position information of the two input data. The self-attention in the figure is used to process sparse input, while the token-to-image attention and image-to-token attention are cross-attention layers used to process the interaction between sparse input and dense input. The normalization layers in the figure are used to standardize the input. The multi-layer perceptron is used to perform nonlinear transformation on the input to enhance the expressiveness of the model. Three adapters are used in each bidirectional attention block. The third adapter layer is located after the token-to-image attn, and its function is to adjust the image features extracted from the prompt information and optimize the alignment of the features with the prompt information. The fourth adapter layer is used to operate in parallel with the multi-layer perceptron, and its function is to perform deeper nonlinear adaptation of the features. The fifth adapter layer is located after the image-to-token attention, and its function is to optimize the expressiveness of the features and enhance the alignment and fusion effects of the features. During fine-tuning, only the parameters in the adapter are trained, and the other parameters remain frozen.

[0100] Design of loss function

[0101] The loss function used in the SAM training process is a linear combination of focal loss, dice loss, and IoU loss. Focal loss is suitable for binary classification tasks, especially in scenarios where the positive and negative classes are severely unbalanced. Dice loss optimizes the shape overlap between the predicted mask and the true mask, and IoU loss enhances the prediction ability of IoU. The comprehensive loss function combines the above advantages and is suitable for scenarios that require accurate mask and IoU prediction in segmentation tasks.

[0102] ,

[0103] In the above formula, represents the true value region of the target defect, Indicates the segmented area returned by SAM. is the IoU score returned by SAM, The weight used to balance positive and negative samples. The default value is 0.25. Used to adjust the weight of difficult samples. The default value is 2. Used to calculate the number of pixels with a value of 1 in the corresponding area. For calculating the natural logarithm, Used to calculate the total number of pixels in the corresponding area, is the batch size, and is the weight coefficient.

[0104] Interactive fine-tuning strategies

[0105] Following the pre-processing and post-processing strategies used in the SAM inference process, all images and masks are adjusted to , and then pad it with 0 pixels on the bottom right edge to a resolution of 1024×1024.

[0106] In each round of training, nine interactions are simulated for each defect in the image. For the first interaction, a foreground point is randomly selected from the gt as the initial cue point with equal probability, or noise is added to the ground truth box as the initial cue box. The initial cue box is larger than the ground truth box, and the settings are as follows:

[0107] ,

[0108] In the above formula, represents the true value region of the target defect, The outer rectangular box used to calculate gt, that is, the true value box, returned They represent the x-coordinate of the upper left point of the truth box, the y-coordinate of the upper left point of the truth box, the width of the truth box, and the height of the truth box, respectively. Indicates the maximum offset value of the left and right sides of the prompt box relative to the true value box. Indicates the maximum offset value of the upper and lower sides of the prompt box relative to the true value box. Used to calculate arrive A random integer between , , , Indicates the outward extension value of the four sides of the prompt box relative to the true value box. It is the initial prompt box.

[0109] For subsequent interactions, a subsequent prompt point is randomly selected from the error area of ​​gt and pred each time. For the case where gt and pred intersect each other, only the foreground point is selected from the false negative. In subsequent interactions, the previous prompt is retained, and the low-resolution mask map returned by the previous interaction is used as the input of mask_input. mutimask_output is set to False, indicating that only one segmentation mask is returned.

[0110] During the training process, only the first interaction needs to use the image encoder to extract image features, combine the prompt encoder and mask decoder to calculate gradients, and update the parameters of these three modules. Subsequent interactions no longer need to use the image encoder, only the prompt encoder and mask decoder, so subsequent interactions will not update the parameters of the image encoder, only the parameters of the prompt encoder and mask decoder. Because the adapter layer is added to the image encoder and mask decoder, only the parameters of the adapter layer in the image encoder and mask decoder are updated, while other parameters remain frozen.

[0111] Model training steps

[0112] The training steps are as follows Fig.11 As shown, there are the following steps:

[0113] Step 1: Load the pre-trained SAM and initialize it, including loading the pre-trained SAM, adding the first adapter layer and the second adapter layer to the image encoder in the pre-trained SAM, adding the third adapter layer, the fourth adapter layer and the fifth adapter layer to the mask decoder, and setting parameters to update only the adapter layers in the image encoder and the mask decoder.

[0114] Step 2: Load the training data, which is industrial surface defect images. Preprocess the industrial surface defect images and masks.

[0115] Step 3: Start training. In each round of training, simulate nine interactions for a preprocessed image. For the first interaction, randomly select a foreground point from the ground truth region gt of the target object as the initial cue point with equal probability, or add noise to the ground truth box as the initial cue box, and the initial cue box is larger than the ground truth box. For subsequent interactions, randomly select a subsequent cue point from the ground truth region gt of the target object and the error region of the segmentation region pred returned by the fine-tuning SAM. For the case where gt and pred intersect with each other, only select the foreground point from the false negative region, where the false negative region refers to the error region that is in the target region but not in the segmentation region.

[0116] In each interaction, the prompt encoder and mask decoder are needed, and only the first interaction requires the image encoder to calculate the image embedding from the preprocessed image. For the prompt encoder, in addition to the prompt box and prompt point, the low-resolution mask map returned by the previous interaction is used as the input of the mask_input option, and the sparse embedding and dense embedding are returned. For the mask decoder, the image embedding, sparse embedding, and dense embedding are used as input, the multimask_ouput option is set to False, the low-resolution mask map and IoU score are returned, and then the low-resolution mask map is scaled to obtain the mask map, and the three are returned. The loss function is calculated based on the three and the gradient is updated.

[0117] Step 4: After each round of training, save the model and record the average loss.

[0118] The inference steps of the model

[0119] The trained model is compatible with the inference step of SAM, and there is no difference for users. Compared with the training process, the inference process does not update the model parameters, but another difference is that the mask image needs to be post-processed and then returned to the user. The post-processing first crops the 0 pixels at the bottom right edge of the mask image, and then scales it to the original resolution size of the image. The post-processing operation is the inverse of the previous pre-processing operation on the image. The inference step of the original SAM model also has this post-processing operation.

[0120] Experimental Section

[0121] (1) Building a dataset

[0122] Eight open source industrial surface defect datasets were collected, merged, and divided into training and test sets in a ratio of 8:2. For the training set, there are 5660 images with a total of 14084 defects; for the test set, there are 1417 images with a total of 3481 defects.

[0123] (2) Contrastive fine-tuning method

[0124] For the three components of SAM, the image encoder and mask decoder contain Transformer, while the hint encoder does not contain Transformer. The adapter technology (Adapter) and low-rank adaptation technology (LoRA) for fine-tuning large models are usually used for Transformer structure. Therefore, Adapter technology or LoRA technology can be used in the image encoder and mask decoder of SAM. The hint encoder is a very simple module, so its structure is generally not changed. Use interactive strategies to fine-tune SAM. Because the subsequent interaction process will not use the image encoder, but the mask decoder, the subsequent interaction will not update the image encoder, but should update the mask decoder. Therefore, you can only fine-tune the mask decoder, or fine-tune the image encoder and mask decoder, but not only the image encoder. It is worth mentioning that the hint encoder will be fine-tuned, and it is a full parameter fine-tuning. As for how to fine-tune the image encoder or mask decoder, you can use the full fine-tuning method to fine-tune all parameters, or you can use the Adapter technology or LoRA technology in the efficient parameter fine-tuning method to fine-tune some parameters. Therefore, if Fig.12 As shown, there are a total of 10 fine-tuning methods:

[0125] Method 1: Do not fine-tune the image encoder, but fine-tune all parameters of the mask decoder;

[0126] Method 2: Instead of fine-tuning the image encoder, use the adapter technique to fine-tune some parameters of the mask decoder;

[0127] Method 3: Do not fine-tune the image encoder, but use LoRA technology to fine-tune some parameters of the mask decoder;

[0128] Method 4: Fine-tune all parameters of the image encoder and fine-tune all parameters of the mask decoder;

[0129] Method 5: Fine-tune all parameters of the image encoder and use the adapter technique to fine-tune some parameters of the mask decoder;

[0130] Method 6: Fine-tune all parameters of the image encoder and use LoRA technology to fine-tune some parameters of the mask decoder;

[0131] Method 7: Use Adapter technology to fine-tune some parameters of the image encoder and all parameters of the mask decoder;

[0132] Method 8: Use LoRA technology to fine-tune some parameters of the image encoder and all parameters of the mask decoder;

[0133] Method 9: Use Adapter technology to fine-tune some parameters of the image encoder and use Adapter technology to fine-tune some parameters of the mask decoder;

[0134] Method 10: Use LoRA technology to fine-tune some parameters of the image encoder and use LoRA technology to fine-tune some parameters of the mask decoder.

[0135] (3) Training Model

[0136] For each fine-tuning method, the model is trained on the training set for 30 rounds, using the Adam optimizer, with an initial learning rate of 1e-4, and the learning rate is reduced to half every 10 rounds. The loss function uses a linear combination of focal loss, dice loss, and IoU loss. For each round of training, the average loss is recorded, and if the average loss is smaller than the previous one, the model is saved. NVIDIA GeForce RTX 4090 is used for training, and the batch size (batch_size) is 4.

[0137] (4) Test evaluation indicators

[0138] For each fine-tuning method, the trained model with the smallest average loss is used to simulate the following 10 interactions on the test set:

[0139] 0 bbox + 1 point: randomly select a foreground point in gt;

[0140] 0 bbox + 3 point: randomly select a foreground point in gt, and then interactively select the prompt point twice;

[0141] 0 bbox + 5 point: randomly select a foreground point in gt, and then perform four interactive selections of cue points;

[0142] 1.21 bbox + 0 point: use a prompt box that is 1.21 times the area of ​​the real box;

[0143] 1.21 bbox + 1 point: first use a hint box with an area 1.21 times that of the real box, and then interactively select the hint point;

[0144] 1.21 bbox + 3 point: first use a hint box with an area 1.21 times that of the real box, and then perform three interactive selections of hint points;

[0145] 1.21 bbox + 5 point: first use a hint box with an area 1.21 times that of the real box, and then perform five interactive selections of hint points;

[0146] 1.44 bbox + 0 point: Use a prompt box that is 1.44 times the area of ​​the real box;

[0147] 1.44 bbox + 1 point: first use a hint box with an area 1.44 times that of the real box, and then interactively select the hint point;

[0148] 1.44 bbox + 3 point: first use a hint box with an area 1.44 times that of the real box, and then perform three interactive selections of hint points;

[0149] 1.44 bbox + 5 point: First use a hint box with an area 1.44 times that of the real box, and then interactively select hint points five times.

[0150] It is worth noting that the width or height of the prompt box is expanded outward by the same length relative to the true value box. The evaluation index is the average of the dice coefficients of the last interaction of all target areas in the dataset. The experiment is repeated three times and the average of the three experimental results is taken.

[0151] (5) Test results

[0152] Table 2

[0153]

[0154] Table 2 shows the test results, from which we can see that the result of method 9 is the best, that is, the fine-tuning effect of adding the adapter layer to both the image encoder and the mask decoder is the best.

[0155] In summary, the steps of using SAM for interactive image segmentation are first explained, and then combined with the user's interactive habits, it is found that the selection area of ​​the subsequent prompt points will seriously affect the segmentation performance. In addition, there is a natural gap between industrial surface defect pictures and SAM's training pictures, which restricts the segmentation performance of SAM on industrial surface defect pictures. Therefore, based on the above two points, the present invention considers user habits and uses industrial surface defect pictures for interactive fine-tuning, trying to improve the interactive segmentation effect of the SAM model on industrial surface defect pictures. For the three component structures of SAM, the present invention inserts adapter layers in both the image encoder and the mask decoder. During fine-tuning, only the parameters in the adapter layer are updated, while other parameters in the image encoder and the mask decoder are kept frozen. The user's interactive habits are fully considered in the fine-tuning process. It is found through experiments that this fine-tuning technology has the best experimental results among ten compared fine-tuning methods, providing an effective solution for users' interactive image segmentation in industrial surface defect scenes.

[0156] An interactive fine-tuning SAM training device for segmenting industrial surface defect images is provided, which is used to perform an interactive fine-tuning SAM training method for segmenting industrial surface defect images, such as Fig.13 As shown, the device comprises:

[0157] The first loading module 1310 is used to load the pre-trained SAM and initialize it, including loading the pre-trained SAM, adding the first adapter layer and the second adapter layer to the image encoder in the pre-trained SAM, adding the third adapter layer, the fourth adapter layer and the fifth adapter layer to the mask decoder, and setting to update only the parameters of the adapter layers in the image encoder and the mask decoder;

[0158] The second loading module 1320 is used to load training data, the training data is industrial surface defect pictures, and preprocess the industrial surface defect pictures and masks;

[0159] A training module 1330 is used to start training;

[0160] The saving module 1340 is used to save the model and record the average loss after each round of training.

[0161] A computing device is provided, comprising: at least one processor and a memory storing program instructions; when the program instructions are read and executed by the processor, the computing device executes an interactive fine-tuning SAM training method for segmenting industrial surface defect images.

[0162] A readable storage medium storing program instructions is provided. When the program instructions are read and executed by a computing device, the computing device executes an interactive fine-tuning SAM training method for segmenting industrial surface defect images.

Claims

1. An interactive fine-tuning SAM for segmenting industrial surface defect images, characterized in that: include: Image encoder, hint encoder, mask decoder, the image encoder and mask decoder are embedded with an adapter structure, which contains a down-projection, a ReLU activation function, and an up-projection, where the down-projection uses a multi-layer perceptron layer to compress the given embedding to a lower dimension, the ReLU activation function is responsible for non-linear transformation, and the up-projection uses another multi-layer perceptron layer to expand the compressed embedding back to its original dimension. Ways to embed adapter structures in the mask decoder include: In the mask decoder, the bidirectional transformer contains two bidirectional attention blocks. In each bidirectional attention block, the third adapter layer is located after the tag-to-image attention, which is used to adjust the image features extracted from the prompt information and optimize the alignment of the features with the prompt information; the fourth adapter layer is used to operate in parallel with the multi-layer perceptron, which is used to perform deeper nonlinear adaptation of the features; the fifth adapter layer is located after the image-to-tag attention, which is used to optimize the expressiveness of the features and enhance the feature alignment and fusion effects; During fine-tuning, only the parameters in the third, fourth, and fifth adapter layers are trained, and the other parameters remain frozen.

2. The interactive fine-tuning SAM for segmenting industrial surface defect images according to claim 1, characterized in that Ways to embed adapter structures in image encoders include: The image encoder includes multiple transformer blocks, each of which includes a normalization layer, a multi-head attention, a first adapter layer, a normalization layer, a multi-layer perceptron, and a second adapter layer in parallel with the multi-layer perceptron in order; the first adapter layer is located after the multi-head attention and uses a skip connection, and the second adapter layer is used for parallel operation with the multi-layer perceptron and does not use a skip connection; During fine-tuning, only the parameters in the first and second adapter layers are trained, while the parameters of other trained network layers are frozen.

3. An interactive fine-tuning SAM training method for segmenting industrial surface defect images, applicable to the interactive fine-tuning SAM for segmenting industrial surface defect images according to claim 1 or 2, characterized in that: include: Step 1: loading a pre-trained SAM and initializing it, including loading the pre-trained SAM, adding a first adapter layer and a second adapter layer to the image encoder in the pre-trained SAM, adding a third adapter layer, a fourth adapter layer and a fifth adapter layer to the mask decoder, and setting parameters for updating only the adapter layers in the image encoder and the mask decoder; Step 2: Load the training data, which is industrial surface defect images, and preprocess the industrial surface defect images and masks; Step 3: Start training; Step 4: After each round of training, save the model and record the average loss.

4. The interactive fine-tuning SAM training method for segmenting industrial surface defect images according to claim 3 is characterized in that: In step three, in each round of training, nine interactions are simulated for each preprocessed image; for the first interaction, a foreground point is randomly selected from the true value region gt of the target object as the initial prompt point with equal probability, or noise is added to the true value box as the initial prompt box, and the initial prompt box is larger than the true value box; for subsequent interactions, a subsequent prompt point is randomly selected from the error area of ​​the true value region gt of the target object and the segmentation region pred returned by the fine-tuning SAM each time. For the case where the true value region gt and the segmentation region pred intersect with each other, only the foreground point is selected from the false negative area, where the false negative area refers to the error area in the target area but not in the segmentation area.

5. The interactive fine-tuning SAM training method for segmenting industrial surface defect images according to claim 4 is characterized in that: In step three, the prompt encoder and mask decoder are used in each interaction, and only the first interaction uses the image encoder to calculate the image embedding of the preprocessed image; for the prompt encoder, in addition to the prompt box and prompt point, the low-resolution mask map returned by the previous interaction is used as the input of the mask_input option, and sparse embedding and dense embedding are returned; for the mask decoder, the image embedding, sparse embedding and dense embedding are used as input, the multimask_ouput option is set to False, a low-resolution mask map and the corresponding IoU score are returned, and then the low-resolution mask map is scaled to obtain the mask map, and the three are returned. The loss function is calculated based on the three and the gradient is updated.

6. An interactive fine-tuning SAM training device for segmenting industrial surface defect images, used to execute the interactive fine-tuning SAM training method for segmenting industrial surface defect images according to claim 4 or 5, characterized in that: include: A first loading module is used to load a pre-trained SAM and initialize it, including loading the pre-trained SAM, adding a first adapter layer and a second adapter layer to the image encoder in the pre-trained SAM, adding a third adapter layer, a fourth adapter layer and a fifth adapter layer to the mask decoder, and setting parameters for updating only the adapter layers in the image encoder and the mask decoder; The second loading module is used to load training data, the training data is industrial surface defect pictures, and preprocess the industrial surface defect pictures and masks; Training module, used to start training; The saving module is used to save the model and record the average loss after each round of training.

7. The interactive fine-tuning SAM training device for segmenting industrial surface defect images according to claim 6 is characterized in that: In the training module, in each round of training, nine interactions are simulated for each preprocessed image; for the first interaction, a foreground point is randomly selected from the true value region gt of the target object as the initial prompt point with equal probability, or noise is added to the true value box as the initial prompt box, and the initial prompt box is larger than the true value box; for subsequent interactions, a subsequent prompt point is randomly selected from the error area of ​​the true value region gt of the target object and the segmentation region pred returned by the fine-tuning SAM each time. For the case where the true value region gt and the segmentation region pred intersect with each other, only the foreground point is selected from the false negative area, where the false negative area refers to the error area in the target area but not in the segmentation area.

8. A computing device, characterized in that include: at least one processor and a memory storing program instructions; When the program instructions are read and executed by the processor, the computing device executes the interactive fine-tuning SAM training method for segmenting industrial surface defect images as claimed in claim 4 or 5.

9. A readable storage medium storing program instructions, characterized in that: When the program instructions are read and executed by a computing device, the computing device executes the interactive fine-tuning SAM training method for segmenting industrial surface defect images as claimed in claim 4 or 5.

Images