A segmentation method for pterygium lesion area based on Vision Transformer
By combining the Vision Transformer and convolutional neural network methods, the problem of low segmentation accuracy of pterygium lesion areas in traditional methods was solved, high-precision lesion area segmentation was achieved, and the effect of medical image segmentation was improved.
Patent Information
- Application Number
- CN202310254245.6
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2023-03-11
- Publication Date
- 2025-09-26
- Estimated Expiration
- 2043-03-11
AI Technical Summary
Traditional machine learning methods have difficulty in achieving accurate segmentation of pterygium lesion areas in medical images, and deep convolutional networks are insufficient in extracting edge detail information, resulting in unsatisfactory segmentation results.
A segmentation method based on Vision Transformer is adopted, combined with convolutional neural network and attention mechanism. The model is trained using expert-annotated datasets. The detailed information of the lesion area is retained through a multi-layer multi-head attention mechanism and pyramid pooling module, and the cross entropy and Dice Loss functions are used to optimize the segmentation effect.
It achieved accurate segmentation of the pterygium lesion area, improved segmentation accuracy and edge segmentation effect, and reached an average intersection-over-union ratio of 87.43% and an average pixel accuracy of 92.57%.
Smart Images
Figure CN116310335B_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the field of medical image processing and application technology, and specifically relates to a pterygium lesion area segmentation method based on deep learning. Background Art
[0002] The global prevalence of pterygium is approximately 12%. Before 2015, researchers primarily used traditional machine learning to segment target objects. Traditional segmentation methods include thresholding, region growing, and edge detection. However, the accuracy and efficiency of traditional machine learning in medical imaging applications are insufficient for practical application.
[0003] In recent years, numerous studies have utilized deep learning convolution techniques to achieve disease classification and diagnosis, achieving approximately 95% accuracy. However, classification alone cannot accurately localize the lesion area for surgical treatment of pterygium. Currently, there is extensive research both domestically and internationally on the application of convolution techniques in medical segmentation, with segmentation accuracy exceeding that of traditional machine learning. Although deep convolution can extract target information, much edge detail is lost during the convolution process, resulting in suboptimal edge segmentation. Summary of the Invention
[0004] Purpose of the Invention: This invention addresses the problems of limited pterygium data, low segmentation accuracy, and difficulty in segmenting boundaries. This method proposes a segmentation method for pterygium lesions based on Vision Transformer. This method uses Vision Transformer as the backbone and a convolutional neural network as an auxiliary, integrating an attention mechanism. The model is trained using a dataset of pterygium lesion regions annotated by experts, with the goal of extracting complete information about the pterygium region. Based on the characteristics of the model network structure and the requirements of medical image segmentation tasks, a new segmentation method is proposed to achieve accurate pterygium segmentation.
[0005] Technical solution:
[0006] A method for segmenting pterygium lesion areas based on Vision Transformer, comprising an acquisition module, a semantic segmentation network module, and a training module, is provided. The method is characterized by the following steps:
[0007] (1) The diseased anterior segment images constitute a set of pterygium segmentation datasets, the pterygium segmentation datasets are used as raw data samples, and the data acquisition module performs a preprocessing operation on the images in the raw data samples to ensure that the length and height of the images are the same, thereby forming a set of training set images;
[0008] (2) Segmenting the training set images using the semantic segmentation network module, wherein the semantic segmentation network module includes a Vision Transformer network and a convolutional network; the Vision Transformer network processes the training set images using an image block method, obtains an image block association relationship by superimposing multiple layers of a multi-head attention mechanism, and then obtains an image attention map; the convolutional network obtains an image feature map through a multi-layer convolution operation; the image attention map and the image feature map are obtained by matrix addition operation to obtain the image feature map, and a pterygium segmentation map is obtained by a pyramid pooling method;
[0009] (3) using a training module to train the segmentation model, inputting the pterygium segmentation dataset into the semantic segmentation network module for training, adjusting the model parameters by setting the learning rate, loss function method, and learning iteration cycle during training, and finally forming a pterygium lesion area segmentation model;
[0010] The pre-processing operation is:
[0011] The method of the present invention requires an input image size of M×N×3, where M and N are positive integers, and an original image size of H×W, where H and W are positive integers. The image is first scaled to M×((N / H)×W), and then gray edges are added to both sides of the shorter side to convert the image size to M×N.
[0012] The image segmentation method is:
[0013] First, the M×N×3 image is upsampled to M'×N'×3 through an upsampling operation, and the M'×N'×3 image is input to the Vision Transformer. The input image is divided into (M' / Patch)×(N' / Patch) Patch×Patch image blocks, and a trainable position information parameter of size 1×((M' / Patch)×(N' / Patch))×(3×Patch×Patch) is added to each image block.
[0014] The multi-head attention mechanism is:
[0015] The image blocks are input into the multi-head attention mechanism, and the relationship between each image block is calculated through matrix operation to generate a new image feature of size ((M' / Patch)×(N' / Patch))×(3×Patch×Patch). The multi-head attention mechanism is cycled 12-16 times;
[0016] The image features generated by the multi-head attention mechanism are transformed to obtain the image attention map, and a convolution module with a convolution kernel size of 3×3 is connected to obtain an image attention map of size 2048×30×30;
[0017] The convolutional network is:
[0018] The parameters of the ResNet50 model obtained after pre-training on the public dataset ImageNet are used as the initialization parameters of the convolutional network, and the image feature map is extracted through four layers of convolutional modules with different sizes and structures;
[0019] The pyramid pooling method is as follows:
[0020] The image attention map and the image feature map are added together by a matrix operation to obtain a new image feature map, which is input into a pyramid pooling module. The image feature map is first converted to 1 / 4 of its dimension by a convolution operation. Then, four pooling blocks of different sizes are used to perform pooling operations to obtain image feature maps a, b, c, and d respectively. Finally, the image feature maps a, b, c, and d are upsampled and stacked with the new image feature map to obtain image feature map (e).
[0021] The image feature map a, image feature map b, image feature map c, and image feature map d are input into the stage upsampling module, and an image feature map is obtained through upsampling and feature fusion operations, and stacked with the image feature map (e) to obtain a new image feature map (f). Finally, a semantic segmentation image is obtained by performing a convolution operation on the image feature map (f).
[0022] Its characteristics are:
[0023] The loss function includes a cross entropy loss function and a Dice Loss loss function. The cross entropy loss function and the Dice Loss loss function are fused as the loss function of the semantic segmentation network model to minimize the following objective function:
[0024] Loss=Cross Entropy Loss+Dice Loss
[0025] Among them, Cross Entropy Loss represents the cross entropy loss function, and Dice Loss represents the Dice Loss loss function.
[0026] Its characteristics are:
[0027] The learning iteration cycle is 80 cycles. To speed up the training and prevent the weights from being destroyed in the early stage of training, frozen training is used in cycles 0-40 and normal training is used in cycles 40-80. The learning rate is 1e-5.
[0028] Beneficial effects:
[0029] The training is conducted using a pterygium segmentation dataset annotated by experts to ensure the authority of the training.
[0030] ResNet50 is used as the feature extraction network, and the transfer learning method is used for pre-training on the public dataset ImageNet. The deep convolutional network ensures that the model can extract sufficiently complete features of the lesion area.
[0031] The multi-head attention mechanism in Vision Transformer can greatly preserve the detailed information of the lesion area contour by connecting the internal relationships of the image.
[0032] A stage upsampling module is added to the pyramid pooling module. While extracting contextual information using pooling blocks of different sizes, the stage upsampling module retains target detail information through special graph fusion.
[0033] In the network model, cross entropy loss and Dice Loss are fused as the loss function to improve the segmentation effect of pterygium. BRIEF DESCRIPTION OF THE DRAWINGS
[0034] Figure 1 Schematic diagram of the semantic segmentation network structure
[0035] Figure 2 Schematic diagram of the pyramid pooling module structure
[0036] Figure 3 Schematic diagram of the Vision Transformer structure
[0037] Figure 4 Schematic diagram of the stage upsampling structure
[0038] Figure 5 Comparison chart before and after data preprocessing
[0039] Figure 6 Schematic diagram of segmentation results DETAILED DESCRIPTION
[0040] Example: The Vision Transformer-based pterygium lesion region segmentation method provided by the present invention is used to segment the pterygium lesion region. The specific operation is performed as follows:
[0041] 1. Pterygium Dataset
[0042] The pterygium segmentation dataset contains 517 anterior segment images of patients with pterygium (including pterygium symptoms of varying degrees of disease), including 367 training and validation data and 150 test data. The lesion area of each pterygium is personally annotated by an ophthalmologist.
[0043] 2. Data Preprocessing
[0044] The method of the present invention requires an input image size of 473×473×3, and the original image size is H×W (H>W). The image is first scaled to 473×((473 / H)×W)×3, and then the shorter sides are evenly padded with gray edges to convert its size to 473×473×3.
[0045] 3. Training semantic segmentation network
[0046] During network training, data is fed into the network in batches. First, a 473×473×3 image is upsampled to 480×480×3 through an upsampling operation. This 480×480×3 image is then fed into the Vision Transformer. After partitioning the input image into 30×30×3 blocks using convolution with a kernel_size of 16×16 and a stride of 16, the features are flattened to obtain a 900×768×1 sequence of image features. Trainable position information parameters of 900×768×1 are added to each block as the sequence features. These sequence features are fed into a multi-head attention mechanism, which transforms the sequence into q, k, and v matrices. Matrix operations are used to calculate the relationships between each block and generate new 900×768 image features. The multi-head attention mechanism is then repeated 12 times. The generated image features are normalized and linearly transformed to obtain a 30×30×768 image attention map. The image attention map is convolved through a convolution module with a convolution kernel size of 3×3 to obtain an image attention map of size 30×30×2048 (A);
[0047] Secondly, the convolutional network uses ResNet50, and adopts the transfer learning method. The parameters of the ResNet50 model obtained by training on ImageNet are used as the initialization parameters of the convolutional network. The image of size 473×473×3 is input into the ResNet50 network. After passing through four different convolutional modules, each convolutional module contains three different convolution operations, which are cycled 3 times, 4 times, 6 times, and 3 times respectively, the image feature map of size 30×30×2048 is obtained (B).
[0048] The image attention map (A) and the image feature map (B) are element-wise added to obtain a new image feature map (C), which is then input into the pyramid pooling module. A convolution operation with a kernel_size of 1×1 is first used to convert the image feature map (C) dimension to 1 / 4 of the input dimension. Pooling operations are then performed using 1×1, 2×2, 3×3, and 6×6 pooling blocks to obtain image feature maps a, b, c, and d, respectively. Finally, the image feature maps a, b, c, and d are upsampled to 30×30×512 and stacked with the image feature map (C) in a channel-by-channel stacking manner to obtain the image feature map (D).
[0049] At the same time, the image feature maps a, b, c, and d are input into the stage upsampling module. Through the upsampling operation, the image feature maps a, b, c, and d are fused to obtain an image feature map of size 30×30×512, and the image feature map (D) is stacked with the channel to obtain a new feature map (e1).
[0050] Finally, the image feature map of the third layer of ResNet50 is repeated in the pyramid pooling module and the stage upsampling module to obtain the image feature map (e2), and the feature fusion operation is performed on it and the new feature map (e1) to obtain the pterygium semantic segmentation map;
[0051] The pterygium semantic segmentation map and the real pterygium semantic segmentation map are subjected to loss function calculation, and the network parameters are updated. The present invention combines the cross entropy loss function and the Dice Loss loss function as the loss function of the semantic segmentation network model, minimizing the following objective function:
[0052] Loss=Cross Entropy Loss+Dice Loss
[0053] Among them, Cross Entropy Loss represents the cross entropy loss function, and Dice Loss represents the Dice Loss loss function.
[0054] Let y = y truth , y′=y pred , define the following Cross Entropy Loss objective function:
[0055] Cross Entropy Loss=-y·log(y′)-(1-y)·log(1-y)
[0056] The more accurate the pixel classification is, the smaller the loss is.
[0057] Let A and B represent the predicted and true contour region point sets respectively, and define the following Dice Loss objective function:
[0058]
[0059] The greater the overlap ratio between the predicted lesion area and the actual lesion area, the smaller the loss.
[0060] 4. Analysis of processing results
[0061] This method uses the following two performance metrics to quantify the processing results: single-class intersection over union (IoU), mean intersection over union (MIoU) and single-class pixel accuracy (PA), mean pixel accuracy (MPA). The calculation formula is as follows:
[0062]
[0063]
[0064] p i represents the segmented area, g i Represents the true region. The intersection over union (IOU) is obtained by dividing the intersection of the true value and the predicted value by their union; the mean intersection over union (MIOU) is calculated by calculating the IOU of each category (including the background category) and averaging it across all categories.
[0065]
[0066]
[0067] p ii Indicates the number of pixels predicted correctly, p ij Represents the number of pixels predicted as class j for class i. Pixel accuracy (PA) represents the proportion of correctly labeled pixels to the total number of pixels; average pixel accuracy is calculated by calculating the proportion of correctly classified pixels for each class and then averaging it across all classes.
[0068] The results of this method on the pterygium test set were MIOU: 87.43%, MPA: 92.57%, IOU: 79.44%, and PA: 87.16%. Extensive use demonstrates that the Vision Transformer-based segmentation method for pterygium lesions proposed in this paper has high segmentation performance. This is of great significance in the medical field.
[0069] As described above, although the present invention has been shown and described with reference to certain preferred embodiments, it is not to be construed as limiting the invention itself. Various changes in form and details may be made thereto without departing from the spirit and scope of the invention as defined in the appended claims.
Claims
1. A method for segmenting pterygium lesion areas based on Vision Transformer, comprising an acquisition module, a semantic segmentation network module, and a training module. The method uses the data acquisition module, the semantic segmentation network module, and the training module to segment images of the anterior segment of the diseased eye, and is characterized by: The specific steps are as follows: (1) The diseased anterior segment images constitute a set of pterygium segmentation datasets, the pterygium segmentation datasets are used as raw data samples, and the data acquisition module performs a preprocessing operation on the images in the raw data samples to ensure that the length and height of the images are the same, thereby forming a set of training set images; (2) Segmenting the training set images using the semantic segmentation network module, wherein the semantic segmentation network module includes a Vision Transformer network and a convolutional network; the Vision Transformer network processes the training set images using an image block method, obtains an image block association relationship by superimposing multiple layers of a multi-head attention mechanism, and then obtains an image attention map; the convolutional network obtains an image feature map through a multi-layer convolution operation; the image attention map and the image feature map are obtained by matrix addition operation to obtain the image feature map, and a pterygium segmentation map is obtained by a pyramid pooling method; (3) using a training module to train the segmentation model, inputting the pterygium segmentation dataset into the semantic segmentation network module for training, adjusting the model parameters by setting the learning rate, loss function method, and learning iteration cycle during training, and finally forming a pterygium lesion area segmentation model; The pre-processing operation is: The required input image size is M×N×3, where M and N are positive integers. The original image size is H×W, where H and W are positive integers. The image is first scaled to M×((N / H)×W), and then the shorter sides are padded with gray edges to convert the size to M×N. The image segmentation method is: First, the M×N×3 image is upsampled to M'×N'×3 through an upsampling operation, and the M'×N'×3 image is input to the Vision Transformer. The input image is divided into (M' / Patch)×(N' / Patch) Patch×Patch image blocks, and a trainable position information parameter of size 1×((M' / Patch)×(N' / Patch))×(3×Patch×Patch) is added to each image block. The multi-head attention mechanism is: The image blocks are input into the multi-head attention mechanism, and the relationship between each image block is calculated through matrix operation to generate a new image feature of size ((M' / Patch)×(N' / Patch))×(3×Patch×Patch). The multi-head attention mechanism is cycled 12-16 times; The image features generated by the multi-head attention mechanism are transformed to obtain the image attention map, and a convolution module with a convolution kernel size of 3×3 is connected to obtain an image attention map of size 2048×30×30; The convolutional network is: The parameters of the ResNet50 model obtained after pre-training on the public dataset ImageNet are used as the initialization parameters of the convolutional network, and the image feature map is extracted through four layers of convolutional modules with different sizes and structures; The pyramid pooling method is as follows: The image attention map and the image feature map are added together by a matrix operation to obtain a new image feature map, which is input into a pyramid pooling module. The image feature map is first converted to 1 / 4 of its dimension by a convolution operation. Then, four pooling blocks of different sizes are used to perform pooling operations to obtain image feature maps a, b, c, and d respectively. Finally, the image feature maps a, b, c, and d are upsampled and stacked with the new image feature map to obtain image feature map (e). The image feature map a, image feature map b, image feature map c, and image feature map d are input into the stage upsampling module, and an image feature map is obtained through upsampling and feature fusion operations, and stacked with the image feature map (e) to obtain a new image feature map (f). Finally, a semantic segmentation image is obtained by performing a convolution operation on the image feature map (f).
2. The segmentation method according to claim 1, wherein: The loss function includes a cross entropy loss function and a Dice Loss loss function. The cross entropy loss function and the Dice Loss loss function are fused as the loss function of the semantic segmentation network model to minimize the following objective function: Loss=Cross Entropy Loss+Dice Loss Among them, Cross Entropy Loss represents the cross entropy loss function, and Dice Loss represents the Dice Loss loss function.
3. The segmentation method according to claim 1, wherein: The learning iteration cycle is 80 cycles. To speed up the training and prevent the weights from being destroyed in the early stage of training, frozen training is used in cycles 0-40 and normal training is used in cycles 40-80. The learning rate is 1e-5.