Self-attention segmentation method for diabetic retinopathy fundus images

Abstract

The application discloses a self-attention segmentation method for diabetic retinopathy fundus images, and comprises the following steps: S1, collecting diabetic retinopathy fundus images, and obtaining a first feature map set after pretreatment; S2, inputting the first feature map set into a swin-transformer network to obtain a third feature map set; and S3, inputting the third feature map set into a prediction layer network to obtain a prediction image, and completing the self-attention segmentation of the diabetic retinopathy fundus images. The application can perform fine segmentation on high-resolution images, and has better effect than traditional U-Net architecture, M-Net, Swin-Unet architecture and the like. Global information can be acquired more. The application uses patch to perform picture cutting and then performs picture prediction, so that the neural network size can be effectively reduced, and the neural network training speed can be accelerated.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of image processing, and particularly relates to a self-attention segmentation method for diabetic retinopathy fundus images. BACKGROUND

[0002] At present, the neural network method has made outstanding achievements in the semantic segmentation, target detection and hierarchical classification of images. In the intelligent segmentation and identification of diabetic retinopathy fundus images, the neural network method also plays a great role. However, most of the current researches still stay in the end-to-end picture hierarchical prediction, and the direct picture hierarchical result cannot be convincing. Therefore, the diabetic retinopathy image classification has changed from the end-to-end direct classification to the feature extraction by the semantic segmentation network, so that the neural network result is more explainable.

[0003] Many scholars and medical experts cooperate to accumulate a large amount of labeled data sets, especially the data sets with lesion segmentation labels, which promotes the rapid development of the diabetic retinopathy neural network diagnosis method. With the progress of equipment, the resolution and clarity of the fundus images are continuously improved, and the amount of accumulated data is also continuously increased. However, due to the limitation of the convolutional neural network, the current method cannot meet the requirement of more fine segmentation, and cannot obtain more rich information in the image. Therefore, we need a better neural network method to learn more rich information in the image and realize more fine fundus image segmentation.

[0004] In the field of image segmentation, the most classic network is U-Net. The U-Net network becomes the most classic segmentation algorithm with its creative encoder, decoder and jump fusion structure. However, the traditional U-Net structure adopts the convolution pooling operation for the encoding and decoding operation, the limited receptive field and the lack of global information lead to that even if there is a large amount of data, some information still cannot be effectively learned. In addition, the size of the neural network will be greatly increased with the increase of the resolution of the input image, and the parameter quantity is too large, so that the neural network is difficult to train to fit, and it is difficult to perform more fine segmentation on the picture according to the high-resolution image.

[0005] M-Net proposes a multi-label deep network (M-Net) based on multi-label deep network (M-Net), which is an end-to-end deep learning system. It contains a multi-scale U-shaped convolutional network with a side output layer to learn a discriminative representation and generate a segmentation probability map. It can be used for diabetic retinopathy segmentation work. M-Net uses multi-scale input to increase multi-scale information, so that more surrounding pixel information can be considered when each pixel is predicted. However, it still cannot obtain global information, and it does not solve the problem of increasing model size with increasing picture size. SUMMARY

[0006] In view of the above problems in the prior art, the self-attention segmentation method for diabetic retinopathy fundus images provided by the present application solves the problem that the traditional neural network cannot perform more fine segmentation on the fundus images.

[0007] In order to achieve the above-mentioned purposes, the technical scheme adopted by the present application is as follows: the self-attention segmentation method for diabetic retinopathy fundus images comprises the following steps:

[0008] S1, collecting diabetic retinopathy fundus images, and obtaining a first feature map set after preprocessing;

[0009] S2, inputting the first feature map set into a swin-transformer network to obtain a third feature map set;

[0010] S3, inputting the third feature map set into a prediction layer network to obtain a predicted image, and completing the self-attention segmentation of the diabetic retinopathy fundus images.

[0011] Further, in S1, the preprocessing method is specifically:

[0012] The diabetic retinopathy fundus images are naturally down-sampled by an average pooling layer, the sampled images are constructed into three diabetic retinopathy fundus images of different scales by an encoder, the three diabetic retinopathy fundus images of different scales are sequentially input into a Patch Partition layer and a Linear Embedding layer, the first feature maps corresponding to the three diabetic retinopathy fundus images of different scales are obtained, and then the first feature map set is obtained.

[0013] Further, in S2, the swin-transformer network comprises first to eighth swin-transformer modules;

[0014] The output end of the first swin-transformer module is connected with the input end of the fourth swin-transformer module, the side output end of the first swin-transformer module is connected with the side input end of the second swin-transformer module, the output end of the second swin-transformer module is connected with the input end of the fifth swin-transformer module, and the side output end of the second swin-transformer module is connected with the side input end of the third swin-transformer module.

[0015] An output end of the third swin-transformer module is connected with an input end of the sixth swin-transformer module, and a side output end of the third swin-transformer module is connected with a side input end of the sixth swin-transformer module through a seventh swin-transformer module and an eighth swin-transformer module in sequence;

[0016] A side output end of the sixth swin-transformer module is connected with a side input end of the fifth swin-transformer module, and a side output end of the fifth swin-transformer module is connected with a side input end of the fourth swin-transformer module.

[0017] Further, the first to sixth swin-transformer modules each include a normalization layer, a multi-head self-attention sub-module, a residual connection layer and a two-layer MLP layer with GELU nonlinearity.

[0018] The multi-head self-attention sub-modules of the first to third swin-transformer modules are window-based multi-head self-attention sub-modules, and the multi-head self-attention sub-modules of the fourth to sixth swin-transformer modules are moving window-based multi-head self-attention sub-modules.

[0019] Further, the S2 is specifically:

[0020] A first feature map with a size of WxHx3, a first feature map with a size of W / 2xH / 2x3 and a first feature map with a size of W / 4xH / 4x3 are respectively input into input ends of the first to third swin-transformer modules, an output end of the first swin-transformer module outputs a second feature map with a size of W / 4xH / 4xC to an input end of the fourth swin-transformer module, an output end of the second swin-transformer module outputs a second feature map with a size of W / 8xH / 8x2C to an input end of the fifth swin-transformer module, an output end of the third swin-transformer module outputs a second feature map with a size of W / 16xH / 16x4C to an input end of the sixth swin-transformer module, and a side output end of the third swin-transformer module outputs a second feature map with a size of W / 32xH / 32x8C to a side input end of the sixth swin-transformer module through the seventh to eighth swin-transformer modules in sequence.

[0021] The fourth to sixth swin-transformer modules output a third feature map set according to the input second feature map, wherein C is a channel number, W is an image width, and H is an image height.

[0022] Further, in the S3, the prediction layer network comprises, in sequence, a patch merging module, a linear embedding layer and a self-attention module.

[0023] Further, the S3 is specifically:

[0024] The third feature map is sequentially input into the patch merging module and the linear embedding layer for feature extraction, and the extracted features are input into the self-attention module for fusion to obtain a prediction image, thereby completing the self-attention segmentation of the diabetic retinopathy fundus image.

[0025] The present application has the following advantages:

[0026] (1) The self-attention segmentation method for the diabetic retinopathy fundus image provided by the present application can perform fine segmentation on high-resolution images, and has better effect than the traditional U-Net architecture, M-Net, Swin-Unet architecture and the like. More global information can be obtained.

[0027] (2) The present application uses patch to cut the picture and then performs picture prediction, which can effectively reduce the size of the neural network and speed up the training speed of the neural network. Experiments on the exudation lesions of the IDRiD data set show that, compared with the traditional U-Net architecture, M-Net, Swin-Unet architecture, better results are obtained.

[0028] (3) The present application introduces multi-scale input and multi-scale input and output into the traditional U-Net network to solve the shortcoming that the traditional neural network structure cannot obtain global information, thereby realizing more fine segmentation of the diabetic retinopathy fundus image, and the multi-scale input and output in the structure can increase the width of the network at each layer, avoiding the difficulty of feature extraction caused by too deep network. BRIEF DESCRIPTION OF DRAWINGS

[0029] Figure 1 The present application is a flowchart of the method.

[0030] Figure 2 The present application is a structure diagram of the swin-transformer network and the prediction layer network as a whole.

[0031] Figure 3 The present application is a structure diagram of the swin-transformer network. DETAILED DESCRIPTION

[0032] The specific embodiments of the present application are described below to enable those skilled in the art to understand the present application, but it should be clear that the present application is not limited to the scope of the specific embodiments, and for those skilled in the art, it is obvious that various changes are within the spirit and scope of the present application defined and determined by the appended claims, and all the inventions utilizing the concept of the present application are within the scope of protection.

[0033] As shown in the drawings, Figure 1 In one embodiment of the present application, the self-attention segmentation method for diabetic retinopathy fundus images comprises the following steps:

[0034] S1, collect diabetic retinopathy fundus images, and obtain a first feature map set after preprocessing;

[0035] S2, input the first feature map set into the swin-transformer network to obtain a third feature map set;

[0036] S3, input the third feature map set into the prediction layer network to obtain a predicted image, and complete the self-attention segmentation of the diabetic retinopathy fundus image.

[0037] The overall network diagram of the swin-transformer network and the prediction layer network of the present application is shown in the drawings. Figure 2

[0038] In the S1, the method of preprocessing is specifically:

[0039] The average pooling layer is used for natural down-sampling of the diabetic retinopathy fundus image, the encoder is used for constructing the sampled image into three diabetic retinopathy fundus images of different scales, the three diabetic retinopathy fundus images of different scales are sequentially input into the Patch Partition layer and the Linear Embedding layer, the first feature map corresponding to the three diabetic retinopathy fundus images of different scales is obtained, and then the first feature map set is obtained.

[0040] The average pooling layer can naturally down-sample the input image and construct multi-scale input on the path of the encoder and the decoder.

[0041] In the present embodiment, as shown in the drawings, Figure 3 In the S2, the swin-transformer network comprises first to eighth swin-transformer modules.

[0042] ​The output end of the first swin-transformer module is connected with the input end of the fourth swin-transformer module, the side output end of the first swin-transformer module is connected with the side input end of the second swin-transformer module, the output end of the second swin-transformer module is connected with the input end of the fifth swin-transformer module, and the side output end of the second swin-transformer module is connected with the side input end of the third swin-transformer module.

[0043] The output end of the third swin-transformer module is connected with the input end of the sixth swin-transformer module, and the side output end of the third swin-transformer module is connected with the side input end of the sixth swin-transformer module through the seventh swin-transformer module and the eighth swin-transformer module in sequence.

[0044] The side output end of the sixth swin-transformer module is connected with the side input end of the fifth swin-transformer module, and the side output end of the fifth swin-transformer module is connected with the side input end of the fourth swin-transformer module.

[0045] In the embodiment, the swin-transformer network is a multi-scale U-shaped convolutional network with a side output layer, the width of the swin-transformer network model can be increased, feature loss can be avoided, and a discriminative representation can be learned and a segmentation probability map can be generated.

[0046] The first to sixth swin-transformer modules each comprise a normalization layer, a multi-head self-attention submodule, a residual connection layer and two-layer MLP layers with GELU nonlinearity.

[0047] The multi-head self-attention submodules of the first to third swin-transformer modules are window-based multi-head self-attention submodules, and the multi-head self-attention submodules of the fourth to sixth swin-transformer modules are moving window-based multi-head self-attention submodules.

[0048] Because the resolution of the image is large, the amount of calculation required by the segmentation model is too large, causing network training difficulty, therefore, when training and testing the neural network, the whole picture is not put into the network, but the patch is cut first, then the cut picture is put into the network for training, and the patch is spliced as a prediction image when testing.

[0049] The S2 is specifically:

[0050] The first feature map with the size of WxHx3, the first feature map with the size of W / 2xH / 2x3 and the first feature map with the size of W / 4xH / 4x3 are respectively input into the input end of the first to third swin-transformer modules, the output end of the first swin-transformer module outputs the second feature map with the size of W / 4xH / 4xC to the input end of the fourth swin-transformer module, the output end of the second swin-transformer module outputs the second feature map with the size of W / 8xH / 8x2C to the input end of the fifth swin-transformer module, the output end of the third swin-transformer module outputs the second feature map with the size of W / 16xH / 16x4C to the input end of the sixth swin-transformer module, and the side output end of the third swin-transformer module outputs the second feature map with the size of W / 32xH / 32x8C to the side input end of the sixth swin-transformer module in turn through the seventh to eighth swin-transformer modules;

[0051] The fourth to sixth swin-transformer modules output third feature maps according to the input second feature maps, to obtain a third feature map set, wherein C is a channel number, W is an image width, and H is an image height.

[0052] In the S3, the prediction layer network comprises a patchmerging module, a linearembedding layer and a self-attention module connected in turn.

[0053] The S3 is specifically:

[0054] The third feature maps are input into the patchmerging module and the linearembedding layer in turn for feature extraction, the extracted features are input into the self-attention module for fusion, to obtain a prediction image, and the self-attention segmentation of the diabetic retinopathy fundus image is completed.

[0055] In the embodiment, the generated prediction image is pre-trained using a large-scale data set FGADR, which can effectively solve the problem of lack of prior information of pure Transformeru.

[0056] The application has the beneficial effects that the self-attention segmentation method for diabetic retinopathy fundus images provided by the application can finely segment high-resolution images, and has better effects than traditional U-Net architecture, M-Net, Swin-Unet architecture and the like. Global information can be obtained more.

[0057] The application uses patch to cut pictures and then performs picture prediction, which can effectively reduce the size of the neural network and accelerate the training speed of the neural network. Experiments on exudative lesions are performed on the IDRiD data set, and better effects are obtained compared with traditional U-Net architecture, M-Net, Swin-Unet architecture.

[0058] The application introduces multi-scale input and multi-scale input and output into the traditional U-Net network to solve the shortcoming that the traditional neural network structure cannot obtain global information, so that more fine segmentation of diabetic retinopathy fundus images is realized, and the width of the network can be increased at each layer in the structure of the multi-scale input and output, so that the difficulty of feature extraction caused by too deep network is avoided.

[0059] In the description of the application, it should be understood that the orientations or positional relationships indicated by the terms "center", "thickness", "upper", "lower", "horizontal", "top", "bottom", "inner", "outer", "radial" and the like are based on the orientations or positional relationships shown in the drawings, and are only for the convenience of describing the application and simplifying the description, and therefore cannot be understood as indicating or implying that the devices or elements referred to must have a particular orientation, be constructed and operated in a particular orientation, and therefore cannot be understood as limiting the application. In addition, the terms "first", "second", "third" are only for the purpose of description, and cannot be understood as indicating or implying relative importance or implicitly indicating the number of technical features. Therefore, the features limited by "first", "second", "third" can explicitly or implicitly include one or more features.

Claims

1. A self-attention segmentation method for diabetic retinopathy fundus images, characterized in that, The method comprises the following steps: S1, collecting diabetic retinopathy fundus images, and obtaining a first feature map set after pretreatment; S2, inputting the first feature map set into a swin-transformer network to obtain a third feature map set; S3, inputting the third feature map set into a prediction layer network to obtain a prediction image, and completing the self-attention segmentation of the diabetic retinopathy fundus image; In S1, the method of pretreatment is specifically: The diabetic retinopathy fundus image is naturally down-sampled through an average pooling layer, the sampled image is constructed into three diabetic retinopathy fundus images of different scales through an encoder, and each diabetic retinopathy fundus image of different scales is subjected to color change and slicing processing, the three diabetic retinopathy fundus images of different scales are sequentially subjected to Patch Partition layer and Linear Embedding layer to obtain the first feature map corresponding to each diabetic retinopathy fundus image of different scales, and then the first feature map set is obtained; In S2, the swin-transformer network comprises a first to eighth swin-transformer module; The output end of the first swin-transformer module is connected with the input end of the fourth swin-transformer module, the side output end of the first swin-transformer module is connected with the side input end of the second swin-transformer module, the output end of the second swin-transformer module is connected with the input end of the fifth swin-transformer module, the side output end of the second swin-transformer module is connected with the side input end of the third swin-transformer module; The output end of the third swin-transformer module is connected with the input end of the sixth swin-transformer module, the side output end of the third swin-transformer module is sequentially connected with the side input end of the sixth swin-transformer module through the seventh swin-transformer module and the eighth swin-transformer module; The side output end of the sixth swin-transformer module is connected with the side input end of the fifth swin-transformer module, and the side output end of the fifth swin-transformer module is connected with the side input end of the fourth swin-transformer module; S2 is specifically: The first feature map with a size of WxHx3, the first feature map with a size of W / 2xH / 2x3 and the first feature map with a size of W / 4xH / 4x3 are input into input ends of the first to third swin-transformer modules respectively, an output end of the first swin-transformer module outputs a second feature map with a size of W / 4xH / 4xC to an input end of the fourth swin-transformer module, an output end of the second swin-transformer module outputs a second feature map with a size of W / 8xH / 8x2C to an input end of the fifth swin-transformer module, an output end of the third swin-transformer module outputs a second feature map with a size of W / 16xH / 16x4C to an input end of the sixth swin-transformer module, and a side output end of the third swin-transformer module outputs a second feature map with a size of W / 32xH / 32x8C to side input ends of the seventh to eighth swin-transformer modules in turn; The fourth to sixth swin-transformer modules output third feature maps according to the input second feature maps, to obtain a third feature map set, wherein C is a channel number, W is an image width, and H is an image height; In the S3, the prediction layer network comprises a patch merging module, a linear embedding layer and a self-attention module connected in turn; and the S3 is specifically: The third feature maps are input into the patch merging module and the linear embedding layer in turn for feature extraction, and the extracted features are input into the self-attention module for fusion to obtain a prediction image, thereby completing self-attention segmentation of the diabetic retinopathy fundus image. 2.The self-attention segmentation method for diabetic retinopathy fundus images of claim 1, wherein, The first to sixth swin-transformer modules each comprise a normalization layer, a multi-head self-attention submodule, a residual connection layer and two-layer MLP layers with GELU nonlinearity. The multi-head self-attention submodules of the first to third swin-transformer modules are window-based multi-head self-attention submodules, and the multi-head self-attention submodules of the fourth to sixth swin-transformer modules are moving window-based multi-head self-attention submodules.

Images