A fast and memory-friendly image segmentation method based on a SAM model

Abstract

A fast and GPU-friendly image segmentation method based on a SAM model, the image segmentation method comprising the following steps: first, establishing a U-SAM model based on SAM model training; the U-SAM model is trained by the following steps: step S1: establishing a SAM model and a U-SAM model; the semantic and scale-aware image encoder of the U-SAM model is obtained by compressing the image encoder of the SAM model through knowledge distillation, so as to improve the semantic perception ability of the SAM; step S2: freezing the relevant parameters of the SAM model, retaining the prompt segmentation ability, training the twin token mask decoder of the U-SAM model, and solving the problem of difficult control of prompt embedding. Second, using the U-SAM model to realize prompt segmentation and general image segmentation; thus, the prompt and general segmentation tasks can be realized through a fast and GPU-friendly framework.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of computer vision, and particularly to a fast and memory-friendly image segmentation method based on a SAM model. BACKGROUND

[0002] From image scene parsing to panoramic image segmentation, and now to general image segmentation, simultaneously processing multiple segmentation tasks is a long-term focus in the field of computer vision. General image segmentation mainly includes instance segmentation, semantic segmentation and panoramic segmentation, and the main difference between these tasks is the method of grouping pixels. Early panoramic segmentation models combined 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 Transformer (DETR) solved this problem by jointly modeling countable targets and non-countable targets through queries, attempting to implement 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 through various means, 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), which is pre-trained on a large-scale dataset (SA-1B), has become a basic model for image segmentation and has shown strong generalization performance in various image segmentation tasks.

[0003] Migrating SAM to specific tasks and optimizing the performance of SAM models are currently hot topics. The former improves the performance of SAM on tasks with weaker generalization ability through techniques such as fine-tuning or adapters. The latter solves the challenges of applying interactive segmentation frameworks to other segmentation tasks by extending the structure of SAM.

[0004] Through class-agnostic training methods, SAM can segment classes that do not exist in the training data, but this 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 scenario that requires predicting target masks, but it cannot be adapted to traditional image segmentation tasks that require simultaneously predicting the masks and corresponding classes of all objects in an image, such as instance segmentation, semantic segmentation and panoramic segmentation.

[0005] To address the above deficiencies, existing methods can be 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 large amount of computing resources (68 hours of training with 256 A100 GPUs), so it is not very feasible.

[0007] Secondly, the semantic information is extracted from the low-level features of the SAM and the corresponding prompt embeddings are generated by the object prompter for predicting the category and the mask. This method can greatly reduce the required computing resources, but the low-level features extracted by the visual self-attention model (ViT) based image encoder have high similarity with the high-level features, so the low-level features of the SAM also lack semantic information. In addition, the object prompter method cannot determine the effectiveness of the generated prompt embeddings during the learning process, and the segmentation performance of the SAM depends heavily on the quality of the prompt embeddings. Different types of prompt embeddings correspond to different forms and meanings, and using different types or quantities of prompts for the same object will result in different quality masks. At the same time, since the SAM aims to finely segment a single object according to the prompt, the special structure design makes it unable to process multiple objects in an image at the same time, and it can only process them one by one at the expense of speed or process multiple objects at the same time by image replication at the expense of memory. This has a high requirement for computing resources or memory, and therefore has low practicability. SUMMARY

[0008] The purpose of the present application is to provide a fast and memory-friendly image segmentation method based on a SAM model, which adopts a semantic-aware and memory-efficient universal segment anything model (U-SAM) and can realize prompt and traditional image segmentation tasks.

[0009] In order to achieve the above purpose, the solution of the present application is:

[0010] A fast and memory-friendly image segmentation method based on a SAM model, the image segmentation method comprising the following steps:

[0011] Firstly, a U-SAM model is trained and established based on a SAM model; the U-SAM model is trained by the following steps:

[0012] Step S1: establishing a SAM model and a U-SAM model; compressing the image encoder of the SAM model through knowledge distillation to obtain the semantic and scale-aware image encoder of the U-SAM model;

[0013] Step S2: freezing the relevant parameters of the SAM model to retain its original promptable segmentation capability to obtain the prompt encoder of the U-SAM model; and training the double token mask decoder of the U-SAM model on the downstream dataset COCO;

[0014] Secondly, the U-SAM model is used to realize promptable segmentation and general image segmentation;

[0015] The prompted segmentation process is: given an input image and an artificial prompt, the semantic and scale-aware image encoder extracts and outputs global segmentation features from the input image to the twin token mask decoder; the prompt encoder encodes the artificial prompt into a prompt embedding and outputs it to the twin token mask decoder; finally, the twin token mask decoder generates a corresponding prompt mask using the global segmentation features and the prompt embedding, and outputs the segmentation result;

[0016] The general image segmentation process is: given an input image, the semantic and scale-aware image encoder extracts and outputs global segmentation features from the input image to the twin token mask decoder, and extracts multi-scale semantic features of the intermediate layer features of the input image and outputs them to the twin token mask decoder; then the twin token mask decoder extracts feature information conducive to the segmentation target from the global segmentation features, and fuses the position encoding of the encoder in the twin token mask decoder to generate a unified token for general segmentation, and then generates a corresponding image segmentation mask using the global segmentation features and the unified token, and outputs the segmentation result; at the same time, the twin token mask decoder extracts feature information conducive to target recognition from the multi-scale semantic features, and fuses the feature information conducive to the segmentation target to obtain the target class of the target.

[0017] Further, the prompted segmentation process is represented as: ;

[0018] Wherein, is an artificial prompt, represents the first prompt; is a prompt embedding, represents the first prompt corresponding prompt embedding; is an input image, ; ; represents a prompt mask, ; represents a prompt encoder; 2IE represents a semantic and scale-aware image encoder; represents a twin token mask decoder;

[0019] Wherein, 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 input artificial prompts, represents the dimension of the image feature; represents the real number field;

[0020] A general image segmentation procedure is represented as: ;

[0021] wherein, is an input image, ; ; represents multi-scale semantic features; represents a target class of a target, ; represents an image segmentation mask, ; represents a semantic and scale-aware image encoder; represents a twin token mask decoder;

[0022] wherein, 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 dimension of the image feature; represents a real number field; represents the number of unified tokens in the twin token mask decoder, represents the number of classes contained in the downstream dataset.

[0023] Further, the step S1 comprises the following specific steps:

[0024] Step S1.1, performing data augmentation on the input image, scaling by the width-height ratio and randomly performing horizontal flipping on the distillation dataset SA-1B;

[0025] Step S1.2, inputting the augmented image into the SAM model image encoder and the semantic and scale-aware image encoder of the U-SAM model, respectively, for knowledge distillation.

[0026] Further, in step S1.2, a semantic and scale-aware image encoder based on a residual network structure is used as a student model, and an image encoder of a SAM model is used as a teacher model. The knowledge of the image encoder of the SAM model is transferred to the semantic and 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 and scale-aware image encoder.

[0027] Further, 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 and scale-aware image encoder. Then, the process of knowledge distillation is supervised by calculating the mean square error loss of the extracted features of the two.

[0028] Further, the semantic and scale-aware image encoder comprises a residual neural network and a neck; wherein the residual neural network comprises four residual blocks Res-1, Res-2, Res-3, Res-4 for extracting intermediate layer features containing different levels of semantic information; the neck is composed of convolutional layers and dilated convolutional layers for channel dimension alignment and receptive field enlargement, respectively;

[0029] The twin token mask decoder comprises a twin molecule decoder and a mask decoder, and the twin molecule decoder comprises a token generator and a semantic perceiver;

[0030] The step S2 comprises the following specific steps:

[0031] Step S2.1, on the downstream dataset COCO, data augmentation is performed on the input image, scaling is performed according to the width-height ratio and random horizontal flipping is performed, and the enhanced image is input into the semantic and scale-aware image encoder in batches; then the global segmentation features of the image are extracted through the semantic and scale-aware image encoder, and the multi-scale semantic features of the image are extracted from the intermediate layer features output by the Res-1, Res-2, Res-3, Res-4 residual blocks of the semantic and scale-aware image encoder; the output of the four residual blocks is aligned in channel dimension by the partial convolutional layers and the receptive field is enlarged by the dilated convolutional layers; the outputs of different stages of the semantic and scale-aware image encoder are further optimized by using the convolutional layers with different receptive fields;

[0032] Step S2.2, the twin molecule decoder of the twin token decoder extracts feature information conducive to target recognition and feature information conducive to segmented target from the multi-scale semantic features and global segmentation features of the image by using the token generator and the semantic perceiver; and the target class and the image segmentation mask are obtained by the semantic perceiver and the mask decoder, respectively;

[0033] Step S2.3, the segmentation loss is calculated by using the given standard segmentation mask and the image segmentation mask predicted by the U-SAM model, and the classification loss is calculated by using the given standard target class and the target class predicted by the U-SAM model, to supervise the output of the U-SAM model, and finally the parameters of the U-SAM model are updated by the adaptive stochastic gradient descent method to realize the segmentation and recognition of the input image.

[0034] Further, in step S2.2, the token generator comprises an encoder for enhancing the global segmentation features and a decoder 1 for generating a unified token; the encoder first extracts feature information conducive to segmented target from the global segmentation features, and uses its position encoding to enhance the spatial position information, then inputs the decoder 1 to automatically generate a unified token, and further outputs an image segmentation mask through the mask decoder according to the unified token;

[0035] The semantic perceiver comprises a feature pyramid network for extracting semantic features and a decoder 2 for obtaining a target category; the feature pyramid network extracts feature information conducive to target identification from multi-scale semantic features, and then combines feature information conducive to target segmentation to obtain the target category of the corresponding segmented target through the decoder 2.

[0036] Further, in step S2.3, the U-SAM model uses binary cross-entropy loss and Dice loss to represent the segmentation loss: , respectively, set and are the weighted coefficients of the binary cross-entropy loss and the Dice loss; the U-SAM model uses a multivariate cross-entropy loss to represent the classification loss , is the weighted coefficient of the multivariate cross-entropy loss; the final overall loss is the sum of the segmentation loss and the classification loss: ;

[0037] wherein is the binary cross-entropy loss, , represents the number of samples, represents whether the i-th sample is a target category, represents the probability that the i-th sample is a target category; is the Dice loss, which has good performance for scenes with severe positive and negative sample imbalance, , and respectively represent the given target segmentation mask and the predicted segmentation mask; is the multivariate cross-entropy loss, , and represent whether the i-th sample is a background class and the probability that it is a background class, and respectively represent whether the i-th sample is the j-th target category and the probability that it is the j-th target category, is the loss coefficient of the background class, and c is the number of categories of the downstream dataset.

[0038] The application also provides a computer system comprising 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 realize the fast and GPU-friendly image segmentation method based on the SAM model.

[0039] The application further provides a computer-readable storage medium, wherein at least one instruction, at least one program, a code set or an instruction set is stored in the storage medium, and the at least one instruction, the at least one program, the code set or the instruction set is loaded and executed by a processor to realize the fast and GPU-friendly image segmentation method based on the SAM model.

[0040] After the technical scheme is adopted, a semantic perception and memory efficient universal segmentation everything model (U-SAM) can realize prompt and general segmentation tasks through a fast and GPU-friendly framework.

[0041] The semantic and scale-aware image encoder (S2IE) proposed by the application can effectively solve the problem that the low-level features of SAM lack semantic information, and the twin token mask decoder (T2MD) proposed by the application can effectively solve the difficulty of learning prompt embedding and the problem of sacrificing GPU for multiple target segmentation.

[0042] The semantic and scale-aware image encoder (S2IE) proposed by the application can effectively solve the problem that the low-level features of SAM lack semantic information, and the twin token mask decoder (T2MD) proposed by the application can effectively solve the difficulty of learning prompt embedding and the problem of sacrificing GPU for multiple target segmentation.

[0043] The application has carried out extensive experiments on instance segmentation, semantic segmentation, panoramic segmentation and interactive segmentation, and verified the creativity of the application relative to the prior art. BRIEF DESCRIPTION OF DRAWINGS

[0044] Figure 1 This invention distinguishes the U-SAM model from existing methods.

[0045] Figure 2 This is a general framework diagram of the U-SAM model proposed in this invention;

[0046] Figure 3 This is an interactive segmentation visualization comparison chart of the U-SAM model proposed in this invention and other models;

[0047] Figure 4 This is a visual comparison chart of the U-SAM model and the SAM model proposed in this invention;

[0048] Figure 5 This is a visualization result of the general image segmentation of the U-SAM model proposed in this invention;

[0049] Figure 6 This is a flowchart of the image segmentation method proposed in this invention;

[0050] Figure 7 This is a flowchart of the U-SAM model training method proposed in this invention. Detailed Implementation

[0051] 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 with reference to the accompanying drawings. Obviously, the described embodiments are merely some embodiments of the present application, and not all embodiments. All other embodiments obtained by those skilled in the art based on the embodiments of the present application without creative effort are within the scope of protection of the present application.

[0052] This invention leverages the powerful image segmentation capabilities and zero-shot transfer 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 class capability and cannot be applied to general image segmentation tasks. It also addresses issues such as... Figure 1 The existing methods described above require significant computational resources through retraining or using complex structures, and also suffer from memory requirements due to the simultaneous processing of multiple targets using a memory-for-speed approach. This approach aims to enable the model to quickly adapt to different image data distributions and overcome the drawback of needing to reuse specific image data for fine-tuning for different problems.

[0053] wherein the SAM model is an existing one-cut-all model, and the structure of the SAM model comprises an image encoder, a prompt encoder, and a mask decoder. The SAM model is trained on a large dataset using an interactive learning and prompt learning strategy. When the SAM model is used, a user needs to provide prompt information for a target to be segmented, such as points inside the target, a bounding box of the target, or a doodle mask. The SAM model has good zero-shot generalization capability and can be migrated to various image segmentation tasks.

[0054] The embodiment is based on a one-cut-all model (SAM model) to establish a unified one-cut-all model (U-SAM model) to realize a fast and memory-friendly image segmentation method.

[0055] As Figure 2 , the U-SAM model comprises a prompt encoder, a semantic and scale-aware image encoder (S2IE), and a twin token mask decoder (T2MD).

[0056] The semantic and scale-aware image encoder comprises four residual blocks Res-1, Res-2, Res-3, and Res-4.

[0057] The twin token mask decoder comprises a twin molecule decoder and a mask decoder.

[0058] The image segmentation method comprises the following steps:

[0059] First, the U-SAM model is established based on the training of the SAM model; see Figure 7 The training method of the U-SAM model can be divided into the following steps:

[0060] Step S1: Establishing the SAM model and the U-SAM model; on the distillation dataset SA-1B, the image encoder of the SAM model is compressed by knowledge distillation to obtain the semantic and scale-aware image encoder of the U-SAM model.

[0061] Step S2: Freezing the related parameters of the SAM model, retaining its original promptable segmentation capability, obtaining the prompt encoder of the U-SAM model; and on the downstream dataset COCO, training the twin token mask decoder of the U-SAM model.

[0062] Second, the U-SAM model is used to realize promptable segmentation and general image segmentation.

[0063] See Figure 6 The promptable segmentation process can be represented as:

[0064]

[0065] ;

[0066] Specifically, similar to the promptable segmentation process of the visual segmentation base model SAM, given an input image and a human prompt , a semantic and scale-aware image encoder first extracts global segmentation features from the input image , then a prompt encoder encodes the human prompt into a prompt embedding , and finally a twin token mask decoder generates a corresponding prompt mask using the global segmentation features and the prompt embedding , and outputs the segmentation result.

[0067] wherein, is the human prompt, denotes the i-th prompt; is the prompt embedding, denotes the i-th prompt corresponding to the 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 input human prompts, represents the dimension of the image feature. The process of general image segmentation can be represented as:

[0068]

[0069]

[0070] .

[0071] Unlike promptable segmentation, general image segmentation is only given an input image, and the semantic and scale-aware image encoder outputs global segmentation features, as well as multi-scale semantic features of the image intermediate layer features . Subsequently, the twin token mask decoder extracts feature information beneficial to the segmentation target from the global segmentation features , and fuses the position encoding to generate a unified token for general segmentation , and then generates a corresponding image segmentation mask using the global segmentation features and the unified token , and outputs the segmentation result.

[0072] ​​The twin token mask decoder extracts feature information conducive to target recognition from multi-scale semantic features , and fuses the feature information conducive to target segmentation to obtain the target category of the target .

[0073] Among them, represents the number of unified tokens in the twin token mask decoder, represents the number of categories contained in the data set.

[0074] The training method of the U-SAM model is specifically as follows:

[0075] Step S1 includes the following specific steps:

[0076] Step S1.1, on the distillation data set SA-1B, data augmentation is performed on the input image, scaling is performed according to the width-height ratio, and horizontal flipping is randomly performed.

[0077] Step S1.2, input the enhanced image into the SAM model image encoder and the semantic and scale perception image encoder of the U-SAM model respectively for knowledge distillation; the semantic and scale perception image encoder based on the residual network structure is used as the student model, the SAM image encoder is used as the teacher model, the knowledge of the SAM image encoder is transferred to the semantic and scale perception 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 the SAM image encoder and the U-SAM semantic and scale perception image encoder.

[0078] The semantic and scale perception image encoder of the application aims to enhance the semantic perception ability of the SAM model by integrating multi-scale feature extraction capability through knowledge distillation. On the one hand, the high and low layer features extracted by the convolutional neural network have lower feature similarity than the visual self-attention model, so the low layer feature distribution will not tend to the visual basic model features with semantic loss in the distillation process. On the other hand, the residual neural network (ResNet) structure can extract good multi-scale features, which is conducive to target positioning and recognition.

[0079] Specifically, in step S1.2, given an input image, the teacher model SAM extracts features from the data enhanced image through its image encoder , and the student model U-SAM extracts global segmentation features through the above semantic and scale perception image encoder. Subsequently, the process of knowledge distillation is supervised by calculating the mean square error loss of the extracted features. Among them, and represent the height and width of the image feature map respectively, dimensions representing image features, representing a real number field.

[0080] Wherein step S2 can be roughly subdivided into the following three steps:

[0081] Step S2.1, data augmentation is performed on the input image on the downstream dataset COCO, scaling by width and height ratio and randomly horizontal flipping, and the augmented image is input into the semantic and scale-aware image encoder in batches. Then the global segmentation features of the image are extracted through the semantic and 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, Res-4 residual blocks;

[0082] Specifically, in step S2.1, the semantic and 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, respectively, for extracting four different levels of semantic information. The neck is composed of convolutional layers and dilated convolutional layers, which are respectively used for channel dimension alignment and receptive field enlargement, to further optimize the outputs of different residual blocks of the semantic and scale-aware image encoder, to obtain the multi-scale semantic features and the global segmentation features of the image.

[0083] Specifically, as Figure 2 The upper right part, given the augmented image, the semantic and scale-aware image encoder first extracts the multi-scale semantic features of the image through the residual neural network wherein respectively represent the output features of the four residual blocks of the residual network. Then the global segmentation features .

[0084] The twin token mask decoder includes a twin molecule decoder and a mask decoder, and the twin molecule decoder includes a token generator and a semantic perceiver.

[0085] In step S2.2, the obtained multi-scale semantic features and global segmentation features are input into the Twin Token Mask Decoder (T2MD). The Twin Token Mask Decoder 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 global segmentation features, respectively. Then, the obtained target recognition feature information 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, which is also the image segmentation mask predicted by the U-SAM model.

[0086] To avoid the difficulties of directly learning cue embeddings, the bi-token mask decoder of this invention abandons the method of generating cue embeddings from image features through a target cue generator. Instead, it employs a set of additional token generators and semantic perceptrons to output the segmentation mask and target category of the image. This is because traditional image segmentation tasks perform a global partitioning of the image, independent of user cues, and replacing image duplication with token redundancy effectively reduces the memory requirements during training and inference. Furthermore, to minimize the impact of semantic information on the model's original segmentation ability, the bi-token mask decoder decouples the segmentation and classification tasks through a bi-token decoder.

[0087] Specifically, the upper part of the bimolecular decoder, called the Token Generator (T.Gen), contains the encoder used to enhance global segmentation features. ) and decoder 1 for generating unified tokens ( The generated unified token is then used by the mask decoder to generate an image segmentation mask. The lower half of the bimolecular decoder is called the Semantic Perception (S.Per), which contains a feature pyramid network for extracting semantic features. ) and decoder 2 for obtaining the target category ( To further improve model performance, 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 cue, the twin token mask decoder reduces training memory by 54% and accelerates inference speed by 43% while maintaining essentially the same performance.

[0088] like Figure 2 The lower right portion shows the global segmentation features output by a semantically and scale-aware image encoder. ) and the residual neural network ) multi-scale semantic features of the intermediate layer output ), the encoder ) first extracts feature information conducive to segmentation targets from the global segmentation features ), and uses position encoding ) to enhance spatial position information, and then inputs the decoder 1 to automatically generate a unified token, and further outputs an image segmentation mask ) according to the unified token through the mask decoder. At the same time, the feature pyramid network extracts feature information conducive to target recognition from the multi-scale semantic features ), and then combines the feature information conducive to target segmentation to obtain the target class ) corresponding to the segmentation target through the decoder 2.

[0089] Step S2.3, calculate the segmentation loss using the given standard segmentation mask and the image segmentation mask predicted by the U-SAM model, and calculate the classification loss using the given standard target class and the target class predicted by the U-SAM model, and finally update the parameters of the U-SAM model by the adaptive stochastic gradient descent method.

[0090] 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 class predicted by the U-SAM model and the given standard target class is calculated to supervise the output of the U-SAM model. Update the U-SAM model parameters using the adaptive stochastic gradient descent method to realize the segmentation and recognition of the input image;

[0091] For the segmentation loss, the U-SAM model uses binary cross-entropy loss and Dice loss instead of focal loss: , respectively and are the weighting coefficients of the binary cross-entropy loss and the Dice loss.

[0092] For the classification loss, use the multivariate cross-entropy loss to represent , 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: ;

[0093] where is the binary cross-entropy loss, , represents the number of samples, represents whether the i-th sample is a target class, represents the probability that the i-th sample is a target class. is Dice loss, which has good performance in the scene of serious imbalance between positive and negative samples, , and denote the given target segmentation mask and the predicted segmentation mask, respectively; is the multi-class cross-entropy loss, , and denote whether the i-th sample is a background class and the probability of being a background class, and denote whether the i-th sample is the j-th target class and the probability of being the j-th target class, is the loss coefficient of the background class, and c is the number of classes in the downstream dataset. At the same time, in order to save memory, not all points on the mask are involved in the calculation of the segmentation loss, but a certain number of points are randomly sampled to calculate the segmentation loss.

[0094] To verify the effectiveness and advancement of the method, the embodiment verifies the fast and memory-friendly image segmentation method based on the SAM model. It should be noted that for the foregoing method embodiments, in order to facilitate the description, they are all expressed as a series of action combinations, but those skilled in the art should know that the present application is not limited by the order of the described actions, because according to the present application, certain steps can be performed in other order or simultaneously.

[0095] The image encoder of the SAM model has a visual-oriented inductive bias, which lays the foundation for the robust segmentation and generalization ability of the SAM model, but the related computational overhead leads to very slow processing of images. Compared with the serial input used in the prompt 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.

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

[0097]

[0098] Table 1 shows the running speed of the U-SAM model on a single NVIDIA GeForce RTX 3090 GPU compared with other models. The parameter amount of the U-SAM model is reduced by 93.0% compared with SAM-H, and only accounts for 48.3% of SAM-B. At the same time, the single prompt inference running speed of the present application is 10 times faster than SAM-H and 2.7 times faster than SAM-B, reaching a similar speed to FastSAM and MobileSAM. It is worth noting that whether it is SAM-H with the slowest single prompt inference speed or MobileSAM with the fastest single prompt inference speed, the difference in average single prompt running speed is only 0.1-0.2 milliseconds when using the 64x64 "cut everything" method. This is because during the interaction process, the image only needs to pass through the image encoder once, making its inference speed real-time amortization.

[0099] Table 2. Comparison of U-SAM and other models on various zero-shot tasks. The compared models have not been trained on the corresponding tasks, nor have they seen any images or annotations from related datasets.

[0100]

[0101] The global optimal (Optimal Dataset Scale, ODS), single image optimal (Optimal Image Scale, OIS), average precision (Average Precision, AP), and recall at 50% precision (Recall at 50% precision, R50) of U-SAM on edge detection are evaluated on the Berkeley Segmentation Data Set (BSDS500).

[0102] We first generate a mask probability map through the SAM model, and then use Non Max Suppression (NMS) to filter redundant masks. Subsequently, the Sobel operator is used to generate an edge probability map, and finally edge NMS is used for post-processing.

[0103] As shown in Table 2, U-SAM is more close to SAM-H than FastSAM and MobileSAM on each index, and the numerical distribution is basically consistent with MobileSAM and SAM-H.

[0104] We evaluate the average recall (AR) of the present invention on the LVIS dataset for object proposal generation.

[0105] As SAM, we do not crop the image, nor do we filter low-quality masks using the Intersection of Union (IoU) and stability scores. We only filter redundant masks through the remaining parameters. If the number of masks exceeds one thousand when testing, we take the one thousand with the highest average of IoU and stability scores as the result.

[0106] Table 2 verifies that the average recall AR of the present invention is relatively lower than that of SAM, but we find that the present invention has basically the same preference for masks of different sizes as SAM, and the average recall AR of masks of various sizes is basically the same. MobileSAM, however, has a stronger preference for large targets, so the average recall AR of large masks is relatively higher.

[0107] At the same time, we evaluate the average precision (AP) of the present invention on the LVIS dataset for instance segmentation. As SAM, we first generate bounding boxes and corresponding target classes through the visual self-attention-based detection model (ViTDet), and then use the detection boxes (target bounding boxes generated by the detection model) as the prompts of the present invention to generate instance segmentation masks.

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

[0109] Table 3. Comparison of U-SAM based on COCO with other object prompt-based models on Tesla V100 GPU test data. The training time is in days.

[0110]

[0111] Table 3 further compares the performance of this invention with other target-based prompting models on the COCO dataset. It is worth noting that the full Remote Sensing Prompter (RSPrompter) requires excessive memory; the RSPrompter used here is a parameter-tuned version. We retained all components of the RSPrompter but reduced its training memory consumption by adjusting the parameters. Specifically, we adjusted the output of the multi-scale features to 4, reduced the number of encoder layers in the pixel decoder to 1, and reduced the number of decoder layers in the mask decoder to 4.

[0112] This invention demonstrates significant advantages over RSPrompter in all aspects. Using only 10% of the parameters, it achieves performance improvements of 5.2 (AP), 6.3 (MIoU), and 6.2 (PQ) in instance, semantic, and panoptic segmentation tasks, respectively, while also realizing a 13x speedup. Furthermore, compared to RSPrompter based on SAM-B or SAM-L, the performance advantages of this invention are even more pronounced, with its greatest advantage being a 61% reduction in memory requirements and a significant reduction in training time.

[0113] Furthermore, this invention aligns the distribution of image encoder output features between its lightweight feature extraction network and the SAM model through knowledge distillation. For example... Figure 3 The image shows a visualization comparison between U-SAM and other models. The red boxes highlight regions successfully segmented by this invention that were missed by other models, 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 segmented regions belong to uncountable (stuff) targets, which to some extent illustrates that U-SAM's ability to identify uncountable targets is closer to that of SAM than MobileSAM and FastSAM. The blue boxes highlight some SAM biases learned by this invention.

[0114] Specifically, both SAM and this invention consider the clouds in the upper right corner of the first row and the sky, as well as the two pieces of ground in the third row, to be separate entities. FastSAM and MobileSAM, however, consider them to be the same entity. Furthermore, this invention's scratches in the sky in the third row and its segmentation of buildings in the second row are consistent with SAM's. These findings indicate that, compared to FastSAM and MobileSAM, this invention's segmentation preferences are more aligned with SAM's, further validating the experimental conclusions.

[0115] Figure 4For the visual comparison between U-SAM and SAM, it is verified that the segmentation results of the present application have stronger consistency with SAM than FastSAM and MobileSAM. However, due to the compression of the image encoder of SAM by knowledge distillation, the present application also loses part of the performance. The first row shows that the present application generates the same segmentation results as SAM, so the segmentation results of the present application are basically consistent with SAM for part of the images. The second row verifies that the present application cannot successfully segment smaller countable (things) targets, such as the person on the left and the cup on the right of the second row. This to some extent shows that the instance segmentation performance of the present application may not be ideal. The third row reflects that the present application is more inclined to segment adjacent targets of the same type as a whole, such as the mountains and trees on the left and the buildings and stairs on the right of the third row. And these regions mostly belong to uncountable targets, so this property is conducive to completing general image segmentation.

[0116] Figure 5 The general image segmentation visualization results of U-SAM on the downstream dataset are shown, which can verify that the present application is effective on the downstream dataset.

Claims

1. A fast and GPU-friendly image segmentation method based on a SAM model, characterized in that: the image segmentation method comprises the following steps: Step 1: Establishing a U-SAM model based on SAM model training; the U-SAM model is trained by the following steps: Step S1: Establishing a SAM model and a U-SAM model; compressing the image encoder of the SAM model through knowledge distillation to obtain a semantic and scale-aware image encoder of the U-SAM model; Step S2: Freezing the relevant parameters of the SAM model, retaining its original promptable segmentation capability, obtaining a prompt encoder of the U-SAM model; and training a twin token mask decoder of the U-SAM model on a downstream dataset COCO; Step 2: Using the U-SAM model to realize promptable segmentation and general image segmentation; The process of promptable segmentation is: given an input image and an artificial prompt, the semantic and scale-aware image encoder extracts and outputs global segmentation features from the input image to the twin token mask decoder; The prompt encoder encodes the artificial prompt into a prompt embedding and outputs it to the twin token mask decoder; finally, the twin token mask decoder generates a corresponding prompt mask using the global segmentation features and the prompt embedding, and outputs the segmentation result; The process of general image segmentation is: given an input image, the semantic and scale-aware image encoder extracts and outputs global segmentation features from the input image to the twin token mask decoder, and extracts multi-scale semantic features of the intermediate layer features of the input image and outputs them to the twin token mask decoder; Then the twin token mask decoder extracts feature information conducive to the segmentation target from the global segmentation features, and fuses the position encoding of the encoder in the twin token mask decoder to generate a unified token for general segmentation, and then generates a corresponding image segmentation mask using the global segmentation features and the unified token, and outputs the segmentation result; at the same time, the twin token mask decoder extracts feature information conducive to target recognition from the multi-scale semantic features, and fuses the feature information conducive to the segmentation target to obtain the target class of the target; The semantic and scale-aware image encoder comprises a residual neural network and a neck; wherein the residual neural network comprises four residual blocks Res-1, Res-2, Res-3, and Res-4 for extracting intermediate layer features containing different levels of semantic information; The neck is composed of convolution layers and atrous convolution layers, which are respectively used for channel dimension alignment and receptive field enlargement; The twin token mask decoder comprises a twin molecule decoder and a mask decoder, and the twin molecule decoder comprises a token generator and a semantic perceiver.

2. The fast and GPU-friendly image segmentation method based on the SAM model according to claim 1, characterized in that:

3. The fast and GPU-friendly image segmentation method based on the SAM model according to claim 1, characterized in that: The segmentation procedure can be prompted to represent as: ; wherein, is a human prompt, represents the prompt; is a prompt embedding, represents the prompt corresponding prompt embedding; is an input image, ; ; represents a prompt mask, ; represents a prompt encoder; represents a semantic and scale-aware image encoder; represents a twin token mask decoder; wherein, 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 input human hints, represents the dimension of the image feature; denotes the real number field; A general image segmentation procedure is represented as: ; wherein, is an input image, ; ; denotes a multi-scale semantic feature; denotes a target class of a target, ; denotes an image segmentation mask, ; denotes a semantic and scale-aware image encoder; denotes a twin token mask decoder; wherein, 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 dimension of the image feature; denotes the real number field; represents the number of unified tokens in the twin token mask decoder, represents the number of classes contained in the downstream dataset. Step S1 comprises the following specific steps: Step S1.1, on the distillation dataset SA-1B, data augmentation is performed on the input image, scaling is performed according to the width-height ratio, and random horizontal flipping is performed. ​ Step S1.2, input the enhanced image into the SAM model image encoder and the semantic and scale perception image encoder of the U-SAM model respectively 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 and scale perception 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 and scale perception 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 and scale perception image encoder.

5. The fast and 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-enhanced image through its image encoder, and the student model extracts global segmentation features through the semantic and scale perception image encoder. The process of knowledge distillation is then supervised by calculating the mean square error loss of the extracted features.

6. The fast and memory-friendly image segmentation method based on the SAM model according to claim 1, characterized in that: The step S2 includes the following specific steps: Step S2.1, data augmentation is performed on the input image on the downstream dataset COCO, scaling is performed according to the width-height ratio, and horizontal flipping is randomly performed. The enhanced image is input into the semantic & scale perception image encoder in batches; then the global segmentation features of the image are extracted through the semantic & scale perception image encoder, and the multi-scale semantic features of the image are extracted from the intermediate layer features output by the Res-1, Res-2, Res-3, and Res-4 residual blocks of the semantic & scale perception image encoder; the output parts of the four residual blocks are convolved to align the channel dimensions and increase the receptive field through the empty convolution layer; the convolution layers with different receptive fields are used to further optimize the outputs of different stages of the semantic & scale perception image encoder; Step S2.2, the twin molecule decoder of the twin token decoder extracts feature information conducive to target recognition and feature information conducive to segmented target from the multi-scale semantic features and global segmentation features of the image using the token generator and semantic perceiver; and the target class and image segmentation mask are obtained through the semantic perceiver and mask decoder respectively; 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 class and the target class 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 by the adaptive stochastic gradient descent method to realize the segmentation and recognition of the input image.

7. The fast and 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 the global segmentation features and a decoder 1 for generating unified tokens; the encoder first extracts feature information conducive to segmented targets from the global segmentation features, and uses its position encoding to enhance the spatial position information, then inputs the decoder 1 to automatically generate unified tokens, and further outputs the image segmentation mask through the mask decoder according to the unified tokens; The semantic sensor comprises a feature pyramid network for extracting semantic features and a decoder 2 for obtaining a target category; the feature pyramid network extracts feature information conducive to target recognition from multi-scale semantic features, then combines feature information conducive to target segmentation, and obtains the target category corresponding to the segmented target through the decoder 2.

8. The fast and 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 represent the segmentation loss: Set up separately and The weighting coefficients are the binary cross-entropy loss and the Dice loss; the U-SAM model uses multivariate cross-entropy loss to represent the classification loss. , These are the weighting coefficients for the multivariate cross-entropy loss; the final overall loss is the sum of the segmentation loss and the classification loss. ; wherein is a binary cross-entropy loss, , denotes the number of samples, denotes whether the ith sample is the target class, denotes the probability that the ith sample is the target class; is a Dice loss, , and denote the given target segmentation mask and the predicted segmentation mask, respectively; is a multi-class cross-entropy loss, , and denote whether the ith sample is the background class and the probability that the ith sample is the background class, and denote whether the ith sample is the jth target class and the probability that the ith sample is the jth target class, respectively, is a loss coefficient for the background class, c is the number of classes of the downstream dataset; at the same time, in order to save the memory, not all points on the mask participate in the calculation of the segmentation loss, but a certain number of points are randomly sampled to calculate the segmentation loss.

9. A computer system, characterized by The computer system comprises a processor and a memory, and the memory stores at least one instruction, at least one program, a code set or an instruction set, which are loaded and executed by the processor to realize the fast and GPU-friendly image segmentation method based on the SAM model according to 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, which are loaded and executed by the processor to realize the fast and GPU-friendly image segmentation method based on the SAM model according to any one of claims 1 to 8.

Images