Target segmentation method and system based on SAM model
By performing parameter adjustment and replacing the coder of the SAM model image encoder as a target detection encoding module, the SAM model has poor segmentation effect on infrared images and the interactive click -through dependence. Essence
Patent Information
- Application Number
- CN202411445829.2
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-10-16
- Publication Date
- 2025-01-28
AI Technical Summary
The segmentation results of existing SAM models on infrared images are not accurate enough, and interactive clicks are required to segment in a specified area, resulting in low usage efficiency.
By fine -tuning the image encoder of the SAM model by fine -tuning module to adjust the parameter of the image encoder, it allows it to obtain a good target segmentation effect on infrared images. Click.
It achieves good target segmentation effects on infrared images, while maintaining generalization capabilities, improving the target segmentation efficiency, and greatly improving model training efficiency.
Smart Images
Figure CN119380020A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of machine vision technology, and in particular to a target segmentation method and system based on a SAM model. Background Art
[0002] SAM (Segment Anything Model) is a model used for image segmentation. Users only need to interactively click on the area of interest in the image to obtain the corresponding segmented area.
[0003] like Figure 1 As shown in the figure, SAM includes an image encoder, a mask decoder, and a prompt encoder. The image encoder is used to encode the image to obtain an image embedding layer, and the prompt encoder encodes prompt information such as point, box, mask, text, etc. to obtain the corresponding embedding layer. Finally, the mask decoder integrates the embedding layers obtained by the image encoder and the prompt encoder to obtain the predicted mask of the target.
[0004] However, due to the domain differences between visible light images and infrared images, the segmentation results of SAM on infrared images are not accurate enough, and interactive clicking is required to segment the specified area, resulting in low efficiency. Summary of the invention
[0005] The purpose of the present invention is to provide a target segmentation method and system based on the SAM model, which can obtain good target segmentation effect on infrared images, while allowing the target segmentation model to maintain generalization ability, and at the same time, there is no need to obtain segmentation results through interactive clicking, thereby greatly improving the target segmentation efficiency.
[0006] To achieve the above object, the present invention provides the following technical solutions:
[0007] On the one hand, a target segmentation method based on a SAM model is provided, which comprises the following steps:
[0008] The preprocessed image is input into an image encoder whose parameters have been adjusted to output image features;
[0009] Fuse the image features to obtain the target detection data fusion features;
[0010] Perform target detection based on target detection data fusion features to obtain target detection output results and automatically output prompt information;
[0011] Performing target segmentation according to the prompt information to output a predicted mask of the target;
[0012] After the model iterative training is completed, the target segmentation model is output;
[0013] And, the target in the current image is segmented through the output target segmentation model, and the target segmentation result is output.
[0014] Preferably, the parameters of the image encoder of the SAM model are adjusted through a fine-tuning module.
[0015] Preferably, the image encoder of the SAM model is adjusted in parameters by a fine-tuning module, comprising the following steps:
[0016] Freezing weight parameters of the image encoder;
[0017] The weight parameters of the fine-tuning module are adjusted, and the adjusted weight parameters of the fine-tuning module are merged with the weight parameters of the image encoder, and the weight parameters obtained by merging are used as the weight parameters of the image encoder.
[0018] Preferably, the image features are fused through a feature pyramid network model.
[0019] Preferably, the target detection output result includes one or more of a target positioning frame and a target category.
[0020] Preferably, during the iterative training of the model, the total loss function L total Supervise model training, and L total =α*L mask +β*L reg +λ*L cls Among them, L mask is the mask loss, L reg is the regression loss function, L cls is the classification loss function, and α, β, and λ are adjustment coefficients.
[0021] Preferably, prompt information is received through a mask decoder of the SAM model, and target segmentation is performed according to the prompt information to output a predicted mask of the target.
[0022] Preferably, during iterative model training, parameters of the mask decoder are frozen.
[0023] On the other hand, there is also a target segmentation system, which includes: a fine-tuning module, a target detection encoding module, an image encoder of a SAM model, a mask decoder of the SAM model, a loss function construction module and a target segmentation module;
[0024] The fine-tuning module is used to adjust the parameters of the image encoder of the SAM model, and input the pre-processed image into the image encoder whose parameters have been adjusted;
[0025] The image encoder whose parameters have been adjusted is used to receive the preprocessed image and output image features;
[0026] The target detection encoding module is used to receive the image features output by the image decoder, fuse the image features to obtain target detection data fusion features, perform target detection according to the target detection data fusion features to obtain target detection output results, and automatically output prompt information;
[0027] The mask decoder is used to receive prompt information output by the target detection encoding module, and perform target segmentation according to the prompt information to output a predicted mask of the target;
[0028] Loss function construction module, which is used to construct the total loss function L total , and supervise the model training through the total loss function until the training meets the preset conditions, and then output the target segmentation model finally used for target segmentation;
[0029] And a target segmentation module, which is used to segment the target in the current image according to the target segmentation model ultimately used for target segmentation, and output the target segmentation result.
[0030] Preferably, the target detection encoding module includes:
[0031] A feature fusion module, which is used to fuse image features;
[0032] The target box is used to detect the fusion features of the target detection data to output the target positioning box;
[0033] The classification head cls is used to detect the fusion features of the target detection data to output the target category;
[0034] And the prompt header prompt, which is used to automatically output prompt information.
[0035] In summary, the present invention has the following beneficial effects compared with the prior art:
[0036] The present invention fine-tunes some parameters of the image encoder of the SAM model through a fine-tuning module, so that it can also obtain a good target segmentation effect on infrared images, while allowing the target segmentation model to maintain generalization ability; at the same time, the target detection encoding module replaces the prompt encoder of the SAM model, so that it can automatically generate and output prompt information while obtaining the target detection output result, so that the segmentation result is obtained without interactive clicking, which greatly improves the target segmentation efficiency;
[0037] At the same time, during the model training process, the relevant parameters of the image encoder and mask decoder are frozen, and the parameters are no longer updated. Only the fine-tuning module and the target detection encoding module are trained, which can greatly improve the model training efficiency. BRIEF DESCRIPTION OF THE DRAWINGS
[0038] Figure 1 It is a structural schematic diagram of a SAM model in the prior art;
[0039] Figure 2 It is a structural schematic diagram of the target segmentation system in the present invention;
[0040] Figure 3 It is the target segmentation result image obtained after being processed by the target segmentation method of the present invention. DETAILED DESCRIPTION
[0041] The technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, rather than all the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative work are within the scope of protection of the present invention.
[0042] Example 1
[0043] like Figure 1 As shown, this embodiment provides a target segmentation method based on the SAM model, which includes the following steps:
[0044] S1, constructing a training data set and a target segmentation model, wherein the training data set includes a plurality of images and annotation information corresponding to each image, and preprocessing the images in the training data set;
[0045] In this embodiment, the image P includes a visible light image and / or an infrared image, and each image P contains a target; the annotation information includes one or more of a target positioning box, a target category (cls), and a target mask (mask) obtained through image segmentation. The target positioning box (box) can be generated based on a target detection algorithm such as a YOLO series algorithm (such as YOLOv5, YOLOv8, etc.), an SSD algorithm (Single Shot MultiBox Detector), or RetinaNet, and the target in the image is annotated in a VOC format;
[0046] The target may be a common category with high detection and recognition accuracy, such as a human body, a vehicle, an animal, a building, an electric power equipment, etc. Specifically, the electric power equipment includes one or more of an oil storage cabinet, a cooler, a submersible oil pump, a relay, a metal expander, an insulator, a connection fitting, a voltage-grading ring, a lightning arrester, a hydraulic gauge, an oil temperature gauge, a conductive joint, a conductive lead, etc.;
[0047] The preprocessing includes one or more operations such as resizing, flipping, and normalizing the image P;
[0048] The target segmentation model includes a fine-tuning module 100, a target detection encoding module 200, an image encoder (Image Encoder) 300 of a SAM model, and a mask decoder (Mask Decoder) 400 of a SAM model;
[0049] S2, adjusting the parameters of the image encoder 300 of the SAM model through the fine-tuning module 100, and inputting the preprocessed image P into the image encoder 300 after the parameter adjustment, so as to output image features through the image decoder 300, and the image features may be implicit feature vectors;
[0050] Specifically, the fine-tuning module 100 is used to adjust the parameters of the image encoder 300 of the SAM model, including the following steps:
[0051] Freezing the weight parameters of the image encoder 300, that is, the weight parameters of the image encoder 300 remain unchanged;
[0052] Adjust the weight parameters of the fine-tuning module 100, merge the adjusted weight parameters of the fine-tuning module 100 with the weight parameters of the image encoder 300, and use the merged weight parameters as the weight parameters of the image encoder 300 to complete the parameter adjustment of the image encoder 300 of the SAM model;
[0053] Furthermore, this embodiment uses Vision Transformer (ViT) as the image encoder 300 of the SAM model. At the same time, the fine-tuning module 100 fine-tunes the model parameters through models such as LoRA (Low-Rank Adaptation of Large Language Models), Textual Inversion, and DreamBooth.
[0054] S3, the target detection encoding module 200 receives the image features output by the image decoder 300, and fuses the image features to obtain target detection data fusion features;
[0055] Specifically, in this embodiment, the target detection encoding module 200 extracts image features through a backbone feature extraction network (Backbone), and the backbone feature extraction network includes one or more network models such as the image encoder 300 of the SAM model, Resnet, VGG, and SwinTransformer;
[0056] Further, the target detection encoding module 200 in the present embodiment includes a feature fusion module 201 having a network model such as a feature pyramid (Feature Pyramid Networks, FPN), thereby the image features can be fused through the network model such as the feature pyramid of the feature fusion module 201 to obtain the target detection data fusion feature; preferably, the feature pyramid network model (FPN) includes one or more of a pyramid attention network (Pyramid Attention Network, PAN), a bidirectional weighted feature pyramid network (Bi-directional Feature Pyramid Networks, BiFPN), etc.;
[0057] S4, the target detection encoding module 200 performs target detection according to the target detection data fusion feature to obtain a target detection output result, and automatically outputs prompt information, wherein the target detection output result includes one or more of a target positioning box (box) and a target category (cls);
[0058] Specifically, the target detection encoding module 200 in this embodiment includes a target box 202, a classification head cls 203 and a prompt head prompt 204. The target box 202 and the classification head cls 203 are displayed and trained by the target detection head Head. The target detection data fusion features are further detected by the target box 202 and the classification head cls 203 to output the target positioning box (box) and the target category (cls) correspondingly, and prompt information is automatically outputted by the prompt head prompt 204. The prompt information may be an implicit feature vector.
[0059] S5. The mask decoder 400 of the SAM model receives the prompt information output by the target detection encoding module 200, and performs target segmentation according to the prompt information to output a predicted mask of the target, that is, a binary mask image;
[0060] S6. Repeat steps S2-S5 to perform iterative model training and construct the total loss function L at the same time. total , and L total =α*L mask +β*L reg +λ*L cls , and supervise the model training through the total loss function until the training meets the preset conditions, and then output the target segmentation model;
[0061] Among them, L mask It is the mask loss, which can be dice loss, and only needs to supervise the predicted mask result of the output; L reg is a regression loss function. In this embodiment, GIoU can be used as the regression loss function L reg , L cls is the classification loss function. In this embodiment, Focal Loss can be used as the classification loss function L cls , and the regression loss function L reg , classification loss function L cls They are collectively used as the loss of the target detection encoding module. α, β, and λ are adjustment coefficients and can be set according to actual needs. For example, in this embodiment, the values of α, β, and λ can be 1, 0.35, and 0.02, respectively.
[0062] During the iterative model training, the relevant parameters of the mask decoder 400 are also frozen, and the parameters are no longer updated. Only the fine-tuning module 100 and the target detection encoding module 200 are trained.
[0063] At the same time, when the fine-tuning module 100 is a LoRA fine-tuning module, the forward propagation process of the model is:
[0064] y=(W0+ΔW)x=W0x+BAx
[0065] Among them, x and y are the input feature map and the output feature map respectively, and x∈R d , y∈R d , d is the number of bits of the feature map, which is generally 64 or 128; W0 is the weight matrix of the image encoder 300 of the SAM model; ΔW is the weight matrix of the fine-tuning module, and ΔW=BA, B and A are two low-rank matrices respectively; when performing iterative model training, the weight matrix W0 of the image encoder 300 is frozen, and no gradient update is performed. Only matrix A and matrix B are trained, and the ranks of matrix A and matrix B can be set to 4. During model inference, since BA is merged into W0, the inference delay is not increased;
[0066] S7. Segment the target in the current image using the target segmentation model output in step S6, and output the target segmentation result, wherein the current image includes an infrared image.
[0067] from Figure 3 It can be seen that parts (a) and (c) are the original infrared images of two types of power equipment, respectively. After being processed by the target segmentation model in this embodiment, as shown in parts (b) and (d), the segmented power equipment can be accurately identified by the detection box, and the generalization performance is good.
[0068] The SAM model is usually trained on a large number of visible light images. Although the target segmentation effect on infrared images can be achieved by adjusting the weights of the SAM model through all parameters, the generalization performance is poor. In this embodiment, the fine-tuning module is used to fine-tune some parameters of the image encoder of the SAM model, so that a good target segmentation effect can be obtained on infrared images, and at the same time, the target segmentation model maintains the generalization ability, making the model particularly suitable for target segmentation of infrared images; at the same time, in this embodiment, the target detection encoding module replaces the prompt encoder (PromptEncoder) of the SAM model. The target detection encoding module can automatically generate and output prompt information while obtaining the target detection output result. Therefore, there is no need to obtain the segmentation result through the interactive clicking method of the prompt encoder (Prompt Encoder) in the SAM model, which greatly improves the target segmentation efficiency;
[0069] Furthermore, during the model training process, the relevant parameters of the image encoder and mask decoder are frozen and no longer updated. Only the fine-tuning module and the target detection encoding module need to be trained, which can greatly improve the model training efficiency.
[0070] Embodiment 2:
[0071] This embodiment provides a target segmentation system for implementing the target segmentation method described in Embodiment 1, such as Figure 2 As shown, it includes: a fine-tuning module 100, a target detection encoding module 200, an image encoder 300 and a mask decoder 400 of a SAM model, a loss function construction module and a target segmentation module;
[0072] The fine-tuning module 100 is used to adjust the parameters of the image encoder 300 of the SAM model, and input the pre-processed image P into the image encoder 300 whose parameters have been adjusted. The specific process is the same as step S2.
[0073] The image encoder 300 after parameter adjustment is used to receive the pre-processed image P and output image features;
[0074] The target detection encoding module 200 is used to receive the image features output by the image decoder 300, fuse the image features to obtain target detection data fusion features, perform target detection according to the target detection data fusion features to obtain target detection output results, and automatically output prompt information;
[0075] The mask decoder 400 is used to receive the prompt information output by the target detection encoding module 200, and perform target segmentation according to the prompt information to output a predicted mask of the target;
[0076] Loss function construction module, which is used to construct the total loss function L total , and supervise the model training through the total loss function until the training meets the preset conditions, and then output the target segmentation model finally used for target segmentation;
[0077] And a target segmentation module, which is used to segment the target in the current image according to the target segmentation model ultimately used for target segmentation, and output the target segmentation result.
[0078] Furthermore, the target detection encoding module 200 includes:
[0079] A feature fusion module 201, which has a network model such as Feature Pyramid Networks (FPN) for fusing image features;
[0080] The target box 202 is used to detect the fusion features of the target detection data to output a target positioning box;
[0081] The classification head cls 203 is used to detect the fusion features of the target detection data to output the target category (cls);
[0082] And the prompt header prompt 204 is used to automatically output prompt information.
[0083] In summary, the present invention fine-tunes some parameters of the image encoder of the SAM model through the fine-tuning module, so that it can also obtain a good target segmentation effect on infrared images, and at the same time, the target segmentation model maintains the generalization ability, so that the model is particularly suitable for target segmentation of infrared images; at the same time, the target detection encoding module replaces the prompt encoder of the SAM model, so that it can automatically generate and output prompt information while obtaining the target detection output result, so there is no need to obtain the segmentation result through the interactive clicking method of the prompt encoder in the SAM model, which greatly improves the target segmentation efficiency;
[0084] During the model training process, the relevant parameters of the image encoder and mask decoder are frozen and no longer updated. Only the fine-tuning module and target detection encoding module need to be trained, which can greatly improve the model training efficiency.
[0085] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principle of the present invention should be included in the protection scope of the present invention.
Claims
1. A target segmentation method based on the SAM model, characterized in that: The steps include: The preprocessed image is input into an image encoder whose parameters have been adjusted to output image features; Fusing image features to obtain target detection data fusion features; Perform target detection based on target detection data fusion features to obtain target detection output results and automatically output prompt information; Performing target segmentation according to the prompt information to output a predicted mask of the target; After the model iterative training is completed, the target segmentation model is output; And, the target in the current image is segmented through the output target segmentation model, and the target segmentation result is output.
2. The target segmentation method according to claim 1, characterized in that: The image encoder parameters of the SAM model are adjusted through the fine-tuning module.
3. The target segmentation method according to claim 2, characterized in that: The fine-tuning module is used to adjust the parameters of the image encoder of the SAM model, including the following steps: Freezing weight parameters of the image encoder; The weight parameters of the fine-tuning module are adjusted, and the adjusted weight parameters of the fine-tuning module are merged with the weight parameters of the image encoder, and the weight parameters obtained by merging are used as the weight parameters of the image encoder.
4. The target segmentation method according to claim 1, characterized in that: Image features are fused through the feature pyramid network model.
5. The target segmentation method according to claim 1, characterized in that: The target detection output result includes one or more of a target positioning frame and a target category.
6. The target segmentation method according to claim 1, characterized in that: When the model is iteratively trained, the total loss function L total Supervise model training, and L total =α*L mask +β*L reg +λ*L cls ; Among them, L mask is the mask loss, L reg is the regression loss function, L cls is the classification loss function, and α, β, and λ are adjustment coefficients.
7. The target segmentation method according to claim 1, characterized in that: The prompt information is received by the mask decoder of the SAM model, and the target is segmented according to the prompt information to output the predicted mask of the target.
8. The target segmentation method according to claim 7, characterized in that: The parameters of the mask decoder are frozen while the model is iteratively trained.
9. A target segmentation system, characterized in that: include: Fine-tuning module, object detection encoding module, image encoder of SAM model, mask decoder of SAM model, loss function construction module and object segmentation module; The fine-tuning module is used to adjust the parameters of the image encoder of the SAM model, and input the pre-processed image into the image encoder whose parameters have been adjusted; The image encoder whose parameters have been adjusted is used to receive the preprocessed image and output image features; The target detection encoding module is used to receive the image features output by the image decoder, fuse the image features to obtain target detection data fusion features, perform target detection according to the target detection data fusion features to obtain target detection output results, and automatically output prompt information; The mask decoder is used to receive prompt information output by the target detection encoding module, and perform target segmentation according to the prompt information to output a predicted mask of the target; Loss function construction module, which is used to construct the total loss function L total , and supervise the model training through the total loss function until the training meets the preset conditions, and then output the target segmentation model finally used for target segmentation; And a target segmentation module, which is used to segment the target in the current image according to the target segmentation model ultimately used for target segmentation, and output the target segmentation result.
10. The target segmentation system according to claim 9, characterized in that: The target detection encoding module includes: A feature fusion module, which is used to fuse image features; The target box is used to detect the fusion features of the target detection data to output the target positioning box; The classification head cls is used to detect the fusion features of the target detection data to output the target category; And the prompt header prompt, which is used to automatically output prompt information.