Fast and video-memory-friendly image segmentation method based on SAM model

Abstract

The invention relates to a rapid and video-memory-friendly image segmentation method based on an SAM model, and the method comprises the following steps: 1, building a U-SAM model based on the training of the SAM model; the U-SAM model is trained by the following steps: step S1, establishing an SAM model and a U-SAM model; a semantic amp of the U-SAM model is obtained through an image encoder of the knowledge distillation compression SAM model; the scale perception image encoder is used for improving the semantic perception capability of the SAM; and S2, freezing related parameters of the SAM model, retaining prompt segmentation capability, training a double-token mask decoder of the U-SAM model, and solving the problem that prompt embedding is difficult to control. 2, using the U-SAM model to realize prompt segmentation and general image segmentation; therefore, prompt and universal segmentation tasks can be realized through a rapid and video-friendly framework.

Description

Technical Field

[0001] The invention relates to the technical field of computer vision, and in particular to a fast and video memory friendly image segmentation method based on a SAM model. Background Art

[0002] From image scene parsing to panoramic image segmentation, and now to general image segmentation, processing multiple segmentation tasks simultaneously has been a long-term hot topic in the field of computer vision. General image segmentation mainly includes instance segmentation, semantic segmentation, and panoramic segmentation. The main difference between these tasks is the method of grouping pixels. Early panoramic segmentation models merged the masks of instance segmentation models and semantic segmentation models through post-processing, resulting in problems such as model complexity and computational redundancy. Subsequently, the detection visual self-attention model (DETR) solved this problem by jointly modeling countable objects (thing) and uncountable objects (stuff) through queries, trying to achieve instance, semantic, and panoramic segmentation in a single model. However, its performance still lags behind segmentation models trained for specific tasks. Later models improved performance in various ways, but none of them could unify these tasks through a single training process. Therefore, current work focuses on expanding the number of tasks or establishing a unified training framework for multiple tasks. Among them, the Segment Anything Model (SAM) pre-trained on a large-scale dataset (SA-1B) has become the basic model for image segmentation, showing strong generalization performance in a variety of image segmentation tasks.

[0003] Migrating SAM to specific tasks and optimizing the specific performance of SAM models are currently hot topics. The former improves the performance of SAM on tasks with weaker generalization capabilities through fine-tuning or adapter techniques. The latter solves the challenge of applying the interactive segmentation framework to other segmentation tasks by expanding the structure of SAM.

[0004] Through a category-independent training method, SAM can segment categories that do not exist in the training data, but it also causes SAM to not pay attention to the semantic information of the target during the learning process. Therefore, SAM can be applied to any scene that needs to predict the target mask, but it cannot adapt to traditional image segmentation tasks that need to simultaneously predict the masks and corresponding categories of all objects in the image, such as instance segmentation, semantic segmentation, and panoptic segmentation.

[0005] In order to solve the above defects, the existing methods are mainly divided into two categories.

[0006] First, retraining a SAM-based model is the simplest way to solve the semantic missing problem, but it requires a lot of computing resources (256 A100 graphics cards are required for 68 hours of training), so it is less feasible.

[0007] Secondly, the object prompter extracts semantic information from the low-level features of SAM and generates corresponding prompt embeddings for predicting categories and masks. This method can greatly reduce the required computing resources, but since the low-level features of the image extracted by the image encoder based on the visual self-attention model (ViT) have high similarity with the high-level features, the low-level features of SAM also lack semantic information. In addition, the object prompter method cannot determine the validity of the generated prompt embedding during the learning process, and the segmentation performance of SAM depends heavily on the quality of the prompt embedding. Different types of prompt embeddings correspond to different forms and meanings. Using different types or numbers of prompts for the same target will result in masks of different quality. At the same time, since the goal of SAM is to finely segment a single target based on the prompt, the special structural design makes it impossible to process multiple targets in an image at the same time. It can only sacrifice speed to process one by one or sacrifice video memory to process multiple targets at the same time by copying the image. This has high requirements on computing resources or video memory, so it is less practical. Summary of the invention

[0008] The purpose of the present invention is to provide a fast and memory-friendly image segmentation method based on the SAM model, which adopts a semantically aware and memory-efficient Universal Segment Anything Model (U-SAM) to achieve both prompt and traditional image segmentation tasks.

[0009] In order to achieve the above object, the solution of the present invention is: A fast and memory-friendly image segmentation method based on a SAM model, the image segmentation method comprising the following steps: The first step is to establish a U-SAM model based on SAM model training; the U-SAM model is trained using the following steps: Step S1: Establish a SAM model and a U-SAM model; compress the image encoder of the SAM model through knowledge distillation to obtain a semantic & scale-aware image encoder of the U-SAM model; Step S2: Freeze the relevant parameters of the SAM model, retain its original hint segmentation capability, and obtain the hint encoder of the U-SAM model; and train the twin token mask decoder of the U-SAM model on the downstream dataset COCO; The second step is to use the U-SAM model to achieve suggestible segmentation and universal image segmentation; The hintable segmentation process is as follows: given an input image and an artificial hint, the semantic & scale-aware image encoder extracts and outputs global segmentation features from the input image to a twin token mask decoder; the hint encoder encodes the artificial hint into a hint embedding and outputs it to a twin token mask decoder; finally, the twin token mask decoder generates a corresponding hint mask using the global segmentation features and the hint embedding, and outputs the segmentation result; The general image segmentation process is as follows: given an input image, a semantic & scale-aware image encoder extracts and outputs global segmentation features from the input image to a twin token mask decoder, and simultaneously extracts multi-scale semantic features of the input image intermediate layer features and outputs them to a twin token mask decoder; subsequently, the twin token mask decoder extracts feature information that is beneficial to the segmentation target from the global segmentation features, and fuses the positional encoding of the encoder in the twin token mask decoder to generate a unified token for general segmentation, and then uses the global segmentation features and the unified token to generate a corresponding image segmentation mask, and outputs the segmentation result; simultaneously, the twin token mask decoder extracts feature information that is beneficial to target recognition from the multi-scale semantic features, and fuses the feature information that is beneficial to the segmentation target to obtain the target category of the target.

[0010] Furthermore, the segmentation process can be expressed as: ; in, For manual prompts, Indicates Tips: To embed the prompt, Indicates Tips The corresponding prompt embedding; is the input image, ; ; Indicates the prompt mask. ; Represents a prompt encoder; 2IE stands for semantic & scale-aware image encoder; represents a Gemini Token Mask Decoder; in, and Represent the height and width of the input image respectively, and Represent the height and width of the image feature map respectively, Represents the number of manual prompts entered. The dimension representing the image features; represents the field of real numbers; The general image segmentation process is expressed as: ; in, is the input image, ; ; Represent multi-scale semantic features; represents the target category of the target, ; represents the image segmentation mask, ; Representation semantic & scale-aware image encoder; represents a Gemini Token Mask Decoder; in, and Represent the height and width of the input image respectively, and Represent the height and width of the image feature map respectively, The dimension representing the image features; represents the field of real numbers; represents the number of unified tokens in the Gemini Token Mask Decoder, Represents the number of categories contained in the downstream dataset.

[0011] Furthermore, the step S1 includes the following specific steps: Step S1.1, on the distillation dataset SA-1B, perform data augmentation on the input image, scale it according to the aspect ratio, and randomly flip it horizontally; In step S1.2, the enhanced image is input into the SAM model image encoder and the U-SAM model semantic & scale-aware image encoder for knowledge distillation.

[0012] Furthermore, in step S1.2, a semantic & scale-aware image encoder based on a residual network structure is used as a student model, and the image encoder of the SAM model is used as a teacher model. The knowledge of the image encoder of the SAM model is transferred to the semantic & scale-aware image encoder of the U-SAM model through knowledge distillation, and a simple mean square error loss is used to supervise the distillation process of the SAM image encoder and the U-SAM semantic & scale-aware image encoder.

[0013] Furthermore, in step S1.2, the teacher model extracts features from the data-augmented image through its image encoder, and the student model extracts global segmentation features through the semantic & scale-aware image encoder; then the knowledge distillation process is supervised by calculating the mean square error loss of the extracted features of both.

[0014] Furthermore, the semantic & scale-aware image encoder includes a residual neural network and a neck; wherein the residual neural network includes four residual blocks Res-1, Res-2, Res-3, and Res-4, which are used to extract intermediate layer features containing semantic information at different levels; the neck is composed of a convolutional layer and a hole convolutional layer, which are used for channel dimension alignment and increasing the receptive field respectively; The twin token mask decoder comprises a bimolecular decoder and a mask decoder, wherein the bimolecular decoder comprises a token generator and a semantic perceiver; The step S2 comprises the following specific steps: Step S2.1, on the downstream dataset COCO, perform data augmentation on the input image, scale it according to the aspect ratio and randomly flip it horizontally, and input the augmented image into the semantic & scale-aware image encoder in batches; then extract the global segmentation features of the image through the semantic & scale-aware image encoder, and extract the multi-scale semantic features of the image through the intermediate layer features output by the Res-1, Res-2, Res-3, and Res-4 residual blocks of the semantic & scale-aware image encoder; align the channel dimensions of the output partial convolution layers of the four residual blocks and increase the receptive field through the hollow convolution layer; use convolution layers with different receptive fields to further optimize the outputs of the semantic & scale-aware image encoder at different stages; Step S2.2, the twin-molecule decoder of the twin token decoder uses the token generator and the semantic perceptron to extract feature information that is beneficial to target recognition and feature information that is beneficial to target segmentation from the multi-scale semantic features and global segmentation features of the image; and obtains the target category and the image segmentation mask through the semantic perceptron and the mask decoder respectively; In step S2.3, the segmentation loss is calculated using the given standard segmentation mask and the image segmentation mask predicted by the U-SAM model, and the classification loss is calculated using the given standard target category and the target category predicted by the U-SAM model to supervise the output of the U-SAM model. Finally, the parameters of the U-SAM model are updated through the adaptive stochastic gradient descent method to achieve segmentation and recognition of the input image.

[0015] Further, in step S2.2, the token generator includes an encoder for enhancing the global segmentation feature and a decoder 1 for generating a unified token; the encoder first extracts feature information that is beneficial to the segmentation target from the global segmentation feature, and uses its position encoding to enhance the spatial position information, and then inputs the decoder 1 to automatically generate a unified token, and further outputs an image segmentation mask through a mask decoder according to the unified token; The semantic sensor includes a feature pyramid network for extracting semantic features and a decoder 2 for obtaining target categories; the feature pyramid network extracts feature information that is beneficial to target recognition from multi-scale semantic features, and then combines the feature information that is beneficial to target segmentation to obtain the target category of the corresponding segmented target through the decoder 2.

[0016] Furthermore, in step S2.3, the U-SAM model uses binary cross entropy loss and Dice loss to represent the segmentation loss: , respectively set and is the weighted coefficient of binary cross entropy loss and Dice loss; the U-SAM model uses multivariate cross entropy loss to represent the classification loss , is the weighting coefficient of the multivariate cross entropy loss; the final overall loss is the sum of the segmentation loss and the classification loss: ; in is the binary cross entropy loss, , represents the number of samples, Indicates whether the i-th sample is the target category, Indicates the probability that the i-th sample is the target category; It is the Dice loss, which has good performance in scenarios where positive and negative samples are severely unbalanced. , and Represent the given target segmentation mask and the predicted segmentation mask respectively; is the multivariate cross entropy loss, , and Indicates whether the i-th sample is a background class and the probability of being a background class, and They respectively represent whether the i-th sample is the j-th target category and the probability of being the j-th target category, is the loss coefficient of the background class, and c is the number of categories in the downstream dataset.

[0017] The present invention also provides a computer system, which includes a processor and a memory, wherein the memory stores at least one instruction, at least one program, a code set or an instruction set, and the at least one instruction, the at least one program, the code set or the instruction set is loaded and executed by the processor to implement the fast and video memory friendly image segmentation method based on the SAM model.

[0018] The present invention also provides a computer-readable storage medium, wherein the storage medium stores at least one instruction, at least one program, a code set or an instruction set, and the at least one instruction, the at least one program, the code set or the instruction set are loaded and executed by a processor to implement the fast and video memory friendly image segmentation method based on the SAM model.

[0019] After adopting the above technical solution, a semantically aware and memory-efficient unified segmentation model (Universal Segment Anything Model, U-SAM) of the present invention can realize prompts and general segmentation tasks through a fast and video memory friendly framework. The present invention proposes a semantic & scale-aware image encoder (Semantic & Scale-aware Image Encoder, S2IE) to address the problem that the low-level features of SAM lack semantic information, and improves the semantic perception ability of SAM by integrating multi-scale feature extraction capabilities. The single prompt reasoning speed of the present invention is 10 times faster than SAM, and the performance on multiple zero-shot tasks exceeds MobileSAM. In addition, the present invention proposes a twin token mask decoder (Twin Token Mask Decoder, T2MD) to solve the problem of difficulty in controlling the effectiveness of prompt embedding and the need to sacrifice video memory to process multiple targets. This method can not only avoid the difficulty of directly learning prompt embedding, but also save a lot of video memory by replacing image replication with token replication. At the same time, the twin token mask decoder also decouples segmentation and classification tasks, while maintaining segmentation performance and enhancing classification capabilities. Compared with the target prompter, the twin token mask decoder reduces the training video memory by 54% and speeds up by 43% while the performance is basically unchanged. In summary, the present invention has the following advantages: The semantically-aware and memory-efficient unified segmentation everything model (U-SAM) proposed in this paper can achieve both suggestive and general segmentation tasks in a compressed and unified framework. Compression refers to the reduction of model parameters through distillation; unification refers to the model being able to complete both suggestive and general segmentation tasks.

[0020] The semantic & scale-aware image encoder S2IE proposed in the present invention can effectively solve the problem that the low-level features of SAM lack semantic information. At the same time, the proposed twin token mask decoder T2MD effectively solves the difficulty of learning cue embedding for the target cue detector and the need to sacrifice video memory to achieve multi-target segmentation.

[0021] The present invention has conducted extensive experiments on instance segmentation, semantic segmentation, panoptic segmentation, and interactive segmentation, verifying the creativity of the present invention over the prior art. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] Figure 1 The difference between the U-SAM model proposed in the present invention and the existing methods; Figure 2 The overall framework diagram of the U-SAM model proposed in the present invention; Figure 3 An interactive segmentation visualization comparison diagram of the U-SAM model proposed in the present invention and other models; Figure 4 A visual comparison diagram of the U-SAM model and the SAM model proposed in the present invention; Figure 5 The general image segmentation visualization result of the U-SAM model proposed in the present invention; Figure 6 It is a flowchart of the image segmentation method proposed by the present invention; Figure 7 This is a flow chart of the U-SAM model training method proposed in the present invention. DETAILED DESCRIPTION

[0023] In order to enable those skilled in the art to better understand the present application, the technical solutions in the embodiments of the present application will be clearly and completely described below in conjunction with the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, rather than all of the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without making creative work are within the scope of protection of the present application.

[0024] The present invention utilizes the powerful image segmentation capability and zero-sample migration performance of the SAM model to obtain a fast and memory-friendly image segmentation model, thereby overcoming the problem that the SAM model lacks output category capability and cannot be applied to general image segmentation tasks. Figure 1 The existing methods shown above require a lot of computing resources through retraining or using complex structures, and use video memory to exchange speed to process multiple targets at the same time. This enables the model to quickly adapt to different image data distributions and overcomes the shortcomings of reusing specific image data for fine-tuning for different problems.

[0025] The SAM model is an existing segmentation model, which consists of three parts: image encoder, prompt encoder and mask decoder. The SAM model is trained on a large dataset using interactive learning and prompt learning strategies. When using the SAM model, the user needs to provide prompt information for the target to be segmented, such as points inside the target, target bounding box or graffiti-style mask. The SAM model has good zero-sample generalization ability and can be migrated to various image segmentation tasks.

[0026] This embodiment establishes a unified segmentation everything model (U-SAM model) based on the segmentation everything model (SAM model) to achieve a fast and memory-friendly image segmentation method.

[0027] like Figure 2 The U-SAM model includes a hint encoder, a semantic & scale-aware image encoder (Semantic & Scale-aware Image Encoder, S2IE) and a twin token mask decoder (Twin Token Mask Decoder, T2MD).

[0028] The semantic & scale-aware image encoder includes four residual blocks Res-1, Res-2, Res-3, Res-4.

[0029] The twin token mask decoder includes a twin molecule decoder and a mask decoder.

[0030] The image segmentation method comprises the following steps: The first step is to build a U-SAM model based on SAM model training; see Figure 7 , the training method of the U-SAM model can be divided into the following steps: Step S1: Establish a SAM model and a U-SAM model; on the distillation dataset SA-1B, compress the image encoder of the SAM model through knowledge distillation to obtain the semantic & scale-aware image encoder of the U-SAM model; Step S2: Freeze the relevant parameters of the SAM model, retain its original hint segmentation capability, and obtain the hint encoder of the U-SAM model; and train the twin token mask decoder of the U-SAM model on the downstream dataset COCO.

[0031] The second step is to use the U-SAM model to achieve suggestible segmentation and universal image segmentation; See also Figure 6 , which can be shown as follows:

[0032] ; Specifically, it is the same as the suggestible segmentation process of the visual segmentation basic model SAM. Given an input image and manual prompts , first semantic & scale-aware image encoder Extract global segmentation features from the input image , then prompt the encoder Encoding artificial cues into cue embeddings , and finally the twin token mask decoder Generate corresponding hint masks using global segmentation features and hint embedding , and output the segmentation result.

[0033] in, For manual prompts, Indicates Tips: To embed the prompt, Indicates Tips The corresponding prompt embedding; and Represent the height and width of the input image respectively, and Represent the height and width of the image feature map respectively, Represents the number of manual prompts entered. Represents the dimension of image features.

[0034] The general image segmentation process can be expressed as:

[0035] .

[0036] Different from hintable segmentation, general image segmentation only gives an input image, while the semantic & scale-aware image encoder outputs global segmentation features and extracts multi-scale semantic features of the image intermediate layer features. The twin token mask decoder then extracts feature information that is beneficial to the segmentation target from the global segmentation feature. , and fused with position encoding Generate unified tokens for universal segmentation , and then use the global segmentation features and unified tokens to generate the corresponding image segmentation mask , and output the segmentation result.

[0037] At the same time, the twin token mask decoder extracts feature information that is beneficial to target recognition from multi-scale semantic features. , and fuses the feature information that is beneficial to segmenting the target to obtain the target category .

[0038] in, represents the number of unified tokens in the Gemini Token Mask Decoder, Represents the number of categories contained in the dataset.

[0039] The training method of the U-SAM model is as follows: Step S1 includes the following specific steps: In step S1.1, on the distillation dataset SA-1B, the input image is augmented, scaled according to the aspect ratio, and randomly flipped horizontally.

[0040] Step S1.2, the enhanced image is input into the SAM model image encoder and the U-SAM model semantic & scale-aware image encoder for knowledge distillation; the semantic & scale-aware image encoder based on the residual network structure is used as the student model, and the SAM image encoder is used as the teacher model. The knowledge of the SAM image encoder is transferred to the semantic & scale-aware image encoder of the U-SAM model through knowledge distillation, and a simple mean square error (MSE) loss is used ( ) to supervise the distillation process of SAM image encoder and U-SAM semantic & scale-aware image encoder.

[0041] The purpose of the semantic & scale-aware image encoder of the present invention is to enhance the semantic perception ability of the SAM model by integrating multi-scale feature extraction capabilities through knowledge distillation. On the one hand, the high- and low-level features extracted by the convolutional neural network have lower feature similarity than the visual self-attention model, so the low-level feature distribution will not tend to the visual basic model features with semantic loss during the distillation process. On the other hand, the residual neural network (ResNet) structure can extract good multi-scale features, which is conducive to the positioning and recognition of targets.

[0042] Specifically, in step S1.2, given an input image, the teacher model SAM passes through its image encoder ( ) Extract features from the data augmented image , the student model U-SAM extracts global segmentation features through the above semantic & scale-aware image encoder The knowledge distillation process is then supervised by calculating the mean square error loss of the extracted features. and Represent the height and width of the image feature map respectively, Represents the dimension of image features, Represents the field of real numbers.

[0043] Step S2 can be roughly divided into the following three steps: Step S2.1: On the downstream dataset COCO, perform data augmentation on the input image, scale it according to the aspect ratio and randomly flip it horizontally, and input the augmented image into the semantic & scale-aware image encoder in batches. Then, the global segmentation features of the image are extracted through the semantic & scale-aware image encoder, and the multi-scale semantic features of the image are extracted through the intermediate layer features output by the Res-1, Res-2, Res-3, and Res-4 residual blocks; Specifically, in step S2.1, the semantic & scale-aware image encoder includes a residual neural network (ResNet) and a neck (Neck). The residual neural network includes four residual blocks Res-1, Res-2, Res-3, and Res-4, which are respectively used to extract semantic information at four different levels. The neck consists of a convolutional layer and a hole convolutional layer, which are respectively used for channel dimension alignment and increasing the receptive field, so as to further optimize the output of different residual blocks of the semantic & scale-aware image encoder and obtain multi-scale semantic features and global segmentation features of the image.

[0044] Specifically, Figure 2 In the upper right part, given the enhanced image, the semantic & scale-aware image encoder first extracts the multi-scale semantic features of the image through a residual neural network. ,in Represent the output features of the four residual blocks of the residual network respectively. Then the global segmentation feature is obtained through the neck .

[0045] The twin token mask decoder comprises a bimolecular decoder and a mask decoder, wherein the bimolecular decoder comprises a token generator and a semantic perceiver; In step S2.2, the obtained multi-scale semantic features and global segmentation features are input into a Twin Token Mask Decoder (T2MD), which includes a dual-branch decoder (D.Dec) and a mask decoder (M.Dec). The dual-branch decoder uses a token generator and a semantic perceptron to extract feature information that is beneficial to target recognition and feature information that is beneficial to target segmentation from the multi-scale semantic features and the global segmentation features, respectively; then the obtained feature information of target recognition is input into the semantic perceptron to obtain the target category, and the feature information of the segmented target is passed through the mask decoder to obtain the image segmentation mask predicted by the U-SAM model, that is, the image segmentation mask predicted by the U-SAM model.

[0046] In order to avoid the difficulty of directly learning the hint embedding, the twin token mask decoder of the present invention abandons the method of generating hint embedding from image features through the target hinter, and adopts a set of additional token generators and semantic perceptrons to output the segmentation mask and target category of the image. Because the traditional image segmentation task is a global division of the image, it is independent of the user's hint, and replacing image replication with token redundancy can effectively reduce the requirements for video memory size during training and reasoning. In addition, in order to reduce the impact of semantic information on the original segmentation ability of the model, the twin token mask decoder decouples the segmentation and classification tasks through a bimolecular decoder.

[0047] Specifically, the upper part of the bimolecular decoder is called the Token Generator (T.Gen), which contains an encoder ( ) and a decoder 1 for generating a unified token ( ). The generated unified tokens are then passed through the mask decoder to further generate image segmentation masks. The lower part of the bimolecular decoder is called the semantic perception (S.Per), which contains a feature pyramid network ( ) and decoder 2 for obtaining the target category ( ). To further improve the performance of the model, we extract high-quality semantic information to enhance the features of the input semantic perceptron, and use positional encoding to enhance the features of the input token generator. Compared with the target prompter, the twin token mask decoder can reduce 54% of training memory and speed up the inference by 43% while maintaining basically the same performance.

[0048] like Figure 2 Bottom right, given the global segmentation features output by the semantic & scale-aware image encoder ( ) and residual neural network ( )Multi-scale semantic features output by the intermediate layer ( ), encoder( ) First, extract feature information that is beneficial to the segmentation target from the global segmentation features ( ), and use positional encoding ( ) enhances the spatial position information, and then inputs the decoder 1 to automatically generate a unified token, and further outputs the image segmentation mask through the mask decoder according to the unified token ( At the same time, the feature pyramid network extracts feature information that is beneficial to target recognition from multi-scale semantic features ( ), and then combined with the feature information that is conducive to target segmentation, the target category corresponding to the segmented target is obtained through decoder 2 ( ).

[0049] In step S2.3, the segmentation loss is calculated using the given standard segmentation mask and the image segmentation mask predicted by the U-SAM model, and the classification loss is calculated using the given standard target category and the target category predicted by the U-SAM model. Finally, the parameters of the U-SAM model are updated using the adaptive stochastic gradient descent method.

[0050] Specifically, in step S2.3, the cross entropy loss and Dice loss of the segmentation mask predicted by the U-SAM model and the given standard segmentation mask are calculated, and the cross entropy loss of the target category predicted by the U-SAM model and the given standard target category are calculated to supervise the output of the U-SAM model, and the U-SAM model parameters are updated using the adaptive stochastic gradient descent method to achieve segmentation and recognition of the input image; For the segmentation loss, the U-SAM model uses binary cross entropy loss and Dice loss instead of focal loss: , respectively set and is the weighted coefficient of binary cross entropy loss and Dice loss.

[0051] For classification loss, multivariate cross entropy loss is used to express , is the weighting coefficient of the multivariate cross entropy loss; the final overall loss is the sum of the segmentation loss and the classification loss: ; in is the binary cross entropy loss, , represents the number of samples, Indicates whether the i-th sample is the target category, Indicates the probability that the i-th sample is the target category; It is the Dice loss, which has good performance in scenarios where positive and negative samples are severely unbalanced. , and Represent the given target segmentation mask and the predicted segmentation mask respectively; is the multivariate cross entropy loss, , and Indicates whether the i-th sample is a background class and the probability of being a background class, and They respectively represent whether the i-th sample is the j-th target category and the probability of being the j-th target category, is the loss coefficient of the background class, and c is the number of categories in the downstream dataset. At the same time, in order to save video memory, when calculating the segmentation loss, not all points on the mask are involved in the loss calculation, but a certain number of points are randomly sampled to calculate the segmentation loss.

[0052] In order to verify the effectiveness and advancement of this method, this embodiment verifies the proposed fast and memory-friendly image segmentation method based on the SAM model. It should be noted that for the aforementioned method embodiments, for the sake of simplicity, they are all expressed as a series of action combinations, but those skilled in the art should know that the present invention is not limited to the described order of actions, because according to the present invention, some steps can be performed in other orders or simultaneously.

[0053] The image encoder of the SAM model has a vision-oriented inductive bias, which lays the foundation for the robust segmentation and generalization capabilities of the SAM model, but the associated computational overhead causes it to process images very slowly. Compared with the serial input adopted in promptable segmentation, general image segmentation is usually parallel input. Therefore, when extended to general image segmentation, the speed and performance of the SAM model are affected by the single prompt inference speed and the overall performance of the model.

[0054] Table 1: Running speed (ms / image) of our model compared with other models. E: All modes of SAM. *: 32×32 is the default setting of SAM.

[0055]

[0056] Table 1 gives a comparison of the running speed of the U-SAM model with other models on a single NVIDIA GeForce RTX 3090 GPU. The number of parameters of the U-SAM model is 93.0% lower than that of SAM-H and only 48.3% of that of SAM-B. At the same time, the single-prompt reasoning running speed of the present invention is 10 times faster than that of SAM-H and 2.7 times faster than that of SAM-B, reaching a speed similar to that of the Fast Segment Everything model (FastSAM) and the Mobile Segment Everything model (MobileSAM). It is worth noting that whether it is SAM-H with the slowest single-prompt reasoning speed or MobileSAM with the fastest single-prompt reasoning speed, when using the 64×64 "split everything" method, the difference in average single-prompt running speed is only 0.1-0.2 milliseconds. This is because during the interaction process, the image only needs to pass through the image encoder once, so that its reasoning speed is amortized in real time.

[0057] Table 2. Comparison of U-SAM with other models on various zero-shot tasks. None of the compared models have been trained on the corresponding tasks, nor have they seen any images or annotations from the relevant datasets.

[0058]

[0059] The global best (Optimal Dataset Scale, ODS), single image best (Optimal ImageScale, OIS), average precision (Average Precision, AP) and recall rate (Recall at50% precision, R50) of U-SAM in edge detection are evaluated on the Berkeley Segmentation Data Set (BSDS500).

[0060] We first generate a mask probability map using the SAM model, and then use non-maximum suppression (NMS) to filter redundant masks. Subsequently, the Sobel operator is used to generate an edge probability map, and finally edge non-maximum suppression (edge ​​NMS) is used for post-processing.

[0061] As shown in Table 2, U-SAM is closer to SAM-H than FastSAM and MobileSAM in all indicators, and its numerical distribution is basically consistent with MobileSAM and SAM-H.

[0062] We evaluate the average recall (AR) of our method on object proposal generation on a large-scale fine-grained vocabulary instance segmentation (LVIS) dataset.

[0063] Similar to SAM, we do not perform image cropping, nor do we use the intersection of union (IoU) and stability score to filter low-quality masks. We only filter redundant masks through the remaining parameters. If the number of masks exceeds one thousand during testing, we take the one thousand masks with the largest mean IoU and stability score as the result.

[0064] Table 2 verifies that the average recall rate AR of the present invention and MobileSAM is relatively lower than that of SAM, but we find that the preference of the present invention for masks of different sizes is basically consistent with that of SAM, and the difference in the average recall rate AR of masks of various sizes is basically the same. MobileSAM has a stronger preference for large targets, so the average recall rate AR of large masks will be relatively higher.

[0065] At the same time, the average precision (AP) of the present invention on instance segmentation is evaluated on the LVIS dataset. Similar to SAM, we first generate bounding boxes and corresponding target categories through a detection model based on visual self-attention (ViTDet), and then use the detection box (the target bounding box generated by the detection model) as a hint for the present invention to generate instance segmentation masks.

[0066] Table 2 verifies that the performance of the present invention is closer to SAM than FastSAM and MobileSAM. At the same time, we find that MobileSAM has a stronger bias towards large targets, while the present invention is more friendly to small targets. Therefore, on both datasets, the average precision (AP-S) of the present invention for small targets exceeds that of MobileSAM, but the average precision (AP-L) of the present invention for large targets is lower than that of MobileSAM.

[0067] Table 3. Comparison of U-SAM on COCO with other object cue-based models. Training time on Tesla V100 GPUs for test data is in days.

[0068]

[0069] Table 3 further compares the performance of the present invention with other target prompting models on the COCO dataset. It is worth noting that the full remote sensing prompter (RSPrompter) requires too much memory, and the RSPrompter used here is a parameter-adjusted version. We retain all components of RSPrompter, but reduce its training memory consumption by adjusting parameters. Specifically, we adjust the output of multi-scale features to 4, and reduce the number of encoder layers of the pixel decoder to 1, and the number of decoder layers of the mask decoder to 4.

[0070] Compared with RSPrompter, the present invention shows significant advantages in all aspects. Using only 10% of the parameters, the performance in instance, semantic and panoptic segmentation tasks is improved by 5.2 (AP), 6.3 (MIoU) and 6.2 (PQ), respectively, while achieving a 13-fold speed increase. In addition, compared with RSPrompters based on SAM-B or SAM-L, the performance advantage of the present invention is even more obvious. Its biggest advantage is that it reduces memory requirements by 61% and significantly reduces training time.

[0071] In addition, the present invention aligns the distribution of the image encoder output features of its lightweight feature extraction network and the SAM model through knowledge distillation. Figure 3As shown in the figure, it is a visual comparison diagram of U-SAM and other models. The red box emphasizes the areas that the present invention successfully segmented but other models missed, such as the road in the first row, the stone platform in the second row, and the trees in the last row. Most of these areas segmented by the present invention belong to uncountable (stuff) targets, which to some extent shows that compared with MobileSAM and FastSAM, U-SAM's recognition ability for uncountable targets is closer to SAM. The blue box emphasizes some SAM biases learned by the present invention.

[0072] Specifically, both SAM and the present invention believe that the cloud and sky in the upper right corner of the first row and the two pieces of ground in the third row do not belong to the same entity, while FastSAM and MobileSAM believe that they belong to the same entity. On the other hand, the present invention is consistent with SAM in the scratches in the sky in the third row and the segmentation of the building in the second row. These all reflect that compared with FastSAM and MobileSAM, the segmentation preference of the present invention is more consistent with that of SAM, which further verifies the conclusion of the experiment.

[0073] Figure 4 It is a visual comparison chart of U-SAM and SAM, which verifies that the segmentation results of the present invention are more consistent with SAM than FastSAM and MobileSAM. However, due to the image encoder that compresses SAM through knowledge distillation, the present invention also loses some performance. The first row shows that the present invention generates the same segmentation results as SAM, so for some images, the segmentation results of the present invention are basically consistent with SAM. The second row verifies that the present invention cannot successfully segment smaller countable (things) targets, such as the person on the left of the second row and the cup on the right. This shows to a certain extent that the instance segmentation performance of the present invention may not be ideal. The third row reflects that the present invention is more inclined to segment adjacent targets of the same type into a whole, such as the mountains and trees on the left of the third row, and the buildings and stairs on the right. Most of these areas are uncountable targets, so this property is conducive to completing general image segmentation.

[0074] Figure 5 The general image segmentation visualization results of U-SAM on downstream datasets are demonstrated, which can verify that the present invention is effective on downstream datasets.

Claims

1. A fast and memory-friendly image segmentation method based on the SAM model, characterized by: The image segmentation method comprises the following steps: The first step is to establish a U-SAM model based on SAM model training; the U-SAM model is trained using the following steps: Step S1: Establish a SAM model and a U-SAM model; compress the image encoder of the SAM model through knowledge distillation to obtain a semantic & scale-aware image encoder of the U-SAM model; Step S2: Freeze the relevant parameters of the SAM model, retain its original hint segmentation capability, and obtain the hint encoder of the U-SAM model; and train the twin token mask decoder of the U-SAM model on the downstream dataset COCO; The second step is to use the U-SAM model to achieve suggestible segmentation and universal image segmentation; The hintable segmentation process is as follows: given an input image and a manual hint, a semantic & scale-aware image encoder extracts and outputs global segmentation features from the input image to a twin token mask decoder; The prompt encoder encodes the artificial prompt into a prompt embedding and outputs it to the Gemini Token Mask Decoder; finally, the Gemini Token Mask Decoder uses the global segmentation features and the prompt embedding to generate the corresponding prompt mask and output the segmentation result; The general image segmentation process is as follows: given an input image, a semantic & scale-aware image encoder extracts and outputs global segmentation features from the input image to a twin token mask decoder, and simultaneously extracts multi-scale semantic features of the intermediate layer features of the input image and outputs them to a twin token mask decoder; Subsequently, the twin token mask decoder extracts feature information that is beneficial to target segmentation from the global segmentation features, and fuses the positional encoding of the encoder in the twin token mask decoder to generate a unified token for universal segmentation. The global segmentation features and the unified token are then used to generate the corresponding image segmentation mask and output the segmentation result. At the same time, the twin token mask decoder extracts feature information that is beneficial to target recognition from multi-scale semantic features, and fuses the feature information that is beneficial to target segmentation to obtain the target category of the target.

2. A fast and memory-friendly image segmentation method based on the SAM model according to claim 1, characterized in that: The segmentation process can be expressed as: ; in, For manual prompts, Indicates Tips: To prompt embedding, Indicates Tips The corresponding prompt embedding; is the input image, ; ; Indicates the prompt mask. ; Represents a prompt encoder; Representation semantic & scale-aware image encoder; represents a Gemini Token Mask Decoder; in, and Represent the height and width of the input image respectively, and Represent the height and width of the image feature map respectively, Represents the number of manual prompts entered, The dimension representing the image features; represents the field of real numbers; The general image segmentation process is expressed as: ; in, is the input image, ; ; Represent multi-scale semantic features; represents the target category of the target, ; represents the image segmentation mask, ; Representation semantic & scale-aware image encoder; represents a Gemini Token Mask Decoder; in, and Represent the height and width of the input image respectively, and Represent the height and width of the image feature map respectively, The dimension representing the image features; represents the field of real numbers; represents the number of unified tokens in the Gemini Token Mask Decoder, Represents the number of categories contained in the downstream dataset.

3. The fast and video memory friendly image segmentation method based on the SAM model according to claim 1, characterized in that: The step S1 comprises the following specific steps: Step S1.1, on the distillation dataset SA-1B, perform data augmentation on the input image, scale it according to the aspect ratio, and randomly flip it horizontally; In step S1.2, the enhanced image is input into the SAM model image encoder and the U-SAM model semantic & scale-aware image encoder for knowledge distillation.

4. The fast and memory-friendly image segmentation method based on the SAM model according to claim 3, characterized in that: In step S1.2, the semantic & scale-aware image encoder based on the residual network structure is used as the student model, and the image encoder of the SAM model is used as the teacher model. The knowledge of the image encoder of the SAM model is transferred to the semantic & scale-aware image encoder of the U-SAM model through knowledge distillation, and a simple mean square error loss is used to supervise the distillation process of the SAM image encoder and the U-SAM semantic & scale-aware image encoder.

5. The fast and video memory friendly image segmentation method based on the SAM model according to claim 4, characterized in that: In step S1.2, the teacher model extracts features from the data-augmented image through its image encoder, and the student model extracts global segmentation features through the semantic & scale-aware image encoder; The knowledge distillation process is then supervised by calculating the mean squared error loss of the extracted features.

6. The fast and video memory friendly image segmentation method based on the SAM model according to claim 1, characterized in that: The semantic & scale-aware image encoder includes a residual neural network and a neck; wherein the residual neural network includes four residual blocks Res-1, Res-2, Res-3, and Res-4, which are used to extract intermediate layer features containing semantic information at different levels; The neck consists of a convolutional layer and a dilated convolutional layer, which are used to align the channel dimension and increase the receptive field, respectively; The twin token mask decoder comprises a bimolecular decoder and a mask decoder, wherein the bimolecular decoder comprises a token generator and a semantic perceiver; The step S2 comprises the following specific steps: Step S2.1, on the downstream dataset COCO, perform data augmentation on the input image, scale it according to the aspect ratio and randomly flip it horizontally, and input the augmented image into the semantic & scale-aware image encoder in batches; then extract the global segmentation features of the image through the semantic & scale-aware image encoder, and extract the multi-scale semantic features of the image through the intermediate layer features output by the Res-1, Res-2, Res-3, and Res-4 residual blocks of the semantic & scale-aware image encoder; align the channel dimensions of the output partial convolution layers of the four residual blocks and increase the receptive field through the hollow convolution layer; use convolution layers with different receptive fields to further optimize the outputs of the semantic & scale-aware image encoder at different stages; Step S2.2, the twin-molecule decoder of the twin token decoder uses the token generator and the semantic perceptron to extract feature information that is beneficial to target recognition and feature information that is beneficial to target segmentation from the multi-scale semantic features and global segmentation features of the image; and obtains the target category and the image segmentation mask through the semantic perceptron and the mask decoder respectively; In step S2.3, the segmentation loss is calculated using the given standard segmentation mask and the image segmentation mask predicted by the U-SAM model, and the classification loss is calculated using the given standard target category and the target category predicted by the U-SAM model to supervise the output of the U-SAM model. Finally, the parameters of the U-SAM model are updated through the adaptive stochastic gradient descent method to achieve segmentation and recognition of the input image.

7. The fast and video memory friendly image segmentation method based on the SAM model according to claim 6, characterized in that: In step S2.2, the token generator includes an encoder for enhancing global segmentation features and a decoder 1 for generating unified tokens; the encoder first extracts feature information that is beneficial to the segmentation target from the global segmentation features, and uses its position encoding to enhance the spatial position information, and then inputs the decoder 1 to automatically generate a unified token, and further outputs an image segmentation mask through a mask decoder according to the unified token; The semantic sensor includes a feature pyramid network for extracting semantic features and a decoder 2 for obtaining target categories; the feature pyramid network extracts feature information that is beneficial to target recognition from multi-scale semantic features, and then combines the feature information that is beneficial to target segmentation to obtain the target category of the corresponding segmented target through the decoder 2.

8. The fast and video memory friendly image segmentation method based on the SAM model according to claim 6, characterized in that: In step S2.3, the U-SAM model uses binary cross entropy loss and Dice loss to express the segmentation loss: , respectively set and is the weighted coefficient of binary cross entropy loss and Dice loss; the U-SAM model uses multivariate cross entropy loss to represent the classification loss , is the weighting coefficient of the multivariate cross entropy loss; the final overall loss is the sum of the segmentation loss and the classification loss: ; in is the binary cross entropy loss, , represents the number of samples, Indicates whether the i-th sample is the target category, Indicates the probability that the i-th sample is the target category; It is the Dice loss, which has good performance in scenarios where positive and negative samples are severely unbalanced. , and Represent the given target segmentation mask and the predicted segmentation mask respectively; is the multivariate cross entropy loss, , and Indicates whether the i-th sample is a background class and the probability of being a background class, and They respectively represent whether the i-th sample is the j-th target category and the probability of being the j-th target category, is the loss coefficient of the background class, c is the number of categories in the downstream dataset; at the same time, in order to save video memory, when calculating the segmentation loss, not all points on the mask are involved in the loss calculation, but a certain number of points are randomly sampled to calculate the segmentation loss.

9. A computer system, characterized in that: The computer system includes a processor and a memory, wherein the memory stores at least one instruction, at least one program, a code set or an instruction set, and the at least one instruction, the at least one program, the code set or the instruction set are loaded and executed by the processor to implement a fast and video memory friendly image segmentation method based on a SAM model as described in any one of claims 1 to 8.

10. A computer-readable storage medium, characterized in that: The storage medium stores at least one instruction, at least one program, a code set or an instruction set, and the at least one instruction, the at least one program, the code set or the instruction set are loaded and executed by the processor to implement a fast and video memory friendly image segmentation method based on the SAM model as described in any one of claims 1 to 8.

Images