Image segmentation method and device, electronic equipment and computer readable storage medium

Abstract

The invention provides an image segmentation method and device, electronic equipment and a computer readable storage medium, and the method comprises the steps: carrying out the feature extraction of an original image, and obtaining an image feature and an image class identifier; a target prompt feature corresponding to the original image is matched from a knowledge base through the image class identifier, and the knowledge base is pre-constructed and comprises multiple pairs of positive prompt features and negative prompt features; performing feature enhancement processing on the target prompt feature by using the image feature to obtain an enhanced target prompt feature; and performing mask decoding processing on the image features and the enhanced prompt features to obtain a segmentation mask for the original image. According to the invention, segmentation of the original image is guided by acquiring the target prompt features corresponding to the original image from the knowledge base, the effect of accurate segmentation of the original image in practical application is realized, and the target prompt features corresponding to various surface feature categories in the original image are used, so that the segmentation accuracy of the original image is improved. Synchronous segmentation of various ground feature categories of the original image is realized, and the segmentation efficiency is improved.

Description

Technical Field

[0001] The present invention relates to the field of image processing and artificial intelligence technology, and in particular to an image segmentation method, device, electronic device and computer-readable storage medium. Background Art

[0002] Image segmentation refers to the technology and process of dividing an image into several specific areas with unique properties and proposing objects of interest. Remote sensing images have many types of objects, varied scenes and large sizes, which makes remote sensing image segmentation difficult. The accurate segmentation of remote sensing images has an important impact on the subsequent processing and application of remote sensing images.

[0003] For the segmentation of remote sensing images, an image segmentation model is usually obtained by training based on a training data set. However, in practical applications, since the image data to be processed is often quite different from the training data set, when the segmentation model is used to segment the image data to be processed, some objects are still difficult to distinguish clearly, the segmentation effect is poor, and the generalization ability of the segmentation model is insufficient. In addition, in the related art, the semantics of the target are segmented by manually inputting prompt information or setting prompt information, which cannot segment multiple object categories simultaneously, and the segmentation efficiency is low.

[0004] Therefore, it is necessary to develop an image segmentation method that can accurately and efficiently segment various images to be segmented in practical applications. Summary of the invention

[0005] The invention provides an image segmentation method, which obtains target prompt features from a knowledge base to guide the segmentation of the original image, thereby achieving the purpose of automatic and accurate segmentation of images of various types of land objects.

[0006] In a first aspect, an embodiment of the present application provides an image segmentation method, the method comprising:

[0007] S1: obtaining an original image and a knowledge base, wherein the knowledge base is pre-built and contains multiple pairs of positive prompt features and negative prompt features;

[0008] S2: Extract features from the original image to obtain image features and image class identifiers;

[0009] S3: using the image class identifier to match the target prompt feature corresponding to the original image from the knowledge base;

[0010] S4: performing feature enhancement processing on the target prompt feature using the image feature to obtain an enhanced target prompt feature;

[0011] S5: performing mask decoding processing on the image features and the enhanced target prompt features to obtain a segmentation mask for the original image.

[0012] Further, the S3 includes:

[0013] The similarity between the image class identifier and the positive prompt features in the knowledge base is calculated, and the k positive prompt features with the highest similarity ranking and their associated k negative prompt features are used as the target prompt features.

[0014] The construction process of the knowledge base includes:

[0015] Acquire a knowledge image, and annotate multiple types of objects in the knowledge image with masks and set labels respectively, to obtain an annotated knowledge image;

[0016] For the annotated knowledge image, a sub-image in a mask area of ​​a certain type of feature is obtained, and a class identifier of the sub-image is obtained using a pre-trained ViT network to obtain the positive prompt feature;

[0017] A sub-image that is not in a mask area of ​​a certain type of object is obtained, a class identifier of the sub-image is obtained using a pre-trained ViT network, the negative prompt feature is obtained, and the positive prompt feature, the negative prompt feature and the label of the certain type of object are associated and stored.

[0018] The S2 includes:

[0019] Multi-scale feature extraction and feature fusion are performed on the original image to obtain a first image feature, a second image feature and the image class identifier, and the first image feature and the second image feature are used as the image features; wherein the scale of the first image feature is greater than the scale of the second image feature.

[0020] Furthermore, the method for obtaining the first image feature, the second image feature and the image type identifier includes:

[0021] Use the pre-trained CNN network to extract multi-scale features from the original image, and obtain the first sub-feature, the second sub-feature, the third sub-feature, and the fourth sub-feature with decreasing scales in sequence;

[0022] Using the pre-trained ViT network to extract features from the original image, to obtain the image class identifier and the fifth sub-feature;

[0023] Using a feature fusion network to fuse the fourth sub-feature, the third sub-feature, and the fifth sub-feature to obtain the second image feature;

[0024] The second image feature, the second sub-feature and the first sub-feature are fused using a feature fusion network to obtain the first image feature.

[0025] The S4 includes:

[0026] The second image feature and the target prompt feature are input into a Transformer network, and the target prompt feature is enhanced by using a cross-attention mechanism to obtain the enhanced target prompt feature.

[0027] The S5 includes:

[0028] A matrix multiplication operation is performed on the first image feature and the enhanced target prompt feature to obtain a segmentation mask for the original image, wherein the segmentation mask corresponds to a plurality of ground object categories of the original image.

[0029] In a second aspect, an embodiment of the present application provides an image segmentation device, the device comprising:

[0030] An input unit, used for acquiring an original image and a knowledge base, wherein the knowledge base is pre-built and includes a plurality of pairs of positive prompt features and negative prompt features;

[0031] A feature extraction unit, used to extract features from the original image to obtain image features and image class identifiers;

[0032] A prompt matching unit, configured to match a target prompt feature corresponding to the original image from the knowledge base using the image class identifier;

[0033] A prompt embedding unit, used for performing feature enhancement processing on the target prompt feature by using the image feature to obtain an enhanced target prompt feature;

[0034] The mask decoding unit is used to perform mask decoding processing on the image feature and the enhanced target prompt feature to obtain a segmentation mask for the original image.

[0035] In a third aspect, an embodiment of the present application provides an electronic device, comprising: a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the processor executes the computer program, it can execute the method of the first aspect.

[0036] In a fourth aspect, an embodiment of the present application provides a computer-readable storage medium having a computer program stored thereon, wherein when the computer program is executed by a processor, the method of the first aspect can be executed.

[0037] The beneficial effects of the present invention include at least:

[0038] 1. In the image segmentation method provided by the present invention, the positive prompt features and negative prompt features corresponding to the original image are obtained from the knowledge base by using the image class identifier of the original image to guide the segmentation of the original image, thereby achieving the effect of accurately segmenting various original images in practical applications;

[0039] 2. By using the target prompt features corresponding to various types of objects in the original image, the synchronous segmentation of various types of objects in the original image is achieved, thereby improving the segmentation efficiency;

[0040] 3. The first image feature and the second image feature are extracted by using a dual feature extraction network, wherein the first image feature is a fusion of multiple scale features of the original image. By using the first image feature for mask decoding, the multi-scale features are fully utilized to improve the segmentation accuracy. BRIEF DESCRIPTION OF THE DRAWINGS

[0041] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings required for use in the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative work.

[0042] Figure 1 A flowchart of an image segmentation method provided in an embodiment of the present application;

[0043] Figure 2 A flowchart of constructing a knowledge base provided in an embodiment of the present application;

[0044] Figure 3 A schematic structural diagram of an image segmentation model provided in an embodiment of the present application;

[0045] Figure 4 A schematic structural diagram of another image segmentation model provided in an embodiment of the present application;

[0046] Figure 5 A schematic block diagram of an image segmentation device provided in an embodiment of the present application. DETAILED DESCRIPTION

[0047] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field belong to the scope of protection of the present invention.

[0048] Figure 1 A flowchart of the image segmentation method provided in the embodiment of the present application is shown in FIG. Figure 1 As shown, the method includes:

[0049] S1: obtaining an original image and a knowledge base, wherein the knowledge base is pre-built and contains multiple pairs of positive prompt features and negative prompt features;

[0050] S2: Extract features from the original image to obtain image features and image class identifiers;

[0051] S3: using the image class identifier to match the target prompt feature corresponding to the original image from the knowledge base;

[0052] S4: performing feature enhancement processing on the target prompt feature using the image feature to obtain an enhanced target prompt feature;

[0053] S5: performing mask decoding processing on the image features and the enhanced target prompt features to obtain a segmentation mask for the original image.

[0054] According to the above process, the present application first obtains the image features and image class identifiers (ClassToken, Cls Token) of the original image, wherein the image class identifier represents the overall distribution of the object categories of the original image, and then automatically matches the corresponding target prompt features from the knowledge base according to the image class identifier of the original image, and then enhances the target prompt features, and finally performs mask decoding according to the image features and the enhanced target prompt features to obtain a segmentation mask for the original image. Through the method provided by the present invention, it is possible to segment original images of various object categories using the guidance of the knowledge base.

[0055] Next, the above steps are described in detail in combination with the embodiments. First, step S1 is described in detail in combination with the embodiments.

[0056] The original image in this application is the remote sensing image data to be segmented. The original image may include data of multiple bands, for example, red, green, and blue three-band data, or near infrared band data. Optionally, the original image may also include other band data.

[0057] In this application, the knowledge base is pre-constructed based on the knowledge image. The knowledge base contains multiple positive prompt features and multiple negative prompt features, and the positive prompt features and the negative prompt features exist in pairs.

[0058] Optionally, the knowledge image may be three-band remote sensing image data or four-band remote sensing image data. Optionally, the knowledge image may be a remote sensing image similar to an original image in actual application. Optionally, the knowledge image may be a remote sensing image that is not similar to an original image in actual application.

[0059] Figure 2 The flowchart of constructing the knowledge base provided in the embodiment of the present application is as follows: Figure 2 As shown in Figure 1, the process of building a knowledge base includes:

[0060] Step a: Acquire a knowledge image, and annotate various types of objects in the knowledge image with masks and set labels respectively to obtain an annotated knowledge image.

[0061] Specifically, the knowledge image is first obtained, and then all the object categories in the knowledge image are annotated with masks, and different labels are set for each object category to obtain the annotated knowledge image.

[0062] Optionally, the annotation may be manual or automatic.

[0063] Optionally, some of the feature categories in the knowledge image may be marked with masks, for example, the feature categories that occupy a larger area may be marked.

[0064] Step b: for the annotated knowledge image, a sub-image in the mask area of ​​a certain type of feature is obtained, and the class identifier of the sub-image is obtained by using the pre-trained ViT network to obtain the positive prompt feature.

[0065] Specifically, the annotated knowledge image is divided into multiple image blocks, and for a certain type of land object, multiple image blocks within the corresponding mask range are selected and input into the pre-trained ViT network, and the class identifier output by the network is the positive prompt feature. Optionally, the ViT network can use the ViT network in the pre-trained self-distillation with no labels (DINO) model. Among them, the pre-trained DINO model is a model obtained by self-supervised learning training on a remote sensing image dataset, and the backbone network of the DINO model adopts the ViT network. DINO is an existing model and will not be described in detail here.

[0066] The annotated knowledge image is divided into multiple image blocks, and the image blocks in the mask are selected according to a certain feature category annotation, and the initial class identifier is added and sent to the pre-trained ViT network. The network output class identifier is used as the positive prompt feature of the knowledge image. Among them, the initial class identifier is the class identifier carried by ViT after pre-training.

[0067] Step c: obtain a sub-image that is not in the mask area of ​​a certain object category, use the pre-trained ViT network to obtain the class identifier of the sub-image, obtain the negative prompt feature, and associate the positive prompt feature, the negative prompt feature and the label of the certain object category for storage.

[0068] Specifically, the image blocks that are not in the mask are selected according to a certain type of feature category annotation, and the initial class identifier is added and sent to the pre-trained ViT network. The network output class identifier is used as the negative prompt feature of the knowledge image.

[0069] According to steps b to c, corresponding positive prompt features and negative prompt features can be produced according to other land object categories in the same knowledge image.

[0070] Other knowledge images can be processed in the same way as step a to step d, and positive prompt features and negative prompt features are continuously added until the preset conditions are met to obtain the knowledge base. The knowledge base finally obtained includes multiple positive prompt features and multiple negative prompt features corresponding to m types of land object categories. Accordingly, there are m types of labels in the knowledge base, and each pair of positive prompt features and negative prompt features is associated with one of the labels and stored.

[0071] The preset conditions can be set according to actual needs, including but not limited to that the types of ground object categories corresponding to the positive prompt features and negative prompt features in the knowledge base can meet the segmentation requirements of the current original image, or the number of positive prompt features and negative prompt features in the knowledge base meets the use requirements. In addition, in actual use, the knowledge base can be updated in real time.

[0072] Next, the above steps S2-S5 are described in detail in conjunction with the embodiment. Each step implemented in this process can be implemented using an image segmentation model.

[0073] Based on an embodiment of the present application, the structure of the image segmentation model can be found in Figure 3 .like Figure 3 As shown, the image segmentation model includes a feature extraction module, a hint retrieval module, a hint embedding module and a mask decoding module.

[0074] In step S2, the feature extraction module is used to extract features from the original image to obtain image features and image class identifiers. The feature extraction module may be, for example, ViT in the pre-trained DINO model, which has been introduced above and will not be repeated here. In step S3, the prompt retrieval module is used to match the target prompt features corresponding to the original image from the knowledge base according to the image class identifier. The prompt retrieval module may, for example, obtain the target prompt features by calculating the similarity between the image class identifier and the positive prompt features in the knowledge base. Correspondingly, in step S4, the prompt embedding module uses the image features to perform feature enhancement processing on the target prompt features to obtain enhanced target prompt features. The prompt embedding module may, for example, use a cross attention network. In step S5, the mask decoding module is used to perform mask decoding processing on the image features and the enhanced prompt features to obtain a segmentation mask for the original image. The mask decoding module performs matrix multiplication on the image features and the enhanced target prompt features in the feature dimension to obtain an image segmentation mask, which includes the location information of multiple types of objects in the original image.

[0075] In the image segmentation method provided in this embodiment, the image segmentation model obtains target prompt features corresponding to the original image from the knowledge base to guide the model to segment the original image. For original images that were previously difficult to segment, accurate segmentation can also be performed with the help of prompt information in the knowledge base, thereby achieving the effect of accurately segmenting original images of various types of land objects. The image segmentation model has a strong generalization ability.

[0076] To further improve the segmentation accuracy, multi-scale features of the original image can be extracted, and image features can be obtained based on the multi-scale features. Figure 4 .

[0077] like Figure 4 As shown, in step S2, multi-scale feature extraction and feature fusion are performed on the original image to obtain the first image feature, the second image feature and the image class identifier, and the first image feature and the second image feature are used as the image features, wherein the scale of the first image feature is larger than the scale of the second image feature.

[0078] Optionally, the feature extraction module can use a dual feature extraction network. The dual feature extraction network includes a CNN network, a ViT network and a feature fusion network. The CNN network can use a pre-trained CNN network, such as a ResNeSt50d network pre-trained using ImageNet. ResNeSt50d is a model in the ResNeSt series, and its attention module uses a decentralized attention mechanism to improve the representation and generalization capabilities of the model. The ViT network can use ViT in the pre-trained DINO model. For ViT, please refer to the introduction in the foregoing text of the present invention, which will not be repeated here.

[0079] The process of performing multi-scale feature extraction and feature fusion on the original image using the dual feature extraction network may include the following steps S21-S25:

[0080] S21: Use the pre-trained CNN network to perform multi-scale feature extraction on the original image to obtain the first sub-feature, the second sub-feature, the third sub-feature and the fourth sub-feature.

[0081] Specifically, a pre-trained ResNeSt50d network is used, which includes a STEM layer and four cascaded residual blocks, which adopt a distraction block.

[0082] Assume that the resolution of the original image is A, and the original image is input into the ResNeSt50d network. After being processed by the STEM layer, it is input into the first residual block. The image features with resolutions of A / 4, A / 8, A / 16 and A / 32 are extracted respectively through four cascaded residual blocks. The four cascaded residual blocks output the first sub-feature, the second sub-feature, the third sub-feature and the fourth sub-feature respectively.

[0083] S22: Use the pre-trained ViT network to extract features from the original image to obtain the image class identifier and the fifth sub-feature.

[0084] Optionally, the Transformer encoder of the ViT network consists of 12 repeatedly stacked encoder blocks.

[0085] The original image is input into the ViT network, and after being processed by the Patch Embedding layer, it is input into the first encoder block. Each encoder block extracts the image features of A / 16 and inputs the results into the next encoder block. Finally, after 12 encoder blocks, the fifth sub-feature with a resolution of A / 16 is output.

[0086] S23: Using a feature fusion network to fuse the fourth sub-feature, the third sub-feature, and the fifth sub-feature to obtain the second image feature.

[0087] Specifically, the sampling layer is three layers. The fourth sub-feature is upsampled by 2 times using the third sampling layer and then fused with the third sub-feature to obtain a fused third sub-feature, and the fused third sub-feature is transformed in channel dimension and then added to the fifth sub-feature to obtain the second image feature.

[0088] S24: Using a feature fusion network to fuse the second image feature, the second sub-feature and the first sub-feature to obtain the first image feature.

[0089] Specifically, the second image feature is upsampled by 2 times using the second sampling layer and then fused with the second sub-feature to obtain a fused second sub-feature, and the fused second sub-feature is upsampled by 2 times using the first sampling layer and then fused with the first sub-feature to obtain the first image feature.

[0090] In this step, a dual feature extraction network is used to extract the first image feature and the second image feature. The first image feature is a fusion of multiple scale features of the original image, which makes full use of the multi-scale features and can improve the segmentation accuracy. The second image feature can express the semantic features of the original image, and using the second image feature to enhance the prompt feature can obtain a more accurate prompt feature.

[0091] Next, step S3 is described in detail.

[0092] In step S3, the similarity between the image class identifier and the positive prompt features in the knowledge base is calculated to obtain the top k positive prompt features in similarity order, and the k positive prompt features and their associated k negative prompt features are used as the target prompt features.

[0093] Specifically, for step S3, one possible implementation method is to retrieve from the knowledge base using the image class identifier as the retrieval condition, and by calculating the similarity between the image class identifier and the positive prompt features in the knowledge base, sorting according to the similarity, selecting the top k positive prompt features, and using the k positive prompt features and their associated k negative prompt features as target prompt features. The k positive prompt features and negative prompt features correspond to t types of land object categories, wherein each type of land object category corresponds to k1, k2, ..., kt positive prompt features and negative prompt features, respectively, and the sum of k1, k2, ..., kt is k.

[0094] The target prompt feature can correspond to a variety of ground object categories, and subsequent processing steps are performed according to the target prompt feature to obtain segmentation masks of the various ground object categories in the original image.

[0095] However, in this method, the k positive hint features obtained by searching based on the image class identifier may sometimes not cover all the object categories in the original image. Therefore, in some cases, the synchronous segmentation of some elements is achieved.

[0096] In order to segment all the object categories in the original image at one time, that is, the synchronous segmentation of all elements, in step S3, a preferred implementation method can be adopted. Specifically, it is assumed that the m types of object categories in the current knowledge base cover all the types of object categories in the original image. Using the image class identifier and the label in the knowledge base as parallel search conditions, for each label, the r positive prompt features ranked top in similarity with the image class identifier are searched from the knowledge base, and the corresponding r negative prompt features are associated, then a total of m×r positive prompt features and m×r negative prompt features are obtained for the m labels, and the m×r positive prompt features and m×r negative prompt features are used as the target prompt features, where m×r is equal to k.

[0097] For example, suppose that the total feature categories in the current knowledge base include three feature categories: water area, grassland, and building, and the labels are 1, 2, and 3 respectively. For any original image, the image class identifier and labels 1, 2, and 3 are used as parallel conditions, and the top five positive prompt features in similarity ranking are searched from the knowledge base for each label, and then 5 positive prompt features corresponding to water area, grassland, and building are obtained, totaling 15 positive prompt features, and a total of 15 negative prompt features are obtained by association. Among them, if some prompt features cannot be matched or the number is insufficient, they are supplemented with zero. In this way, the target prompt feature obtained can match all the feature categories in the original image. According to the target prompt feature, the subsequent steps can be processed to obtain the segmentation mask corresponding to all the feature categories in the original image, and the full-factor segmentation can be realized.

[0098] In this embodiment, target prompt features corresponding to various types of objects in the original image can be obtained through image class identification, so that the model can simultaneously segment multiple types of objects in the original image, thereby improving segmentation efficiency.

[0099] Next, in step S4, the prompt embedding module shown in FIG. 4 may be used to perform feature enhancement processing on the prompt feature using the second image feature to obtain an enhanced prompt feature.

[0100] Specifically, the hint embedding module adopts a cross attention network. The second image feature and the target hint feature are input into the cross attention network, and the cross attention network uses the cross attention mechanism to perform feature enhancement processing on the target hint feature to obtain an enhanced target hint feature. The target hint feature is mapped to a query vector (Query), and the second image feature is mapped to a key vector (Key) and a value vector (Value). The cross attention mechanism is used to continuously update the Query, and the Key and Value remain unchanged to obtain an enhanced target hint feature.

[0101] In the embodiment of the present invention, the second image feature represents the semantic feature of the original image, and the prompt embedding module only updates the target prompt feature without updating the second image feature, which can reduce the complexity of calculation and will not interfere with the second image feature.

[0102] In step S5, mask decoding is performed on the first image feature and the enhanced target prompt feature to obtain a segmentation mask for the original image. Figure 4 The mask decoding module in the image segmentation model shown is implemented.

[0103] Specifically, the mask decoding module performs matrix multiplication on the first image feature and the enhanced target prompt feature in the feature dimension to obtain multiple positive feature maps and multiple negative feature maps corresponding to multiple object categories. For each object category, a positive final feature map and a negative final feature map are obtained based on the corresponding positive feature maps and negative feature maps. Then, the positive final feature map and the negative final feature map are normalized using the softmax function to obtain a probability map, and the segmentation mask for a type of object in the original image can be obtained from the probability map. The above processing is performed simultaneously on multiple object categories in the original image to obtain multiple probability maps, and a segmentation mask for the original image can be obtained, which contains the location information of multiple object categories in the original image. Finally, the segmentation mask is upsampled to the resolution of the original image to complete the prompt segmentation.

[0104] Optionally, a corresponding positive final characteristic graph is obtained by calculating the average value of several positive characteristic graphs, and a corresponding negative final characteristic graph is obtained by calculating the average value of several negative characteristic graphs. Optionally, a corresponding positive final characteristic graph and a corresponding negative final characteristic graph can be obtained by respectively finding the maximum values ​​of several positive characteristic graphs and several negative characteristic graphs. The present application prefers the average value method.

[0105] Corresponding to step S3, in an optional manner, the mask decoding module performs matrix multiplication on the first image feature and the enhanced target prompt feature in the feature dimension to obtain k positive feature maps and k negative feature maps, corresponding to t types of land object categories, each type of land object category corresponds to k1, k2, ..., kt positive feature maps and negative feature maps. The k1 positive feature maps are averaged to obtain the positive final feature map of the first type of land object category. Similarly, the k2, ..., kt positive feature maps are averaged to obtain the positive final feature maps corresponding to the second to tth types of land object categories. Similarly, the negative final feature maps corresponding to each land object category are obtained. Then, each pair of positive final feature maps and negative final feature maps are normalized using the softmax function to obtain the corresponding t types of land object category probability maps.

[0106] Corresponding to step S3, in a preferred manner, the mask decoding module performs matrix multiplication on the first image feature and the enhanced target hint feature in the feature dimension to obtain m×r positive feature maps and m×r negative feature maps, corresponding to m types of ground object categories, each of which has r positive feature maps and r negative feature maps. The r positive feature maps and r negative feature maps of each ground object category are averaged to obtain m positive final feature maps and m negative final feature maps. Then, each pair of positive final feature maps and negative final feature maps are normalized using the softmax function to obtain probability maps corresponding to the m types of ground object categories.

[0107] Figure 5A schematic block diagram of an image segmentation device provided in an embodiment of the present application, such as Figure 5 As shown in , the device may include: an input unit 501, a feature extraction unit 502, a hint matching unit 503, a hint embedding unit 504 and a mask decoding unit 505. The main functions of each component unit are as follows:

[0108] An input unit 501 is used to obtain an original image and a knowledge base, wherein the knowledge base includes a plurality of pairs of positive prompt features and negative prompt features;

[0109] A feature extraction unit 502 is used to extract features from the original image to obtain image features and image class identifiers;

[0110] A prompt matching unit 503, configured to match a target prompt feature corresponding to the original image from the knowledge base using the image class identifier;

[0111] A prompt embedding unit 504 is used to perform feature enhancement processing on the target prompt feature by using the image feature to obtain an enhanced target prompt feature;

[0112] The mask decoding unit 505 is used to perform mask decoding processing on the image feature and the enhanced prompt feature to obtain a segmentation mask for the original image.

[0113] An embodiment of the present invention further provides an electronic device, comprising: a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the processor executes the program, the steps of any one of the methods described in the aforementioned method embodiments are implemented.

[0114] Furthermore, the electronic device also includes: at least one input device; at least one output device. The above-mentioned memory, processor, input device and output device are connected via a bus. Among them, the input device can be a camera, a touch panel, a physical button or a mouse, etc. The output device can be a display screen. The memory can be a high-speed random access memory (RAM) memory, or a non-volatile memory (non-volatile memory), such as a disk memory. The memory is used to store a set of executable program codes, and the processor is coupled to the memory.

[0115] The embodiment of the present invention further provides a computer-readable storage medium, which may be disposed in the electronic device in the above embodiments, and the computer-readable storage medium may be the electronic device in the above embodiments. The computer-readable storage medium stores a computer program, which, when executed by a processor, implements the steps of any of the methods in the above method embodiments.

[0116] Furthermore, the computer storable medium may also be a U disk, a mobile hard disk, a read-only memory (ROM), a random access memory (RAM), a magnetic disk or an optical disk, and other media that can store program codes.

[0117] It should be noted that the functional modules in the various embodiments of the present invention can be integrated into a processing module, or each module can exist physically separately, or two or more modules can be integrated into one module. The above-mentioned integrated modules can be implemented in the form of hardware or in the form of software functional modules. If the integrated modules are implemented in the form of software functional modules and sold or used as independent products, they can be stored in a computer-readable storage medium. Based on such an understanding, the technical solution of the present invention is essentially or the part that contributes to the prior art or all or part of the technical solution can be embodied in the form of a software product.

[0118] It should be noted that, for the convenience of description, the aforementioned method embodiments are all described as a series of action combinations, but those skilled in the art should be aware that the present invention is not limited by the described action sequence, because according to the present invention, certain steps can be performed in other sequences or simultaneously. Secondly, those skilled in the art should also be aware that the embodiments described in the specification are all preferred embodiments, and the actions and modules involved are not necessarily required by the present invention.

[0119] In the above embodiments, the description of each embodiment has its own emphasis. For parts that are not described in detail in a certain embodiment, reference can be made to the relevant descriptions of other embodiments.

[0120] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principle of the present invention should be included in the protection scope of the present invention.

Claims

1. An image segmentation method, characterized in that: The method comprises: S1: obtaining an original image and a knowledge base, wherein the knowledge base is pre-built and contains multiple pairs of positive prompt features and negative prompt features; S2: Extract features from the original image to obtain image features and image class identifiers; S3: using the image class identifier to match the target prompt feature corresponding to the original image from the knowledge base; S4: performing feature enhancement processing on the target prompt feature using the image feature to obtain an enhanced target prompt feature; S5: performing mask decoding processing on the image features and the enhanced target prompt features to obtain a segmentation mask for the original image.

2. The method according to claim 1, characterized in that The S3 includes: The similarity between the image class identifier and the positive prompt features in the knowledge base is calculated, and the k positive prompt features with the highest similarity ranking and their associated k negative prompt features are used as the target prompt features.

3. The method according to claim 1, characterized in that: The construction process of the knowledge base includes: Acquire a knowledge image, and annotate multiple types of objects in the knowledge image with masks and set labels respectively, to obtain an annotated knowledge image; For the annotated knowledge image, a sub-image in a mask area of ​​a certain type of feature is obtained, and a class identifier of the sub-image is obtained using a pre-trained ViT network to obtain the positive prompt feature; A sub-image that is not in a mask area of ​​a certain type of object is obtained, a class identifier of the sub-image is obtained using a pre-trained ViT network, the negative prompt feature is obtained, and the positive prompt feature, the negative prompt feature and the label of the certain type of object are associated and stored.

4. The method according to claim 1, characterized in that The S2 includes: Multi-scale feature extraction and feature fusion are performed on the original image to obtain a first image feature, a second image feature and the image class identifier, and the first image feature and the second image feature are used as the image features; wherein the scale of the first image feature is greater than the scale of the second image feature.

5. The method according to claim 4, characterized in that The method for obtaining the first image feature, the second image feature and the image type identifier includes: Use the pre-trained CNN network to extract multi-scale features from the original image, and obtain the first sub-feature, the second sub-feature, the third sub-feature, and the fourth sub-feature with decreasing scales in sequence; Using the pre-trained ViT network to extract features from the original image, to obtain the image class identifier and the fifth sub-feature; Using a feature fusion network to fuse the fourth sub-feature, the third sub-feature, and the fifth sub-feature to obtain the second image feature; The second image feature, the second sub-feature and the first sub-feature are fused using a feature fusion network to obtain the first image feature.

6. The method according to any one of claims 4-5, characterized in that: The S4 includes: The second image feature and the target prompt feature are input into a Transformer network, and the target prompt feature is enhanced by using a cross-attention mechanism to obtain the enhanced target prompt feature.

7. The method according to claim 6, characterized in that The S5 includes: A matrix multiplication operation is performed on the first image feature and the enhanced target prompt feature to obtain a segmentation mask for the original image, wherein the segmentation mask corresponds to a plurality of ground object categories of the original image.

8. An image segmentation device, characterized in that: The device comprises: An input unit, used for acquiring an original image and a knowledge base, wherein the knowledge base is pre-built and includes a plurality of pairs of positive prompt features and negative prompt features; A feature extraction unit, used to extract features from the original image to obtain image features and image class identifiers; A prompt matching unit, configured to match a target prompt feature corresponding to the original image from the knowledge base using the image class identifier; A prompt embedding unit, used for performing feature enhancement processing on the target prompt feature by using the image feature to obtain an enhanced target prompt feature; The mask decoding unit is used to perform mask decoding processing on the image feature and the enhanced target prompt feature to obtain a segmentation mask for the original image.

9. An electronic device, characterized in that: The electronic device comprises: A memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the computer program is executed by the processor, the steps of the method according to any one of claims 1 to 7 are implemented.

10. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 7 are implemented.

Images