Image segmentation method based on CNN and Transformer

Abstract

The present invention discloses an image segmentation method based on CNN and Transformer. The present invention passes an image to a CNN module, extracting local low-level features of the image through layer-by-layer convolution and pooling; passes the same image to a Transformer module, which segments the image into appropriate blocks and performs feature extraction, extracting and preserving the global information and long-range dependencies of the image; rearranges the extraction results of the two modules executed in parallel; passes the intermediate results of each layer of the CNN module to an upsampling layer, and passes the results of the last layer through an additional Transformer layer to further enhance the global features; finally, the fused data and the data passed from the CNN module are uniformly passed to the upsampling module, which restores the image resolution by using global and local features. The present invention improves the accuracy of image semantic segmentation.

Description

Technical Field

[0001] The present invention belongs to the field of semantic segmentation of images and relates to a parallel medical image segmentation method combining Convolutional Neural Networks (CNN) and Transformer. Background Art

[0002] Medical image segmentation is essential for the development of healthcare systems, particularly for disease diagnosis and treatment planning. Diagnosis requires specialized personnel to segment tumors and organs within images. However, manual image segmentation consumes significant human resources. Therefore, the use of artificial intelligence (AI) for medical image segmentation has become an important approach to reducing the cost of medical image analysis. However, due to the complexity of medical images, achieving accurate segmentation of complex and intertwined images of organs is extremely difficult. Consequently, there is an urgent need for automated, high-precision image segmentation technology in the medical field.

[0003] Over the past decade, various technologies have emerged that use artificial intelligence to help researchers segment medical images, such as the UNet architecture, DeepLab architecture, and contrastive learning based on fully convolutional neural networks; and Vision Transformer, SwinUNet, and TransUNet, which were inspired by the field of natural language processing. Among them, the UNet architecture has become one of the most widely used architectures due to its good performance and accuracy. Existing UNet-based technologies focus on accurately segmenting tissues or organs with different semantics, or preventing tissues or organs with the same semantic information from being over-segmented, but rarely consider combining the two. For semantic segmentation, the two requirements are completely different. Existing methods may lead to over-segmentation or incomplete segmentation of the same or different organs. This analysis is mainly based on the following facts:

[0004] 1. Technology based on fully convolutional neural networks uses convolution to extract semantic information from images, then compresses the image and enhances its features through downsampling. Multiple layers of downsampling yield semantic information at different resolutions and levels. Upsampling then combines the extracted semantic information at different levels to restore the image resolution. CNNs naturally extract low-level features from images, so this technology can learn features with the same semantic information.

[0005] 2. In contrast, Transformer-based technology first divides the image into blocks of equal size and serializes each block as a piece of semantic information. The encoder encodes this serialized information and its positional information, allowing each block to learn semantic features from each other. This encoded semantic information is then passed to the decoder, which decodes and reconstructs the input data to restore the original image. This technology effectively learns long-range dependencies within the image through mutual learning between serialized blocks.

[0006] However, due to the different directions and focuses of the two technologies, existing technologies often tend to ignore one of the two technologies and excessively pursue better results in one direction, ignoring the connection between the two technologies and failing to find a good balance between the two directions. Summary of the Invention

[0007] In view of the shortcomings of the existing technology, the present invention proposes an image segmentation method based on CNN and Transformer.

[0008] A CNN and Transformer-based image segmentation method, the method specifically includes the following steps:

[0009] Step 1: Extract global semantic features of the image;

[0010] The image is input into the Transformer encoding part, first passing through the patch segmentation layer, then passing through four downsampling layers, where the patch is first fused, and then input into the SwinTransformerBlock to extract global semantic features;

[0011] Step 2: Extract local detail features of the image through the CNN module;

[0012] Step 2-1: Convolution and downsampling. The whole CNN stage performs four downsamplings. The final image format is

[0013] Step 3: Feature reconstruction and merging;

[0014] After both the CNN module and the Transformer module have extracted features, the feature data of the convolutional network needs to be reconstructed because they have different organization specifications for feature labeling data. The reconstructed data is then merged with the output of the latter in the dimension channel. The merged feature data is passed through a Transformer block, and the output result is reconstructed again to make it conform to the data format required for upsampling.

[0015] Step 4: Decode the feature information and obtain the segmentation result;

[0016] Step 4-1: The decoder inputs the merged and reconstructed data into the upsampling layer, gradually restoring the resolution of the feature image. After upsampling, it is merged with the feature images of different levels retained in CNN through jump connections, and then the merged image is convolved. After three consecutive upsampling and three jump connections, the image size is restored to Finally, a patch expansion layer with 4x upsampling is used to restore the image to H×W×C.

[0017] Step 4-3: Segmentation prediction, pass the upsampled feature labels to the linear mapping layer, and output the final pixel-level segmentation prediction result.

[0018] Preferably, the patch segmentation layer is specifically:

[0019] The input RGB image is segmented into non-overlapping patches through the patch segmentation module. Each patch is regarded as a "marker" and its feature is set as the concatenation of the original pixel RGB values;

[0020] Preferably, the four downsampling layers are as follows:

[0021] In the first downsampling layer, the patch merging layer is replaced by a linear embedding layer, and the patch is reconstructed into a dimension of The structure is as follows: the patch merging layer or the linear embedding layer is followed by the SwinTransformerBlock to extract features.

[0022] As an example, the patch merging layer is used to reduce the number of markers and strengthen the global features, and the H and W of the patch become the original The feature dimension becomes twice as large as before; the patch merging layer passes the merged patches to SwinTransformerBlock for feature extraction.

[0023] As a preferred method, the construction method of SwinTransformerBlock is to replace the standard multi-head self-attention module in the original Transformer with a module based on a shift window. Since it is a layered structure, the Transformers in each block must exist in pairs. In the window multi-head self-attention layer of the first Transformer, the patch is divided into four blocks and its global features are learned. The extracted features are then input into the second Transformer. In the second shift window multi-head self-attention layer, the patches are re-divided and combined according to specific rules to extract patch features. The computational complexity of standard multi-head self-attention and shift window-based self-attention is as follows:

[0024] Ω(MSA)=4hwC 2 +2(hw)2 C (1)

[0025] Ω(W-MSA)=4hwC 2 +2M 2 hwC (2)

[0026] Where h is the number of blocks divided in the height direction of the image, and w is the number of blocks divided in the width direction of the image. The computational cost in formula (1) is proportional to the square of the number of patches hw, while the computational cost in formula (2) is linearly related to hw. The computational cost of the Transformer based on the shifted window is significantly better than that of the original Transformer.

[0027] Preferably, the CNN stage performs downsampling four times, specifically:

[0028] The same image is passed to the convolutional network part at the same time. In order to keep pace with the patch size of the Transformer part, the convolutional network part performs 4 downsampling. Before the first downsampling, the image information is extracted by convolution operation, and then downsampled through the first pooling layer to reduce the amount of calculation and expand the receptive field. After pooling, the H and W of the image become and The dimension becomes C. In the next three downsampling, each downsampling will change the H and W of the image to the original The dimension is doubled, and after each downsampling, a copy of the extracted image feature markers is retained for use in the jump connection of the upsampling part during the decoding process.

[0029] Preferably, the upsampling described in step 4-1 is specifically as follows:

[0030] It consists of multiple upsampling steps to decode the hidden features and output the final segmentation mask. After encoding the image through the hybrid encoder, it goes through multiple upsampling blocks to reach the full resolution, where each block consists of a 2× upsampling operator, a 3×3 convolutional layer, and a ReLU layer.

[0031] Preferably, the jump connection is specifically as follows: the cascade upsampling and the hybrid encoder form a U-shaped architecture, direct upsampling of low-resolution patches to high-resolution patches will lose feature information, while the feature map retained by CNN preserves the low-level features of the image, and feature aggregation of different resolution levels is achieved through jump connections. The original image is restored as much as possible by fusing the intermediate features preserved after CNN downsampling with the upsampled features. During the first upsampling, the patch features tend to global semantic information as a whole, so in the first jump connection, the feature map will pass through an additional Transformer block to further enhance the global characteristics of the low-level feature map retained by CNN.

[0032] The effect of the present invention compared with the prior art: the present invention adopts a U-shaped semantic segmentation framework to achieve efficient and accurate semantic segmentation.

[0033] (1) Transformer is used to encode the input tokenized image blocks as the input sequence for extracting the global context. Due to its inherent global self-attention mechanism, it can extract more long-range dependency information and solve the problem of missing global semantic information.

[0034] (2) A coding method that mixes CNN and Transformer is proposed. The improved Transformer, namely SwinTrnasformer, works in parallel with the classic ResNet50, which improves the feature extraction ability of the model and successfully makes up for the problem that Transformer is insufficient in extracting low-level detail features.

[0035] (3) A new semantic segmentation method is proposed, which integrates the improved encoder into the U-Net model. Through the convolution operation of CNN and the global self-attention mechanism of Transformer, it can take into account both the local detail features and global semantic information of the image, improving the problem of incomplete feature extraction in the existing method. In addition, the original features of the image can be restored through skip connections and upsampling in the decoder, further improving the accuracy of semantic segmentation. BRIEF DESCRIPTION OF THE DRAWINGS

[0036] Figure 1 This is the overall flow chart;

[0037] Figure 2 Schematic diagram of the hierarchical Transformer block;

[0038] Figure 3 This is a schematic diagram of the decoder part. DETAILED DESCRIPTION

[0039] In view of the shortcomings of the existing technology, this paper proposes a method for image semantic segmentation in the medical field. The overall process is shown in the attached figure. Figure 1As shown in the figure, it is divided into two parts. The first part is a parallel hybrid encoder composed of CNN and Transformer to encode the input image. First, the image is input into the Transformer module, and the image is segmented through the patch segmentation layer. Then, the patch is downsampled and encoded through the patch fusion layer and Transformer block. At the same time, the same image is subjected to feature extraction through the CNN module to obtain the detailed features of the image. After convolution and downsampling, image feature labels of different levels can be obtained. After the Reshape process, the feature data is reconstructed so that the two features extracted above have the same organizational form. The second part is the decoding part, which merges the feature data extracted by CNN and Transformer and inputs it into the decoder. The decoder gradually restores the original resolution of the image through upsampling, and combines the global semantic information obtained by upsampling with the local detail information obtained by skip connection to restore the original information of the image to the greatest extent. Finally, through a linear mapping layer, the pixel-level segmentation prediction result is obtained based on the extracted image feature labels.

[0040] The present invention will be described in detail below with reference to the accompanying drawings, and the specific steps are as follows:

[0041] Step 1: Extract global semantic features of the image through the Transformer module.

[0042] Step 1-1: Patch Segmentation

[0043] Take an image of the format H×W×3 as input. In the segmentation layer, the image is divided into 4 parts in both the W and H directions and concatenated in the dimension direction. The dimension of the segmented image is 4×4×3=48, and the patch size is

[0044] Step 1-2: Linear embedding, after passing the linear embedding layer, the feature dimension of the patch is projected to an arbitrary dimension C, where C is 64.

[0045] Steps 1-3: Transformer Block Encoding

[0046] As attached Figure 2 As shown, the specific process of the hierarchical structure can be expressed as:

[0047]

[0048]

[0049]

[0050]

[0051] Among them, z l-1is the input, z l+1 For output, and The intermediate results for each hierarchical structure.

[0052] Steps 1-4: Patch merging, similar to the downsampling of convolutional networks, the patch merging layer is used to downsample and increase the dimension.

[0053] Step 2: Extract local image detail information through the CNN module.

[0054] Step 2-1: Convolution and downsampling. The whole CNN stage performs four downsamplings. The final image format is

[0055] Step 2-2: Feature reconstruction and merging, reconstruct the output of CNN into the same structure as the output of Transformer, and then merge the two.

[0056] Step 3: Decode the feature information and obtain the segmentation result.

[0057] Step 3-1: Feature enhancement and reconstruction: pass the merged feature data through a Transformer block, and then reconstruct the output result again to make it conform to the data format required by upsampling.

[0058] Step 3-2: If Figure 3 As shown in Figure 2, the decoder inputs the merged and reconstructed data into the upsampling layer, gradually recovering the resolution of the feature image. After upsampling, it is merged with the feature images of different levels retained in CNN through jump connections, and then the merged image is convolved. After three consecutive upsampling and jump connections, the image size is restored to Finally, a patch expansion layer with 4x upsampling is used to restore the image to H×W×C.

[0059] Step 3-3: Segmentation prediction, pass the upsampled feature labels to the linear mapping layer, and output the final pixel-level segmentation prediction result.

[0060] For the experimental results of image segmentation: the average Dice Similarity Coefficient (DSC) and the average Hausdorff distance are used as evaluation indicators.

[0061] The experimental results of image segmentation are shown in Table 1. All experiments in the table are performed on the Synapse dataset. The proposed method outperforms TransUnet, a method that also combines CNN and Transformer, in both DSC and HD. It also outperforms SwinUnet, the latest method based entirely on SwinTransformer.

[0062] Table 1 Experimental comparison of different image segmentation methods

[0063]

Claims

1. Image segmentation method based on CNN and Transformer; its characteristics are: The specific steps include: Step 1: Extract global semantic features of the image; The image is input into the Transformer encoding part, first passing through the patch segmentation layer, then passing through four downsampling layers, where the patch is first fused, and then input into the SwinTransformerBlock to extract global semantic features; Step 2: Extract local detail features of the image through the CNN module; Step 2-1: Convolution and downsampling. The whole CNN stage performs four downsamplings. The final image format is Step 3: Feature reconstruction and merging; After both the CNN module and the Transformer module have extracted features, the feature data of the convolutional network needs to be reconstructed because they have different organization specifications for feature labeling data. The reconstructed data is then merged with the output of the latter in the dimension channel. The merged feature data is passed through a Transformer block, and the output result is reconstructed again to make it conform to the data format required for upsampling. Step 4: Decode the feature information and obtain the segmentation result; Step 4-1: The decoder inputs the merged and reconstructed data into the upsampling layer, gradually restoring the resolution of the feature image. After upsampling, it is merged with the feature images of different levels retained in CNN through jump connections, and then the merged image is convolved. After three consecutive upsampling and three jump connections, the image size is restored to Finally, the image is restored to H×W×C through a patch expansion layer with 4x upsampling; Step 4-2: Segmentation prediction, pass the upsampled feature labels to the linear mapping layer, and output the final pixel-level segmentation prediction result.

2. The image segmentation method based on CNN and Transformer according to claim 1, characterized in that: The patch segmentation layer is specifically: The patch segmentation module segments the input RGB image into non-overlapping patches, each of which is considered a "marker". Its characteristics are set to the concatenation of the original pixel RGB values.

3. The image segmentation method based on CNN and Transformer according to claim 1, characterized in that: After four downsampling layers, specifically: In the first downsampling layer, the patch merging layer is replaced by a linear embedding layer, and the patch is reconstructed into a dimension of The structure is as follows: the patch merging layer or the linear embedding layer is followed by the SwinTransformerBlock to extract features.

4. The image segmentation method based on CNN and Transformer according to claim 3, characterized in that: The patch merging layer is used to reduce the number of markers and strengthen the global features. The H and W of the patch become the original The feature dimension becomes twice as large as before; The patch merging layer passes the merged patches to SwinTransformerBlock for feature extraction.

5. The image segmentation method based on CNN and Transformer according to claim 3, characterized in that: The construction method of SwinTransformerBlock is to replace the standard multi-head self-attention module in the original Transformer with a module based on shifted windows. Since it is a layered structure, the Transformers in each block must exist in pairs. In the window multi-head self-attention layer of the first Transformer, the patch is divided into four blocks and its global features are learned. The extracted features are then input into the second Transformer. In the second shifted window multi-head self-attention layer, the patches are re-divided and combined according to specific rules to extract patch features. The computational complexity of the standard multi-head self-attention and the shifted window-based self-attention is as follows: Ω(MSA)=4hwC 2 +2(hw) 2 C (1) Ω(W-MSA)=4hwC 2 +2M 2 hwC (2) Where h is the number of blocks divided in the height direction of the image, and w is the number of blocks divided in the width direction of the image. The computational cost in formula (1) is proportional to the square of the number of patches hw, while the computational cost in formula (2) is linearly related to hw.

6. The image segmentation method based on CNN and Transformer according to claim 1, characterized in that: The CNN stage performs four downsampling operations, specifically: The same image is passed to the convolutional network part at the same time. In order to keep pace with the patch size of the Transformer part, the convolutional network part performs 4 downsampling. Before the first downsampling, the image information is extracted by convolution operation, and then downsampled through the first pooling layer to reduce the amount of calculation and expand the receptive field. After pooling, the H and W of the image become and The dimension becomes C. In the next three downsampling, each downsampling will change the H and W of the image to the original The dimension is doubled, and after each downsampling, a copy of the extracted image feature markers is retained for use in the jump connection of the upsampling part during the decoding process.

7. The image segmentation method based on CNN and Transformer according to claim 1, characterized in that: The upsampling described in step 4-1 is specifically as follows: It consists of multiple upsampling steps to decode the hidden features and output the final segmentation mask; after encoding the image through a hybrid encoder, it reaches the full resolution through multiple upsampling blocks, where each block consists of a 2× upsampling operator, a 3×3 convolutional layer and a ReLU layer.

8. The image segmentation method based on CNN and Transformer according to claim 1, characterized in that: The skip connection is specifically as follows: the cascade upsampling and the hybrid encoder form a U-shaped architecture. Direct upsampling of low-resolution patches to high-resolution patches will lose feature information, while the feature map retained by CNN preserves the low-level features of the image. Feature aggregation at different resolution levels is achieved through skip connections. The original image is restored as much as possible by fusing the intermediate features preserved after CNN downsampling with the upsampled features. During the first upsampling, the patch features tend to global semantic information as a whole. Therefore, in the first skip connection, the feature map will pass through an additional Transformer block to further enhance the global characteristics of the low-level feature map retained by CNN.

Images