Image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance

Abstract

The present invention discloses an image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance, belonging to the fields of image segmentation and deep learning. The method is performed by an image semantic segmentation system, which includes a MiT encoder and a lightweight MLP decoder. The MiT encoder includes an overlapping patch embedding module and a Transformer Block module. Under the Transformer network framework, the 2D neighborhood aggregation module calculates the similarity scores of adjacent tokens through Manhattan distance and merges them based on these scores. The token merging can be used with any trained Transformer without retraining the model, which can significantly reduce the computational burden of the attention mechanism and achieve significant acceleration while maintaining high performance values.

Description

Technical Field

[0001] The present invention belongs to the field of image segmentation and deep learning, and more specifically, relates to an image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance. Background Art

[0002] Image segmentation involves pixel-level recognition, which involves labeling the object category of each pixel in an image. Semantic segmentation only cares about the category of each pixel. If two objects of the same category are present in the input, semantic segmentation does not distinguish them as separate objects.

[0003] Image segmentation is a crucial area in computer vision. Image segmentation learning models enable machines to interpret visual information much like the human brain. For example, self-driving cars use semantic segmentation to understand and react to their surroundings in real time. While image segmentation models share some common functions with object detection models, they determine the specific shapes and boundaries of entities within an image. For example, by analyzing medical images and drawing precise boundaries around various objects within them, AI equipped with semantic segmentation can help detect anomalies and even suggest potential diagnoses.

[0004] As a key image processing technology, semantic segmentation research not only significantly improves image recognition accuracy and the ability to understand complex scenes, but also promotes the application of artificial intelligence in multiple important fields. From environmental perception in autonomous driving to intelligent analysis of medical images to efficient processing of remote sensing images, semantic segmentation demonstrates broad application prospects and significant research value.

[0005] Traditional image segmentation methods are typically based on convolutional neural networks (CNNs), which excel at processing local information and dependencies within a specific receptive field. However, with the increasing complexity and diversity of image data, effectively capturing global contextual information has become crucial for improving segmentation accuracy. The Transformer was introduced to image segmentation tasks to overcome the limitations of traditional methods in capturing global context and handling long-range dependencies, and to explore its potential in vision tasks.

[0006] Research on Transformer-based image segmentation aims to explore how to optimize model structures and algorithms to improve segmentation efficiency and computational speed. In the task of semantic segmentation with dense pixels, a fundamental architectural requirement is the ability to process high-resolution images to ensure that details can be detected and accurate predictions can be made. This often comes at a high computational cost. Therefore, it is imperative to find an appropriate trade-off between computational cost and performance. Summary of the Invention

[0007] In view of this, the purpose of the present invention is to address the problem that the semantic segmentation of high-resolution images in the Transformer network architecture is affected by the quadratic computational complexity of attention in the number of tokens. A 2D neighborhood aggregation module is used to reduce the number of tokens at runtime, and an image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance is proposed.

[0008] The technical solution adopted by the present invention to achieve the above-mentioned purpose is: an image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance, the method is executed by an image semantic segmentation system, the image semantic segmentation system includes a MiT encoder and a lightweight MLP decoder, the MiT encoder includes an overlapping patch embedding module and a Transformer Block module, the overlapping patch embedding module is used to extract patch images from the entire image, and flatten all patches and embed them into a high-dimensional vector, the number of Transformer Block modules is four, the four Transformer Block modules are connected in sequence, and the output of the previous Transformer Block module is used as the input of the next Transformer Block module, and the outputs of all Transformer Block modules are transmitted to the lightweight MLP decoder;

[0009] The specific steps include:

[0010] Step S1: The original image is input into the MiT encoder. The overlapping patch embedding module extracts 7*7 pixel patches. Each patch includes not only its own 7*7 pixel area but also the overlapping parts with adjacent patches. All 7*7 pixel patches are flattened and embedded into a high-dimensional vector and sent to the Transformer Block module.

[0011] Step S2: MiT encoder uses a pyramid structure to calculate multi-scale features across four stages, and each stage generates a dimension of feature map, where i∈{1, 2, 3, 4}; H, W and D i are the height of the original image, the width of the original image, and the number of channels after each stage i;

[0012] Step S3: The four feature maps obtained in step S2 are linearly transformed through a lightweight MLP decoder to unify their dimensions, and then upsampled and merged into one feature map. The final prediction result is converted into a pixel-level segmentation mask of the image through the MLP layer.

[0013] Among them, after the data enters the TransformerBlock module, before calculating the multi-head attention, spatial attenuation attention is performed on the key vector K and the value vector V, and then a 2D merging module, that is, a 2D maximum pooling layer with a step size of 2 and a pooling kernel is used to reduce the number of tokens. At the same time, the query vector Q is subjected to 2D neighborhood aggregation, that is, the similarity is calculated by Manhattan distance, and according to the pre-set similarity threshold, the tokens with similarity higher than the similarity threshold are merged, and the number of tokens is controlled by setting the reduction rate r, that is, the percentage of merged tokens; then the multi-head self-attention mechanism is calculated to generate a weighted feature vector; after the multi-head attention, the token merging is canceled, and the output result is passed to the next Transformer Block module.

[0014] Furthermore, in step S2, in order to obtain a pyramid structure, the resolution of the feature map is reduced by performing a 3*3 convolution with a stride of 2 at each stage.

[0015] Furthermore, spatial attenuation is implemented via 2D convolution with stride R.

[0016] According to a specific embodiment of the present invention, the similarity threshold is set to 0.8.

[0017] Furthermore, the lightweight MLP decoder works as follows:

[0018] First, the four feature maps output are unified to dimension C, the formula is:

[0019]

[0020] Where C is the target channel number of the feature map, the i-th feature map F i Linear transformation from C i The i-th feature map is converted from dimension to C dimension

[0021] Secondly, the four feature maps are uniformly upsampled to H / 4*W / 4*C and concatenated together to become H / 4*W / 4*4C. The formula is:

[0022]

[0023] Where H and W are the height and width of the original image respectively;

[0024] Then, H / 4*W / 4*4C is converted to H / 4*W / 4*C through the MLP layer, and the formula is:

[0025]

[0026] Concat is a concatenation operation that connects multiple feature maps together;

[0027] Finally, the MLP layer is used for classification, and the formula is:

[0028] M=Linear(C,N cls )(F)

[0029] where N cls is the number of categories.

[0030] Through the above design scheme, the present invention can bring the following beneficial effects: the present invention proposes an image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance. Under the Transformer network framework, the 2D neighborhood aggregation module calculates the similarity scores of adjacent tokens through Manhattan distance and merges them based on these scores. The token merging can be used with any trained Transformer without the need to retrain the model. It can significantly reduce the computational burden of the attention mechanism and achieve significant acceleration while maintaining high performance values. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] The drawings described herein are used to provide a further understanding of the present invention and constitute a part of this application. The exemplary embodiments of the present invention and their descriptions are used to understand the present invention and do not constitute improper limitations of the present invention. In the drawings:

[0032] Figure 1 This is a framework diagram of the image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance proposed in the present invention;

[0033] Figure 2 Schematic diagram of the TransformerBlock module structure in the present invention;

[0034] Figure 3 Schematic diagram of the structure of the 2D neighborhood aggregation module in the present invention. DETAILED DESCRIPTION

[0035] To make the objects, features, and advantages of the present invention more apparent and understandable, the technical solutions of the present invention are described clearly and completely below in conjunction with the accompanying drawings in the embodiments of the present invention. Obviously, the present invention is not limited to the following embodiments, and specific implementation methods can be determined based on the technical solutions of the present invention and actual conditions. To avoid obscuring the essence of the present invention, well-known methods, processes, and procedures are not described in detail.

[0036] Figure 1A framework diagram of an image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance is shown. The image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance is executed by an image semantic segmentation system. The image semantic segmentation system includes a MixTransformer (MiT) encoder and a lightweight MLP decoder. The MixTransformer (MiT) encoder includes an overlapping patch embedding module and a Transformer Block module. The overlapping patch embedding module is used to extract patch images from the entire image, flatten all patches and embed them into a high-dimensional vector and feed it into the Transformer Block module. There are four Transformer Block modules, and the four Transformer Block modules are connected in sequence. The output of the previous Transformer Block module is used as the input of the next Transformer Block module. At the same time, the outputs of all Transformer Block modules are transmitted to the lightweight MLP decoder.

[0037] It should be noted that the present invention maps the flattened patch vector to a high-dimensional space through a linear layer and inputs the processed high-dimensional vector into the Transformer Block. The linear layer maps the flattened vector to a higher-dimensional space through a weight matrix W and a bias term b. It is expressed as: z = Wx + b, where x is the flattened patch vector, W is the weight matrix of the linear transformation, b is the bias, and z is the high-dimensional embedding vector. This is prior art and will not be described in detail here.

[0038] From left to right, the four Transformer Block modules have output dimensions of The feature map of and H and W are the height and width of the feature map respectively, and C1, C2, C3, and C4 represent the dimensions.

[0039] In the present invention, for the convenience of description, the relative positional relationships of various components are described according to the layout of the drawings in the specification. For example, the positional relationships of left and right are determined according to the layout directions of the drawings in the specification.

[0040] Figure 2The figure shows a schematic diagram of the TransformerBlock module framework. After data enters the TransformerBlock module, spatially attenuated attention is applied to the key vector K and value vector V before multi-head attention is calculated. This spatial attenuation is achieved through a 2D convolution with a stride of R. The number of tokens is then reduced through a 2D merging module, which consists of a 2D max pooling layer with a stride of 2 and a pooling kernel. Simultaneously, 2D neighborhood aggregation is performed on the query vector Q, similarity is calculated using Manhattan distance, and merging is performed based on these scores. A multi-head self-attention mechanism is then applied to generate a weighted feature vector. After the multi-head attention, token merging is canceled, and the output is passed to the next TransformerBlock module.

[0041] It should be noted that in the present invention, after linear transformation of each pixel point in the data set, the query vector Q, key vector K and value vector V corresponding to each pixel point are obtained, which belongs to the existing technology and will not be described in detail.

[0042] Figure 3 Figure 2 shows the structure of the 2D Neighborhood Aggregation module, which calculates similarity scores using Manhattan distance in the two-dimensional space of an image. Darker colors indicate higher similarity scores between two locations. This score is used to prioritize tokens at certain locations during the aggregation process, helping to reduce computational complexity.

[0043] The specific implementation process of the image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance includes the following steps:

[0044] Step S1: The original image is input into the MiT encoder. The overlapping patch embedding module extracts 7*7 pixel patches. Each patch includes not only its own 7*7 pixel area but also the overlap with adjacent patches. This allows for capturing more fine-grained information in the image, which is crucial for dense tasks. However, this also increases the number of tokens and the computational burden. All 7*7 pixel patches are flattened and embedded into a high-dimensional vector, which is then fed into the Transformer Block module.

[0045] Step S2: MiT encoder uses a pyramid structure to calculate multi-scale features across four stages, and each stage generates a dimension of feature map, where i∈{1,2,3,4}; H, W and D i are the height of the original image, the width of the original image, and the number of channels after each stage i.

[0046] Step S3: The four feature maps obtained in step S2 are linearly transformed through a lightweight MLP decoder to unify their dimensions. They are then upsampled and merged into a single feature map. This is then converted through the MLP layer into the final prediction result. The final prediction result is the image segmentation mask, where each pixel is classified into a specific category.

[0047] Step S4: To address the issue of increasing the number of tokens in step S1, a 2D neighborhood aggregation module is designed. Specifically, the 2D neighborhood aggregation module calculates similarity using Manhattan distance and merges nearby tokens with high similarity based on these scores. For example, if the similarity threshold is set to 0.8, only tokens with a similarity higher than 0.8 will be merged. The number of tokens is controlled by setting the reduction rate r, which is the percentage of merged tokens. At the same time, before calculating the multi-head attention, spatial attenuation attention is performed on the key vector K and the value vector V. The 2D merging module is applied after the spatial reduction attention, and the token merging is canceled after the multi-head attention. The output is passed to the next TransformerBlock module.

[0048] In step S2, the MiT encoder avoids interpolation position encoding when inferring images of different resolutions, adapting to any test resolution without compromising performance. The four sequentially connected Transformer Block modules output a hierarchical structure of feature maps, enabling the encoder to generate high-resolution fine-grained features and low-resolution coarse-grained features.

[0049] In step S3, the attention of the lower levels of the lightweight MLP decoder tends to maintain locality, while the attention of the highest level is highly non-local. The lightweight MLP decoder combines local and global attention, aggregates information from different levels, and fuses these multi-level features to generate the final segmentation mask.

[0050] In step S1, compared to other Transformer-based networks such as SETR and Swin, the network architecture proposed in this paper emphasizes robustness and effectiveness. Specifically, it is robust against image interference while also being fast and highly accurate. Because SETR can only produce single-scale feature maps, which cannot effectively incorporate contextual information, a series of MixTransformer backbone networks are used to extract multi-scale features.

[0051] In step S2, to obtain a pyramid structure, the resolution of the feature map is reduced by performing a 3×3 convolution with a stride of 2 at each stage.

[0052] In step S3, a lightweight MLP decoder is used to avoid redundant computation. The key to implementing such a simple decoder is that the layered Transformer encoder has a larger effective receptive field than the traditional CNN encoder.

[0053] The lightweight MLP decoder is mainly divided into four steps:

[0054] First, the four feature maps output are unified to dimension C, the formula is:

[0055]

[0056] Where C is the target channel number of the feature map, the i-th feature map F i Linear transformation from C i The i-th feature map is converted from dimension to C dimension

[0057] Secondly, the four feature maps are uniformly upsampled (Upsample) to H / 4*W / 4*C and concatenated together to become H / 4*W / 4*4C. The formula is:

[0058]

[0059] Where H and W are the height and width of the original image respectively;

[0060] Then, H / 4*W / 4*4C is converted to H / 4*W / 4*C through the MLP layer, and the formula is:

[0061]

[0062] Concat is a concatenation operation that connects multiple feature maps together;

[0063] Finally, the MLP layer is used for classification, and the formula is:

[0064] M=Linear(C,N cls )(F)

[0065] where N cls is the number of categories;

[0066] In step S4, the Manhattan distance is used to calculate the distance between two locations so that the model can learn the relationship between spatial locations. Specifically, the locations are represented as multidimensional vectors, and then the Manhattan distance between each location and other locations is calculated as input features. The Manhattan distance d(A, B) between two points A(x1, y1) and B(x2, y2) on a two-dimensional plane is calculated as follows:

[0067] d(A,B)=|x1-x2|+|y1-y2|

[0068] Two n-dimensional vectors a(x 11 ,x 12 …x 1n ) and b(x 21 ,x 22 ,…x 2n ) between the Manhattan distance d 12 The formula is:

[0069]

[0070] In step S4, spatial attenuation is achieved by applying a 2D convolution with a stride of R. Before applying the convolution, the shape is N i ×D i (in )’s key vector K and value vector V are reshaped into After convolution, the tokens are flattened again; they are then reduced in number using a 2D merging module, consisting of a 2D max pooling layer with a stride of 2 and a pooling kernel, equivalent to 2D pooling. Simultaneously, 2D neighborhood aggregation is applied to the query vector Q, similarity is calculated using Manhattan distance, and merging is performed based on these scores. A multi-head self-attention mechanism is then performed to generate a weighted feature vector. After the multi-head attention, tokens are unmerged by copying them to their original positions, equivalent to 2D interpolation. The output is then passed to the next Transformer Block.

Claims

1. Image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance, characterized by: The method is performed by an image semantic segmentation system, which includes a MiT encoder and a lightweight MLP decoder. The MiT encoder includes an overlapping patch embedding module and a Transformer Block module. The overlapping patch embedding module is used to extract patch images from the entire image and flatten all patches and embed them into a high-dimensional vector. There are four Transformer Block modules, which are connected in sequence, and the output of the previous Transformer Block module is used as the input of the next Transformer Block module. At the same time, the outputs of all Transformer Block modules are transmitted to the lightweight MLP decoder. The specific steps include: Step S1: The original image is input into the MiT encoder. The overlapping patch embedding module extracts 7*7 pixel patches. Each patch includes not only its own 7*7 pixel area but also the overlapping parts with adjacent patches. All 7*7 pixel patches are flattened and embedded into a high-dimensional vector and sent to the Transformer Block module. Step S2: MiT encoder uses a pyramid structure to calculate multi-scale features across four stages, and each stage generates a dimension of The feature map of , 、 and They are the height of the original image, the width of the original image and each stage The number of channels after this; Step S3: The four feature maps obtained in step S2 are linearly transformed through a lightweight MLP decoder to unify their dimensions, and then upsampled and merged into one feature map. The final prediction result is converted into a pixel-level segmentation mask of the image through the MLP layer. Among them, after the data enters the Transformer Block module, before calculating the multi-head attention, the key vector K and the value vector V are spatially attenuated by 2D convolution with a stride of R, and then the number of tokens is reduced by a 2D merging module, that is, a 2D maximum pooling layer with a stride of 2 and a pooling kernel. At the same time, the query vector Q is aggregated by 2D neighborhood, that is, the similarity is calculated by Manhattan distance, and according to the pre-set similarity threshold, the tokens with similarity higher than the similarity threshold are merged, and the number of tokens is controlled by setting the reduction rate r, that is, the percentage of merged tokens; then the multi-head self-attention mechanism is calculated to generate a weighted feature vector; after the multi-head attention, the token merging is canceled, and the output result is passed to the next Transformer Block module.

2. The image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance according to claim 1, characterized in that In step S2, to obtain a pyramid structure, the resolution of the feature map is reduced by performing a 3×3 convolution with a stride of 2 at each stage.

3. The image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance according to claim 1, characterized in that Set the similarity threshold to 0.

8.

4. The image semantic segmentation method based on 2D neighborhood aggregation and Manhattan distance according to claim 1, characterized in that The lightweight MLP decoder works as follows: First, the four feature maps output are unified to the dimension , the formula is: ； in is the target channel number of the feature map, Feature maps Linear transformation from Dimension conversion to Dimension Feature maps ; Secondly, the four feature maps are uniformly upsampled to , and connected together to become , the formula is: ； in and are the height and width of the original image respectively; Then, Transformed by MLP layers , the formula is: ； in For the concatenation operation, multiple feature maps are connected together; Finally, the MLP layer is used for classification, and the formula is: ； in is the number of categories.

Images