A Cataract Image Classification Method and System Based on Self-Supervised Pre-training

Abstract

The present invention discloses a cataract image classification method and system based on self-supervised pre-training, which relates to the field of image classification technology. The specific steps are: constructing a mask template based on a dilated anterior segment photographic image; extracting a slit lamp photographic image from a local data set, and performing image and operation on the slit lamp photographic image and the mask template to obtain a mask image; building an image restoration model; using the image restoration model to perform self-supervised pre-training on the mask image to obtain a self-supervised pre-training weight; constructing an image classification model based on the image restoration model, the self-supervised pre-training weight and a classification module; obtaining an image to be identified, and obtaining a classification result based on the image classification model. The present invention can reduce the dependence on the number of samples through self-supervised pre-training, and can also improve the accuracy of classification and recognition.

Description

Technical Field

[0001] The present invention relates to the technical field of image classification, and more particularly to accurate identification of cataract image classification. Background Art

[0002] Cataract is the main cause of blindness. Studies have shown that the number of patients who become blind due to cataracts will increase by 1 million each year, and the number of cataract patients who need surgical treatment will increase by 4 million each year. This huge patient population has placed a heavy burden on the medical system. The grading of the severity of the cataract nucleus is crucial for determining whether surgery is necessary and for designing surgical plans. Ophthalmologists usually judge the severity of the cataract nucleus by observing the slit lamp photography images of the patient after dilation of the pupil. However, the disadvantage of this diagnostic method is that it relies on the subjective vision of the ophthalmologist and there is a problem of subjective diagnosis. How to more accurately and objectively grade the severity of the cataract nucleus through slit lamp photography images is a technical problem faced by this field.

[0003] In order to objectively and accurately evaluate the severity of cataracts, based on the AREDS (age-related eye diseases study) No. 4 document, standard 1 is defined as normal, standards 2 and 3 are judged as mild, standards 4 and 5 are judged as moderate, and standards 6 and 7 are defined as severe cataracts, with a total of four categories. Using deep learning technology to analyze dilated anterior segment photography images can achieve intelligent grading of cataract nuclear severity, make up for the shortcomings of traditional detection methods, and achieve more efficient intelligent grading. However, there are still certain challenges in directly using deep learning models to process medical images.

[0004] In the existing technology, Qian et al. conducted a three-classification study on cataracts using SqueezeNet based on the transfer learning method. By manually cropping the lesion area and balancing the data set distribution, the diagnostic accuracy of the model reached 96%.

[0005] Zhou et al. used prior knowledge to combine improved Haar features and visible structure features, and their accuracy in grading the severity of cataract nuclei was 89.23%.

[0006] Xu et al. proposed GLA-Net. With a large amount of labeled data to train the model, the model enhanced the classification performance of cataracts by integrating global and local attention, and the diagnostic accuracy of the model was 83.47%.

[0007] However, the above models are simple in classification, need to combine prior information to build the model, and rely heavily on a large amount of labeled data, which cannot meet the requirements for accurate identification of cataract image classification. Therefore, how to accurately identify the classification of cataract images is an urgent problem that technicians in this field need to solve. Summary of the invention

[0008] In view of this, the present invention provides a cataract image classification method and system based on self-supervised pre-training, which overcomes the above-mentioned defects.

[0009] In order to achieve the above object, the present invention provides the following technical solutions:

[0010] A cataract image classification method based on self-supervised pre-training, the specific steps are:

[0011] A mask template is constructed based on the dilated anterior segment camera image; a slit lamp camera image is extracted from a local data set, and an image AND operation is performed between the slit lamp camera image and the mask template to obtain a mask image;

[0012] Build an image restoration model;

[0013] Use the image restoration model to perform self-supervised pre-training on the mask image to obtain self-supervised pre-training weights;

[0014] Build an image classification model based on the image restoration model, self-supervised pre-trained weights, and classification module;

[0015] Obtain the image to be identified and obtain the classification result based on the image classification model.

[0016] Optionally, the mask template is constructed by:

[0017] Acquire a mydriatic anterior segment photographic image, and perform scaling processing on the mydriatic anterior segment photographic image;

[0018] The zoomed mydriatic anterior segment photographic image is used to generate a mask template according to a set mask ratio.

[0019] Optionally, the image restoration model includes an encoder and a decoder; the encoder is divided into five stage parts, namely, a first stage part stage1, a second stage part stage2, a third stage part stage3, a fourth stage part stage4 and a fifth stage part stage5; the structures within the second stage part stage2, the third stage part stage3, the fourth stage part stage4 and the fifth stage part stage5 are repeatedly stacked 3, 4, 6 and 3 times respectively; the decoder includes a feature fusion module, a first upsampling block UpBlock1, a second upsampling block UpBlock2, a third upsampling block UpBlock3, a fourth upsampling block UpBlock4 and a fifth upsampling block UpBlock5.

[0020] Optionally, a convolution layer is provided in the second stage part stage2, the third stage part stage3, the fourth stage part stage4 and the fifth stage part stage5, and the convolution layer includes a convolution operation, a batch normalization operation and an activation function SiLU.

[0021] Optionally, the expression of the activation function SiLU is:

[0022]

[0023] Where x represents the batch normalized tensor.

[0024] Optionally, the structures of each upsampling block adopt the same structure, including a first branch, a second branch, a splicing module and a fusion module; the first branch adopts a Transformer block and a convolutional layer for upsampling; the second branch adopts a transposed convolution for upsampling; the splicing module splices the feature maps output by the first branch and the second branch in the channel dimension; the fusion module uses two convolutional layers to fuse the feature maps output by the splicing module.

[0025] Optionally, a Transformerblock includes a linear layer and a self-attention computation.

[0026] Optionally, the steps for self-attention calculation are:

[0027] Perform self-attention calculation on the output features of the linear layer;

[0028] Linearly map the output results calculated by the self-attention mechanism to obtain the output sequence;

[0029] Reshape the output sequence to obtain the output feature map.

[0030] Optionally, the classification module includes two parallel convolutional layers, Transformerblock, MLP block and AVG block.

[0031] A cataract image classification system based on self-supervised pre-training, comprising:

[0032] A mask template construction module, used for constructing a mask template according to a mydriatic anterior segment photographic image;

[0033] A mask image calculation module is used to extract a slit lamp photography image from a local data set, and perform an image AND operation on the slit lamp photography image and a mask template to obtain a mask image;

[0034] Image restoration model building module, used to build an image restoration model;

[0035] A self-supervised pre-training module is used to perform self-supervised pre-training on the mask image using the image restoration model to obtain self-supervised pre-training weights;

[0036] An image classification model building module is used to build an image classification model based on the image restoration model, self-supervised pre-training weights, and classification module;

[0037] The image recognition module is used to obtain the image to be recognized and obtain the classification result based on the image classification model.

[0038] It can be seen from the above technical solution that compared with the prior art, the present invention has the following beneficial effects:

[0039] The present invention proposes for the first time a cataract grading model based on self-supervised pre-training, which is used to classify the severity of cataracts into four categories. The self-supervised pre-training adopted by the present invention can solve the problem that the deep learning model based on natural image training shows mismatch on medical data, thereby obtaining the initial weight suitable for the experimental data, and then in the case of a small number of samples, the model can overcome the problem of high similarity between samples, and realize efficient grading of the severity of cataract nuclei;

[0040] The present invention designs a hybrid module that combines the advantages of CNN and Transformer, which can effectively capture the relationship between different positions in the feature map and improve the attention to the anterior segment of the eye of the dilated pupil, such as the pupil and iris;

[0041] The present invention constructs a decoder of CNN and Transformer hybrid upsampling, which enables the model to fully explore the potential structure and feature representation of the data without subjective adjustment of the data set, such as manual adjustment of the distribution of the data set and cropping of the lesion area. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are only embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on the provided drawings without paying creative work.

[0043] Figure 1 It is a schematic flow chart of the method of the present invention;

[0044] Figure 2 In the embodiment disclosed in the present invention, the mask rate is Schematic diagram of masking with a mask template;

[0045] Figure 3 is a schematic diagram of a mask template according to an embodiment of the present invention;

[0046] Figure 4 is a schematic diagram of the encoder structure of an embodiment disclosed in the present invention;

[0047] Figure 5 It is a schematic diagram of the structure of a decoder according to an embodiment of the present invention;

[0048] Figure 6 It is a structural diagram of a classification module of an embodiment disclosed in the present invention. DETAILED DESCRIPTION

[0049] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0050] The embodiment of the present invention discloses a cataract image classification method based on self-supervised pre-training, such as Figure 1 As shown, the specific steps are:

[0051] Step 1: construct a mask template based on the dilated anterior segment camera image; extract a slit lamp camera image from a local data set, and perform an image AND operation on the slit lamp camera image and the mask template to obtain a mask image;

[0052] Step 2: Build an image restoration model;

[0053] Step 3: Use the image restoration model to perform self-supervised pre-training on the mask image to obtain self-supervised pre-training weights;

[0054] Step 4: Build an image classification model based on the image restoration model, self-supervised pre-training weights, and classification module;

[0055] Step 5: Obtain the image to be identified and obtain the classification result based on the image classification model.

[0056] In step 1, the steps for obtaining the mask image are:

[0057] First, the dilated anterior segment camera image is uniformly scaled to a size set as needed, and then the mask ratio is determined (the principle of determining the mask ratio is that the set size multiplied by the mask ratio is an integer) to generate a mask template, and finally an image AND operation is performed with the slit lamp camera image to obtain a mask image.

[0058] In this embodiment, the anterior segment photographic image of the dilated pupil is first uniformly scaled to 224x224, and then a mask template with an image size of 224x224 and a mask ratio of one seventh to six sevenths is generated; finally, an image AND operation is performed with the slit lamp photographic image to obtain a mask image. In this embodiment, the mask ratio is The steps to obtain the mask image are as follows: Figure 2 shown.

[0059] Among them, the mask template is as follows Figure 3 As shown, the 224x224 template is divided into 49 grids of size 32x32 according to the mask ratio;

[0060] Randomly select n grids from the 49 grids, and set the pixel values ​​corresponding to the randomly selected n grids to 1, and the pixel values ​​corresponding to the remaining grids to 0;

[0061] Where n represents the number of grids that set the pixel value to 1 under different mask ratios. The calculation formula is as follows:

[0062] n = mask rate * 49;

[0063] In this embodiment, n is 7.

[0064] Among them, the local data set in this embodiment includes a total of 792 slit lamp photography images for model training and testing, of which the number of normal, mild cataract, moderate cataract and severe cataract slit lamp photography images are 125, 218, 320 and 129 respectively;

[0065] In step 2, the image restoration model consists of two parts: an encoder and a decoder; the structure of the encoder is as follows Figure 4 As shown in the figure, it includes five parts: stage1, stage2, stage3, stage4 and stage5, where k7, k3, k1 represent the convolution kernel sizes of 7, 3 and 1 respectively, s2 and s1 represent the step sizes of 2 and 1 respectively, p0, p1 and p3 represent no pixel padding, pixel padding of 1 and pixel padding of 3 respectively, c64, c128, c256, c512 and c2048 represent the number of output channels of 64, 128, 256, 512 and 2048 respectively;

[0066] The structures in stage2 to stage5 are repeatedly stacked 3, 4, 6, and 3 times respectively. Convolutional layer 1, convolutional layer 2, convolutional layer 3, and convolutional layer 4 are only used in the first stacking of each stage. Except for the first stacking, all other stages of each stage are identical mappings.

[0067] The identity mapping represents directly adding the input to the output;

[0068] The structure of each convolutional layer is the same, including convolution operation, batch normalization operation and activation function SiLU;

[0069] The calculation formula corresponding to the activation function SiLU is as follows:

[0070]

[0071] Where x represents the batch normalized tensor;

[0072] The structure of the decoder is as follows Figure 5 As shown, it includes six parts, namely, feature fusion module, UpBlock1, UpBlock2, UpBlock3, UpBlock4 and UpBlock5;

[0073] The feature fusion module consists of two convolutional layers, which perform convolution operations on the output of the encoder;

[0074] Among them, the structure of each UpBlock is the same, consisting of two branches on the left and right. The left branch uses Transformerblock and combines the convolution layer to achieve upsampling, and the right branch uses transposed convolution to achieve upsampling. Then the feature maps output by the two branches are spliced ​​in the channel dimension, and finally two convolution layers are used to perform feature fusion calculation;

[0075] Among them, Transformerblock consists of two parts. The first part uses a linear layer to map the output feature vector x3 to 4 times the input feature vector x2. The second part performs self-attention calculation. The specific operations are as follows:

[0076] First, the feature map x1 to be passed into the Transformerblock is expanded along the feature dimension and converted into a feature vector x2. The size of the feature map x1 is (B, C, H, W), and the size of the feature vector x2 is (B, C, H*W);

[0077] Where B represents the batch size, C represents the number of channels, H and W represent the length and width of the feature map respectively;

[0078] Then the feature vector x2 is mapped to the feature vector x3 using a linear layer, and the size of the feature vector x3 is (B, C, 4*H*W);

[0079] Among them, the mapping formula of the linear layer is as follows:

[0080] x3=w1*x2+b1;

[0081] Among them, w1 and b1 are the multiplier matrix parameters and constant vector parameters of the image restoration network learned and output in the training set;

[0082] The calculation formula for the second part is as follows:

[0083] First, the self-attention mechanism is calculated for the feature vector x3, where the calculation formula for self-attention is as follows:

[0084] Y1=self-attentation(x3)+x3;

[0085] Among them, self-attentation represents the self-attention module, and the calculation formula is as follows:

[0086]

[0087] Among them, Q, K and V are all multiplier matrices, which are the multiplier parameters learned and output by the image restoration network in the training set. Q stands for query, K stands for key, V stands for value, and d k The fixed value is 512.

[0088] Specifically, self-attentation is to perform matrix multiplication of each query vector with the transpose of all key vectors and divide the result by a scaling factor d k Obtain the attention weight, after mapping through the softmax activation function, and finally multiply the attention weight with the corresponding value vector;

[0089] Then, the result calculated by the self-attention mechanism is linearly mapped, and the calculation formula is as follows:

[0090] Y2=(w2*Y1+b2)+Y1;

[0091] Where Y2 represents the final output sequence with a size of (B, C, 4*H*W), w2 and b2 are the multiplier matrix parameters and constant vector parameters of the image restoration network learned output in the training set;

[0092] Finally, Y2 is reshaped and converted into a feature map of (B, C, 2*H, 2*W).

[0093] In step 3, the mask image is self-supervised pre-trained using the image restoration model. The specific steps are:

[0094] The masked image in step 1 is used as the input of the image restoration network, and the mask is normalized, and the unmasked image is used as the output of the image restoration network. After 300 rounds of training, the self-supervised pre-training weights are obtained;

[0095] Among them, the normalized processing is that the mean values ​​of the three RGB channels of the input image are [0.485, 0.456, 0.406], and the corresponding variances are [0.229, 0.224, 0.225].

[0096] In step 4, first migrate the encoder and its weights in step 3, and then build the classification model MeNet in combination with the classification module;

[0097] Among them, the structure of the classification module is as follows Figure 6 As shown in the figure, ⊕ represents concatenating the outputs of the two convolutional layers in the channel dimension. The specific process is as follows:

[0098] First, the feature map output by the encoder is used to halve the number of channels using two convolutional layers. Then, the outputs of the two convolutional layers are concatenated in the channel dimension to restore the original feature map size. The concatenated feature map is then input into the Transformer block for calculation. Finally, the output of the Transformer block is used to linearly transform the feature vector through the MLP. The linear transformation formula is:

[0099] p=w3*Y+b3;

[0100] Where p represents the vector after linear transformation, w3 is the multiplier matrix, which is the multiplier parameter of the classification model in the training set, and b is the constant matrix, which is the constant parameter of the classification model in the training set. p includes 4 elements, corresponding to the probabilities of normal, mild, moderate and severe cataracts, respectively. The largest element among the 4 elements is the category of cataract severity.

[0101] In this embodiment, the verification of the classification model is also included, which uses MeNet to train the dilated anterior segment photography images to determine the severity of the cataract, specifically:

[0102] The local data set is divided into 100 images as the test set. The remaining data is divided into training set and validation set in a ratio of 8:2 and preprocessed, including data enhancement and normalization. Then the cross entropy loss function is used to compare the difference between the classification result and the true result, and the parameter gradient of the loss function with respect to the neural network is calculated. The gradient is used to update the network parameters of the classification model, so that the network classification result is close to the true result of manual annotation, thereby continuously improving the network performance.

[0103] The data enhancement includes color space change, random brightness contrast change, translation and scaling, and random direction rotation of the training set data with probabilities of 0.2, 0.2, 0.2, and 1.0 respectively. The training set and test set data are normalized, and the corresponding means of the three RGB channels are [0.485, 0.456, 0.406] and the corresponding variances are [0.229, 0.224, 0.225].

[0104] The embodiment uses a test set to compare the constructed classification model with the classic model in the field. The classification accuracy of different models for the severity of cataract nucleus is shown in Table 1:

[0105] Table 1 The accuracy of different models in grading the severity of cataract nucleus

[0106] Model Accuracy (%) EfficientNetV2 85.0 NASNet 82.0 RegNet 83.0 ResNeXt 85.0 ResNet50 86.0 ViT 83.0 Swin 85.0 MeNet 91.0

[0107] Table 1 shows the accuracy of cataract nucleus severity classification by the classic model in this field and the MeNet constructed in this embodiment. From the prediction results, based on self-supervised pre-training, the prediction accuracy of MeNet is better than the current model, and compared with the optimal model ResNet50, MeNet's accuracy in cataract nucleus severity classification is improved by 5%.

[0108] A cataract image classification system based on self-supervised pre-training, comprising:

[0109] A mask template construction module, used for constructing a mask template according to a mydriatic anterior segment photographic image;

[0110] A mask image calculation module is used to extract a slit lamp photography image from a local data set, and perform an image AND operation on the slit lamp photography image and a mask template to obtain a mask image;

[0111] Image restoration model building module, used to build an image restoration model;

[0112] A self-supervised pre-training module is used to perform self-supervised pre-training on the mask image using the image restoration model to obtain self-supervised pre-training weights;

[0113] An image classification model building module is used to build an image classification model based on the image restoration model, self-supervised pre-training weights, and classification module;

[0114] The image recognition module is used to obtain the image to be recognized and obtain the classification result based on the image classification model.

[0115] In this specification, each embodiment is described in a progressive manner, and each embodiment focuses on the differences from other embodiments. The same or similar parts between the embodiments can be referred to each other. For the device disclosed in the embodiment, since it corresponds to the method disclosed in the embodiment, the description is relatively simple, and the relevant parts can be referred to the method part.

[0116] The above description of the disclosed embodiments enables one skilled in the art to implement or use the present invention. Various modifications to these embodiments will be apparent to one skilled in the art, and the general principles defined herein may be implemented in other embodiments without departing from the spirit or scope of the present invention. Therefore, the present invention will not be limited to the embodiments shown herein, but rather to the widest scope consistent with the principles and novel features disclosed herein.

Claims

1. A cataract image classification method based on self-supervised pre-training, characterized in that: The specific steps are: A mask template is constructed based on the dilated anterior segment camera image; a slit lamp camera image is extracted from a local data set, and an image AND operation is performed between the slit lamp camera image and the mask template to obtain a mask image; Build an image restoration model; Use the image restoration model to perform self-supervised pre-training on the mask image to obtain self-supervised pre-training weights; Build an image classification model based on the image restoration model, self-supervised pre-trained weights, and classification module; Obtain an image to be identified and obtain a classification result based on an image classification model; The image restoration model includes an encoder and a decoder; the encoder is divided into five stages, namely the first stage stage1, the second stage stage2, the third stage stage3, the fourth stage stage4 and the fifth stage stage5; The structures in the second stage part stage2, the third stage part stage3, the fourth stage part stage4 and the fifth stage part stage5 are repeatedly stacked 3, 4, 6 and 3 times respectively; Decoder It includes a feature fusion module, a first upsampling block UpBlock1, a second upsampling block UpBlock2, a third upsampling block UpBlock3, a fourth upsampling block UpBlock4 and a fifth upsampling block UpBlock5; The structures of each upsampling block adopt the same structure, including the first branch, the second branch, the splicing module and the fusion module; the first branch adopts Transformerblock and convolution layer for upsampling; the second branch adopts transposed convolution for upsampling; the splicing module splices the feature maps output by the first branch and the second branch in the channel dimension; the fusion module uses two convolution layers to fuse the feature maps output by the splicing module.

2. The cataract image classification method based on self-supervised pre-training according to claim 1, characterized in that: The steps to construct the mask template are: Acquire a mydriatic anterior segment photographic image, and perform scaling processing on the mydriatic anterior segment photographic image; The zoomed mydriatic anterior segment photographic image is used to generate a mask template according to a set mask ratio.

3. The cataract image classification method based on self-supervised pre-training according to claim 1, characterized in that: Convolutional layers are provided in the second stage part stage2, the third stage part stage3, the fourth stage part stage4 and the fifth stage part stage5. The convolutional layers include convolution operations, batch normalization operations and activation functions SiLU.

4. The cataract image classification method based on self-supervised pre-training according to claim 3, characterized in that: The expression of the activation function SiLU is: Where x represents the batch normalized tensor.

5. The cataract image classification method based on self-supervised pre-training according to claim 1, characterized in that: Transformerblock includes linear layer and self-attention calculation.

6. The cataract image classification method based on self-supervised pre-training according to claim 5, characterized in that: The steps of self-attention calculation are: Perform self-attention calculation on the output features of the linear layer; Linearly map the output results calculated by the self-attention mechanism to obtain the output sequence; Reshape the output sequence to obtain the output feature map.

7. The cataract image classification method based on self-supervised pre-training according to claim 1, characterized in that: The classification module consists of two parallel convolutional layers, Transformerblock, MLP block and AVG block.

8. A cataract image classification system based on self-supervised pre-training, characterized in that: Executing the cataract image classification method based on self-supervised pre-training as described in claim 1, comprising: A mask template construction module, used for constructing a mask template according to a mydriatic anterior segment photographic image; A mask image calculation module is used to extract a slit lamp photography image from a local data set, and perform an image AND operation on the slit lamp photography image and a mask template to obtain a mask image; Image restoration model building module, used to build an image restoration model; A self-supervised pre-training module is used to perform self-supervised pre-training on the mask image using the image restoration model to obtain self-supervised pre-training weights; An image classification model building module is used to build an image classification model based on the image restoration model, self-supervised pre-trained weights, and classification module; The image recognition module is used to obtain the image to be recognized and obtain the classification result based on the image classification model.

Images