Image encoder, high-precision interactive image segmentation method, electronic device, and storage medium

Abstract

The application discloses an image encoder, a high-precision interactive image segmentation method, an electronic device and a storage medium, solves the problems in the prior art, improves the calculation efficiency and segmentation quality of interactive segmentation, and can realize high-precision interactive segmentation on a high-resolution image. The image encoder comprises one image embedding module, 12 Transformer modules and four cyclic scanning modules which are constructed in series, three Transformer modules are one stage, and one cyclic scanning module is inserted before the last Transformer module in each stage; the Transformer module comprises an attention calculation module and a feedforward neural network module which are connected in series; the attention of the attention calculation module comprises general window attention and translation window attention, and the two kinds of attention are alternately distributed in the 12 Transformer modules. The image encoder mainly uses window attention calculation, constructs the calculation of general window attention and translation window attention, and cooperates with the cyclic scanning module to strengthen the global feature information fusion of the image.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of computer vision, in particular to an image encoder, a high-precision interactive image segmentation method, an electronic device and a storage medium. BACKGROUND

[0002] Under the impetus of massive data, deep networks have shown remarkable performance in various computer vision tasks. The availability of accurately labeled data is crucial for the success of deep networks, however, the process of manual labeling can occupy a large amount of time and computing resources. Therefore, the interactive image segmentation developed gradually has become an indispensable tool for labeling large-scale image datasets. This technology aims to achieve high-quality pixel-level labeling through limited interaction, and the interaction includes one or a combination of multiple methods such as scribble, bounding box, polygon, and click. Among the many interactive image segmentation methods, the click-based method has attracted much attention due to its simplicity and efficiency. Therefore, the present application focuses on the click-based method and attempts to solve the computational and efficiency problems in practical applications.

[0003] In the prior art, the research on click-based interactive image segmentation methods mainly focuses on real-time, low-resolution interactive segmentation models, such as SAM (Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollar, and Ross Girshick. Segment anything. arXiv:2304.02643, 2023.). However, this interactive segmentation model has two problems:

[0004] (1) First, the attention calculation of the visual Transformer (ViT) used by the encoder cannot use the efficient underlying implementation Flatten Attention;

[0005] (2) Second, the global attention of this encoder does not produce a performance improvement that matches its computational effort on high-resolution images.

[0006] Therefore, the SAM cannot perform reasoning on a running environment with insufficient display memory (for example, a common 3090 graphics card) when facing high-resolution images due to the inefficient attention calculation method, and even if it is run on a graphics card with sufficient display memory, it cannot guarantee that the segmentation effect will be improved by using high-resolution images, because the attention calculation of the SAM has a length extrapolation problem. Specifically, when the training process and the test process use images with different resolutions, the model lacks sufficient high-resolution image generalization due to the significant length difference in the image feature sequence it faces during attention calculation. For example, when a 1024x1024 image is used for training and a 2048x2048 image is used for testing as input, the resolution of the input image is inconsistent between the training stage and the testing stage, which may cause differences in numerical distribution during attention calculation, so that the model cannot effectively improve the segmentation effect on high-resolution images but brings great reasoning calculation cost. SUMMARY

[0007] The purpose of the present application is to provide an image encoder, a high-precision interactive image segmentation method, an electronic device and a storage medium, which solve the problems in the prior art, improve the calculation efficiency and segmentation quality of interactive segmentation, and can realize high-precision interactive segmentation on high-resolution images.

[0008] To achieve the above purpose, one of the solutions of the present application is:

[0009] An image encoder includes one image embedding module, 12 Transformer modules and four cyclic scanning modules constructed in series, every three Transformer modules form one stage, and a cyclic scanning module is inserted before the last Transformer module in each stage; the Transformer module includes an attention calculation module and a feedforward neural network module connected in series; the attention of the attention calculation module includes general window attention and translation window attention, and the two kinds of attention are alternately distributed in the 12 Transformer modules;

[0010] The image embedding module converts the image to be labeled into image features through linear transformation by convolution , wherein represents a real space, and represent the height and width of the image, respectively;

[0011] The Transformer module fuses the image features in space through the attention calculation module, and then strengthens the features through the feedforward neural network module;

[0012] The attention calculation module has two types of attention calculations:

[0013] The general window attention is to focus the image features. Reconstruction yields the reconstructed image features. ,in Indicates the predefined window size, superscript This is an abbreviation for reshape, followed by the features of each window. Conventional attention calculations are performed in parallel. After the attention calculations are completed, a reshaping operation is performed to restore the original feature form. Indicates the feature number after windowing;

[0014] The translation window attention is achieved by focusing on the image features. Fill the left and top respectively and A zero vector of units, simultaneously in the image features Fill the right and bottom and The zero vector of one unit is used to obtain the expanded image features. Then, the general window attention is calculated, where This indicates the predefined window size. and These represent the preset number of features for horizontal and vertical axis translation, respectively, with values ​​ranging from 0 to... between.

[0015] The feedforward neural network module first processes the image features After general matrix multiplication mapping, the mapped image features are obtained. superscript This is an abbreviation for mapping; the mapped image features After nonlinear activation calculation, it is converted back into image features via another linear mapping. ;

[0016] The loop module first processes the input image features Flattening is a characteristic of short sequences Then, the short sequence features are... Long sequence features are obtained by concatenating copies three times. superscript and These represent short sequences and long sequences, respectively; then, coefficients are generated through multiple parallel linear mappings. and ,in and is the weight of this linear mapping to generate two coefficients; again the coefficients and are input into the calculation library of the selective state space model to perform the scan calculation process of the sequence features, which is described in the form of iteration , wherein represents the calculation result of the th element in the sequence, represents the feature corresponding to the th element, and have the same dimension, represents the sequence number of the element in the sequence; after traversing and calculating all the elements in the sequence, the is output as the final result of the scan calculation process, and then the merged sequence features are obtained by equal-length division and merging in the form of summation, and finally the merged sequence features are converted into image features as the output of the loop module.

[0017] Preferably, the attention calculation module utilizes the calculation function library Flash Attention to perform the calculation process when performing general window attention calculation.

[0018] The second solution of the present application is:

[0019] A high-precision interactive image segmentation method, applying the decoder and the image encoder, comprising the following steps:

[0020] Step 1: pre-processing the image to be labeled by the image encoder to obtain image features ; then reducing the channel number through convolution calculation to obtain the final image features ;

[0021] Step 2: converting the user input click information into click encoding , and converting the rough mask information provided by the user into mask encoding , to obtain interactive encoding containing click encoding and mask encoding , and the interactive encoding is updated each time the user interacts through clicking;

[0022] Step 3: inputting the interactive encoding and the learnable Query encoding into the decoder, and combining with the final image features in step 1 to obtain Query encoding fused with interactive information​ Up-sampled image features with fusion interaction information The calculation formula is as follows:

[0023] ;

[0024] Wherein Indicates a cross attention operation;

[0025] Step 4, the up-sampled image features with fusion interaction information are multiplied by point operation Query code with fusion interaction information Further fusion, get the predicted mask The calculation formula is as follows:

[0026] ;

[0027] Wherein Indicates the transpose of ;

[0028] Step 5, in the process of model training and standard test, according to the final predicted mask and mask label, find the largest error connected region, select the center pixel of the connected region as the position of the next simulated click, and iterate the next model prediction; In the real test environment, the iterative prediction is carried out according to the user's interactive click feedback, and the segmentation quality is gradually improved.

[0029] The third solution of the application is:

[0030] An electronic device comprises a processor, a memory and an application program; the application program is stored in the memory and is configured to execute the high-precision interactive image segmentation method by the processor.

[0031] The fourth solution of the application is:

[0032] A computer readable storage medium, the storage medium stores a computer program; the computer program is executed in the computer, so that the computer executes the high-precision interactive image segmentation method.

[0033] After adopting the above technical scheme, the application has the following technical effects:

[0034] The present application optimizes the image encoder in the preprocessing stage of a large visual model, can extract high-quality and reusable image features from the image to be labeled, and realizes the fusion of the global feature information of the image to be labeled, so that the model is more competitive on high-resolution images while maintaining the original efficient calculation, and provides an effective solution for high-precision interactive image segmentation; the image encoder integrates multiple windowed attention calculation modes, which significantly improves the ability of the image encoder trained at a regular size resolution to perform high-precision interactive segmentation on high-resolution images.

[0035] The method of the present application uses an image encoder to extract high-quality and reusable image features from the image to be labeled, which can realize interactive segmentation on a personal computer or a mobile device with low power consumption by a lightweight interactive segmentation module, and the image features extracted by the large visual model and the clicks of the labeler will be used as the input of the interactive segmentation module; under the premise of low calculation delay, the overall interactive segmentation model can improve the high-precision segmentation effect on high-resolution images. The model architecture using the method can only use the simplest distillation training, but the trained model can still exceed the teacher model used for distillation on high-resolution images under the premise of lower calculation delay, achieving extremely high segmentation effect. BRIEF DESCRIPTION OF DRAWINGS

[0036] Figure 1 The network structure diagram of the image encoder of the embodiment of the present application is shown in the figure.

[0037] Figure 2 The experimental data table of the embodiment of the present application is shown in the figure. DETAILED DESCRIPTION

[0038] In order to further explain the technical scheme of the present application, the present application will be described in detail through specific embodiments.

[0039] The present application discloses an image encoder for extracting high-quality and reusable image features from an image to be labeled before the interactive process of an interactive image segmentation method. The image encoder mainly uses windowed attention calculation, constructs two attention calculation methods: general windowed attention and translation windowed attention, and cooperates with a cyclic scanning module to strengthen the global feature information fusion of the image.

[0040] REFERENCE Figure 1As shown, the above image encoder mainly includes one image embedding module, 12 Transformer modules and four cyclic scanning modules constructed in series, every three Transformer modules is one stage, and one cyclic scanning module is inserted before the last Transformer module in each stage; the Transformer module includes attention calculation module and feed-forward neural network module (FFN) in series; the attention of the attention calculation module includes general window attention (corresponding to the general module in Figure 1 Figure 1 The above two attentions are alternately distributed in the 12 Transformer modules, for example, the attention calculation module of the first Transformer module uses general window attention, and the second Transformer module uses translation window attention, and vice versa.

[0041] The image embedding module converts the originally three-channel image to be labeled through linear conversion by convolution to obtain image features , where represents the real space, and represent the height and width of the image respectively; and the subsequent 12 Transformer modules serially process the image features .

[0042] The Transformer module receives image features as input and outputs image features of the same dimension; specifically, the image features are first fused in space by the attention calculation module, and then the enhanced features are obtained by the feed-forward neural network module.

[0043] The two attention calculations of the attention calculation module are:

[0044] a. General window attention

[0045] The general window attention windowizes the image features by reshaping operation: the image features are reshaped to obtain the reshaped image features , where represents a predefined window size (preset by the user, for example, set to 16), and the superscript represents the abbreviation of reshape, which is used to distinguish different image features (the superscript is mainly used for distinction, and will not be described below).

[0046] Then, the features of each window ​Conventional attention calculations are performed in parallel. After the attention calculations are completed, a reshaping operation is performed to restore the original feature form. Indicates the feature number after windowing.

[0047] b. Shifting window attention

[0048] The overall process of translation window attention is similar to that of general window attention, the main difference being the windowing method: translation window attention primarily involves applying preliminary image features... Perform a fixed number of translations in both the horizontal and vertical directions, for example, move the object originally located at the... Line number The image features of the column are in the horizontal direction ( (axial direction) and longitudinal direction ( Translate in the axial direction respectively One and Units ( (This represents the predefined window size), so when the window is reshaped and then windowed, it is equivalent to indirectly moving the window, thus achieving the effect of panning the window.

[0049] The translation operation process is simplified in this invention: by using image features Fill the left and top respectively and A zero vector of units, simultaneously in image features Fill the right and bottom and The zero vector of one unit is used to obtain the expanded image features. Then, the general window attention calculation described above is performed, which is equivalent to calculating the image features. The above translation operation was performed; among which, This indicates the predefined window size. and These represent the preset number of features for horizontal and vertical axis translation, respectively, with values ​​ranging from 0 to... between.

[0050] The feedforward neural network module consists of two simple linear mapping layers, first processing image features... After general matrix multiplication mapping, the mapped image features are obtained. superscript This is an abbreviation for mapping; mapping image features. After nonlinear activation calculation, it is converted back into image features via another linear mapping. .

[0051] The loop module first processes the input image features Flattening is a characteristic of short sequences Then, short sequence features Long sequence features are obtained by concatenating copies three times. superscript and These represent short sequence and long sequence, respectively; then, coefficients are generated through multiple parallel linear mappings. and ,in and This linear mapping is used to generate the weights of the two coefficients; then the coefficients are... and The sequence features are scanned and calculated using a computational library of a selective state-space model, and this process is described as iterative. ,in This indicates the first [number] in this sequence. The calculation result of each element, Indicates the first The features corresponding to each element and Having the same dimensions This represents the index of an element in the sequence; it iterates through all elements in the sequence and outputs the index. As the final result of this scanning calculation process, the merged sequence features are obtained by dividing the sequence into equal-length segments and summing them. Finally, the merged sequence features are transformed into image features through a reshaping operation. As the output of the loop module.

[0052] Through the above solution, the present invention has the following outstanding advantages:

[0053] This invention optimizes the image encoder in the preprocessing stage of large-scale visual models, enabling the extraction of high-quality and reusable image features from the image to be labeled, and achieving enhanced global feature information fusion of the image to be labeled. This makes the model more competitive on high-resolution images while maintaining the original high-efficiency computation, providing an effective solution for high-precision interactive image segmentation.

[0054] In some embodiments of the present invention, when performing general window attention calculation, the attention calculation module utilizes the mature calculation function library Flash Attention to run the calculation process, which can reduce the actual video memory usage.

[0055] The application also discloses a high-precision interactive image segmentation method, aiming to improve the processing effect of interactive segmentation on high-resolution images, and using the image encoder to extract high-quality and reusable image features from the image to be labeled, so that the interactive segmentation module can be executed on a low-power device such as a personal computer or a mobile device, and the image features extracted by the large visual model and the clicks of the labeler will be used as the input of the interactive segmentation module. The interactive image segmentation method has the advantages that under the premise of low calculation delay, the overall interactive segmentation model can improve the high-precision segmentation effect on high-resolution images. The specific steps of the method are as follows:

[0056] Step 1, pre-processing the image to be labeled by the image encoder to obtain image features ; and then reducing the channel number by convolution calculation to obtain the final image features . ;

[0057] Step 2, converting the user input click information into click encoding , and converting the rough mask information provided by the user into mask encoding , to obtain interactive encoding containing click encoding and mask encoding , and updating the interactive encoding every time the user interacts by clicking;

[0058] Step 3, sending the interactive encoding and the learnable Query encoding into the decoder, and combining the final image features in step 1 to obtain the Query encoding fused with interactive information and the up-sampled image features fused with interactive information , and the calculation formula is as follows:

[0059] ;

[0060] Wherein represents cross-attention operation;

[0061] Step 4, further fusing the up-sampled image features fused with interactive information and the Query encoding fused with interactive information by dot multiplication operation to obtain the predicted mask , and the calculation formula is as follows:

[0062] ;

[0063] Wherein represents Transpose of;

[0064] Step 5: In the model training and standard testing process, based on the final prediction mask and mask label, find the connected region with the largest error, select the center pixel of the connected region as the simulated next click position, and iterate to perform the next model prediction; in the real test environment, iterative prediction is performed based on the user's interactive click feedback to gradually improve the segmentation quality.

[0065] The technical effects of the present invention will be further illustrated below through simulation experiments and experimental data.

[0066] 1) Simulation conditions

[0067] This invention was developed on the Ubuntu platform, and the developed deep learning framework is based on PyTorch, with Python being the primary language used.

[0068] 2) Simulation content

[0069] Using the COCO2017, LVIS, and HQSeg44K-Training datasets, the SAM model was first used as the teacher model on the first two datasets to distill the large-scale vision model applied to this invention. Then, fine-tuning training was performed on the last dataset using a simulated interactive process. Finally, HQSeg44K-Validation and DAVIS were used as test sets for testing. See [link to documentation]. Figure 2 The figures show the segmentation results of this invention and other methods on various test datasets. A direct numerical comparison reveals that this invention performs best compared to other methods. The evaluation metrics NoC (Number of Clicks) represent the number of clicks required to reach a certain IoU threshold, and SPC (Seconds Per Click) represents the time spent on inference for one click, measuring the model's efficiency in pixel segmentation of potentially large or small objects. The NoC comparison shows that, under the same training conditions, this method surpasses other state-of-the-art methods in segmentation quality, while also exhibiting lower computational latency (SPC) compared to previous best methods, meaning it has a lower computational burden and higher segmentation quality.

[0070] The present invention also discloses an electronic device, including a processor, a memory, and an application program; the application program is stored in the memory and configured to be executed by the processor using the methods described above.

[0071] The application further discloses a computer readable storage medium, the storage medium stores a computer program; the computer program is executed in the computer, and the computer program causes the computer to execute the method.

[0072] The above-mentioned embodiments and drawings are not intended to limit the product form and style of the application, and any appropriate changes or modifications made by those skilled in the art to the application should be considered as not departing from the patent scope of the application.

Claims

1. An image encoder characterized in that: it comprises one image embedding module, 12 Transformer modules and four recurrent scanning modules, which are constructed in series, every three Transformer modules form a stage, and one recurrent scanning module is inserted before the last Transformer module in each stage; the Transformer module comprises an attention calculation module and a feedforward neural network module in series; the attention of the attention calculation module comprises general window attention and translation window attention, and the two kinds of attention are alternately distributed in the 12 Transformer modules; The image embedding module obtains image features by linear conversion on the to-be-labeled image in a convolution manner wherein represents a real number space, and respectively represent the height and width of the image; The Transformer module processes the image features through the attention calculation module. Spatial fusion is performed, and then the features are enhanced through the feedforward neural network module; the two kinds of attention calculation of the attention calculation module are respectively: The general window attention is to apply the attention mechanism to the image features reshaped to get reshaped image features wherein denotes a predefined window size, the superscript denotes the abbreviation of reshape, and then the features of each window are The regular attention calculation is performed in parallel, and after the attention calculation is completed, the reshape operation is performed to return to the original feature form, wherein denotes the sequence number of the features after windowing; The translation window attention is obtained by padding zero vectors of and units on the left and top of the image feature , respectively, while padding zero vectors of and units on the right and bottom of the image feature , respectively, to obtain an expanded image feature , and then performing the calculation of the general window attention, wherein represents a predefined window size, and respectively represent preset horizontal axis and vertical axis translation feature numbers, and the values are between 0 and . The feedforward neural network module first converts the image features After a general matrix multiplication mapping, the mapped image features are obtained where the superscript denotes the abbreviation mapping; the mapped image features are converted back to image features via another linear mapping after a non-linear activation computation The cyclic scanning module first processes the input image features Flattening is a characteristic of short sequences Then, the short sequence features are... Long sequence features are obtained by concatenating the copies three times. superscript and These represent short sequences and long sequences, respectively; then, coefficients are generated through multiple parallel linear mappings. and ,in and This linear mapping is used to generate the weights of the two coefficients; the coefficients and are input into the calculation library of the selective state space model for the scan calculation process of the sequence features, which is described in the form of iteration wherein represents the calculation result of the th element in the sequence, represents the feature corresponding to the th element, has the same dimension as , represents the sequence number of the element in the sequence; after traversing and calculating all the elements in the sequence, the is output as the final result of the scan calculation process, and then the merged sequence features are obtained by equal-length division and merging in the form of summation, and finally the merged sequence features are converted into image features as the output of the cyclic scan module.

2. The image encoder of claim 1 characterized in that: when the attention calculation module performs general window attention calculation, it uses the calculation function library FlashAttention to run this calculation process.

3. A high-precision interactive image segmentation method, applying a decoder and an image encoder as claimed in any one of claims 1 or 2, characterized in that comprising the following steps: Step 1, pre-processing the image to be labeled by the image encoder to obtain image features ; The image features are reduced in channel number by convolution calculation to obtain final image features The image features are reduced in channel number by convolution calculation to obtain final image features ; Step 2: The image encoder processes the user input. Each click information is converted into a click code. The user-provided coarse mask information is converted into mask encoding. Get the click code and mask encoding The interaction code is updated every time the user interacts by clicking; Step 3, combine the interaction code with the learnable Query code into the decoder and combine with the final image feature in Step 1 to get the Query code with fused interaction information with the up-sampled image feature with fused interaction information The formula is as follows: ； wherein denotes a cross-attention operation; Step 4, up-sampled image features fused with interaction information by point-wise multiplication operation Query encoding fused with interaction information Further fusion to get predicted mask The calculation formula is as follows: ； wherein denotes the transpose of Step 5. In the process of model training and standard testing, according to the final prediction mask and the mask label, find the largest error connected region, select the center pixel of the connected region as the position of the next simulated click, and iterate the next model prediction; In the real test environment, iterative prediction is performed according to the user's interactive click feedback to gradually improve the segmentation quality.

4. An electronic device characterized in that: it comprises a processor, a memory and an application program; the application program is stored in the memory and is configured to be executed by the processor to perform the high-precision interactive image segmentation method of claim 3.

5. A computer readable storage medium characterized in that: the storage medium stores a computer program; when the computer program is executed in the computer, the computer program makes the computer execute the high-precision interactive image segmentation method of claim 3.

Images