Multi-label Classification Method for Anterior Segment Diseases Based on Dual-branch Cross-attention Pooling

Abstract

The present invention discloses a multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling. The present invention includes the following steps: (1) performing data preprocessing on slit lamp images; (2) inputting the images into a dual-branch network to output a token sequence and a feature map; (3) performing dual-branch cross-attention pooling on the feature map and the attention map to output the results of disease categories and regions; (4) using the attention map in the CNN branch to perform attention-guided data augmentation on the input of this branch and then inputting it into the main network again; (5) calculating the loss between each result and the label for supervised training, and the trained model can be used for the diagnosis of disease categories and regions. By constructing a dual-branch cross-attention pooling module, the present invention solves the problems of different object sizes and positions in multi-label images and the visual similarity between some features, and can accurately perform multi-label classification of disease categories and regions based on slit lamp images.

Description

Technical Field

[0001] The present invention belongs to the field of computer-aided medical image processing, and relates to a multi-label classification method for anterior segment diseases based on double-branch cross-attention pooling. Background Art

[0002] In the field of ophthalmology, slit lamp examination for the evaluation of the cornea, conjunctiva and lens is the basis for the diagnosis of ocular surface diseases. Early screening of eye diseases can help patients prevent vision impairment and other problems. At present, the existing anterior segment disease diagnosis systems based on slit lamp images mostly target single diseases. Since patients may suffer from multiple eye diseases simultaneously in real life, multi-label classification of anterior segment diseases has great clinical significance and research significance.

[0003] Currently, the methods for multi-label image classification mainly include constructing label correlation, based on object detection bounding boxes and attention mechanisms, etc. The method of calculating the cost by constructing a label co-occurrence matrix is large, and it is easy to cause overfitting when applied to small medical image datasets. The method based on object detection bounding boxes requires high annotation costs and the model is complex. Due to cost and time-consuming reasons, in most cases, doctors cannot densely annotate the lesion areas. The existing methods based on visual attention have been proven to improve the accuracy of multi-label image classification tasks and are applied to medical image datasets, but there are still the following deficiencies: (1) Due to the different positions and sizes of lesions, it is difficult to distinguish and extract features specifically. There is visual similarity between some label objects and it is difficult to distinguish; (2) The number of positive samples in the image is far less than that of negative samples, and there is a serious imbalance between positive and negative samples, resulting in some false negative phenomena that may occur in the prediction results; (3) Most methods ignore the problem of the fine and coarse granularity of labels in multi-label image classification. Summary of the Invention

[0004] In view of the deficiencies of the prior art and the need for classifying anterior segment diseases of slit lamp images, the present invention proposes a multi-label classification method for anterior segment diseases based on double-branch cross-attention pooling.

[0005] To achieve the above object, the present invention includes the following steps:

[0006] S1. Perform data preprocessing on the slit lamp images to unify the image sizes in the dataset;

[0007] S2. Input the images into a double-branch network, and output a slice sequence through the operation of branch 1, and output a feature map through the operation of branch 2;

[0008] S3. Perform double-branch cross-attention pooling on the feature maps and attention maps of branch 1 and branch 2, and output the results of disease types and disease regions;

[0009] S4. Use the attention map output in Branch 2 to perform attention-guided data augmentation on the pre-processed images of this branch, and input the augmented images into the main network again. The obtained results are output as the results of the data augmentation module;

[0010] S5. Calculate the losses of the results output by the main network and the results output by the data augmentation module respectively with the labels, and perform weighted summation of the obtained losses to supervise the training of the model. The model stops training when it reaches the set number of iterations. During this period, the model parameters are saved with the optimal evaluation index results of the validation set;

[0011] Among them, the backbone feature maps of Branch 1 and Branch 2 are extracted by the ResNet50 network model. For Branch 1, the backbone feature map is projected to obtain a slice sequence, which is input into the Transformer module to output a slice sequence. For Branch 2, the backbone feature map is directly used as the feature map and input into the convolutional module to output the feature map. Feature coupling is performed between Branch 1 and Branch 2 through the interaction module.

[0012] Furthermore, in the above multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling, the specific steps of step S3 are as follows:

[0013] S31. Reshape the slice sequence output by Branch 1, and use a convolutional layer with a convolution kernel of 1 to change the number of channels to obtain an aligned feature map, denoted as Aligned Feature Maps.

[0014] At the same time, Branch 2 also outputs a feature map;

[0015] S32. Convolve the feature maps of the two branches respectively using a convolutional layer with a convolution kernel of 1 to obtain attention maps;

[0016] S33. Denote the feature map and attention map of Branch 1 as F1 and A1, and the feature map and attention map of Branch 2 as F2 and A2. Perform bilinear attention pooling on F1 and A1, A2 respectively, and splice the obtained results on the channel with dimension 1, denoted as Feature Matrix 1. Similarly, perform bilinear attention pooling on F2 and A1, A2 respectively, and splice the obtained results on the channel with dimension 1, denoted as Feature Matrix 2;

[0017] S34. Fine-grained disease classification loss. Connect the Feature Matrix 1 and Feature Matrix 2 respectively to a fully connected layer with the number of output channels being the number of disease categories. The number of disease categories is C. The obtained logits pass through the sigmoid activation function to output disease category predictions 1 and disease category predictions 2, and calculate the cross-entropy loss function with the true labels;

[0018] S35. Coarse-grained disease region classification loss. The above C categories belong to the fine-grained level of labels. At the same time, by dividing the diseases according to their regions, coarse-grained labels, that is, region categories, can be obtained, and the number of region categories is R. Therefore, for each input slit lamp image, in addition to the fine-grained disease category label of the image, there is also the coarse-grained disease region label of the image. The calculation method of the coarse-grained disease region prediction is as follows: After the feature matrix 1 and the feature matrix 2, fully connected layers with the number of categories R and ReLU activation functions are connected respectively to obtain the disease region prediction values; after the disease category prediction 1 and the disease category prediction 2, fully connected layers with the number of categories R are connected respectively to obtain the disease region prediction values, and the above-obtained prediction values are added together, and finally input into the sigmoid activation function to obtain the coarse-grained disease region prediction value, and the cross-entropy loss function is calculated with the true label.

[0019] Further, in the above multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling, the step S4 specifically includes the following steps:

[0020] S41. Upsample the attention map output by branch 2 once, and perform attention-guided data augmentation on the image after data preprocessing;

[0021] S42. There are 3 types of attention-guided data augmentation methods, namely cropping, cutmix, and mixup. Cropping means cropping out the region with a higher attention score in the original image after preprocessing and resizing it to the size of the original image; cutmix means cropping the region with a higher attention score in the original image and overlaying it on the upper left corner of the original image; mixup means cropping the region with a higher attention score in the original image and resizing it to the size of the original image, and mixing it with the original image in equal proportion;

[0022] S43. Each of the three attention-guided data augmentation methods can obtain a new input, and this input does not need to be preprocessed again. The three pictures are respectively input into the network, and three results of the data augmentation branch can be obtained.

[0023] Further, in the above multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling, the step S5 specifically includes the following steps:

[0024] S51. Calculate the loss of the output of the main network and the output of the data augmentation module with the label respectively. The loss function of the main network is the cross-entropy loss function, and the data augmentation module is the asymmetry loss function, where the weights of the main network and the data augmentation module are 0.5 respectively. The obtained losses are weighted and summed to supervise the training of the model;

[0025] S52. Stop training when the model is trained to the set number of iterations. During this period, save the model parameters with the optimal evaluation index result of the validation set. The optimal evaluation index of the validation set is defined as the F1-score for predicting the disease type.

[0026] Advantages of the present invention:

[0027] For the multi-label classification method of anterior segment diseases based on dual-branch cross-attention pooling of the present invention, firstly, considering that the positions and sizes of the lesion areas in multi-label images are different, an attention map is selected for feature extraction, and dual-branch cross-attention pooling is used to distinguish easily confused or visually highly similar features, and the multi-label eye images are classified from a fine-grained perspective; secondly, considering the problem of imbalance between positive and negative samples, three different data augmentation methods, namely cropping, cutmix, and mixup, are used to augment the preprocessed original images, highlighting the features of the positive example regions, and using an asymmetric loss function to guide the network to focus on the loss contribution of positive samples during training and reduce the loss contribution of negative samples, which plays a role in alleviating the imbalance between positive and negative samples and correcting some false negative labels; finally, this method considers the coarse and fine-grained problems in multi-label image classification, and realizes the hierarchical classification of multi-label eye images with an end-to-end network structure. The network can simultaneously output accurate disease category information and disease region information. Brief Description of the Drawings

[0028] Figure 1 It is the overall flowchart of the present invention.

[0029] Figure 2 It is the schematic diagram of the overall network structure applied in the present invention;

[0030] Figure 3 It is the schematic diagram of the hierarchical classification module structure for generating disease types and disease regions by dual-branch cross-attention pooling;

[0031] Figure 4 It is the schematic diagram of the three data augmentation module structures guided by attention. Detailed Embodiment

[0032] The technical solutions of the present invention will be further described in detail below with reference to the drawings:

[0033] As Figure 2As shown in the figure, the backbone feature maps of Branch 1 and Branch 2 of the double-branch structure in the main network are extracted by the ResNet50 network model. For Branch 1, the backbone feature map is projected to obtain a slice sequence, which is then input into the Transformer module to output a slice sequence. For Branch 2, the backbone feature map is directly used as the feature map and input into the convolutional module to output a feature map. Feature coupling is performed between Branch 1 and Branch 2 through an interaction module. The interaction module is built based on the feature interaction module in Conformer. This framework uses the Transformer branch to capture global information and the CNN branch to focus on local information, making the overall network more suitable for the small dataset of anterior segment diseases. Different learning paradigms ensure the asymmetry of feature acquisition between the two branches. Considering the particularity of the data in this task and the fact that existing multi-label image classification networks do not consider the visual similarity between each label, the present invention proposes a multi-label classification method for anterior segment diseases based on double-branch cross-attention pooling, which can better classify anterior segment diseases.

[0034] In an embodiment of the present invention, as Figure 1 shown, the following steps are included:

[0035] S1. Preprocess the slit lamp images. Since the input image size of the network needs to be fixed, the original images of different sizes are adjusted in proportion through cropping and scaling to unify the image size to 224×224.

[0036] S2. Input the image into the double-branch network. First, pass it through ResNet50 to obtain the output of its Layer3 as the backbone feature map

[0037] S2.1. For Branch 1, project the backbone feature map to obtain a slice sequence. The projection operation includes: using a convolutional layer with a kernel size of 1 to change the number of channels from 1024 to 384, reshaping the feature map with a resolution of 14×14 into a vector with a length of 196, and changing the channel positions to obtain a slice sequence with a size of 196×384. Input it into the Transformer module to output a slice sequence of the same size.

[0038] S2.2. For Branch 2, directly use the backbone feature map as the feature map and input it into the convolutional module to output a feature map of the same size

[0039] S2.3. Feature coupling is performed between Branch 1 and Branch 2 through the feature interaction module in Conformer.

[0040] S3. Perform double-branch cross-attention pooling on the feature maps and attention maps of Branch 1 and Branch 2 to output the disease types and disease region results;

[0041] The sliced sequence output by Branch 1 After reshaping and using a convolutional layer with a kernel size of 1 to change the number of channels, an aligned feature map is obtained The feature map output by Branch 2 is denoted as

[0042] For the feature maps output by the two branches, convolutional layers with a kernel size of 1 are respectively applied to obtain the corresponding attention maps. The number of channels of the attention maps is the number of disease categories C. The calculation method is as follows

[0043] A = f(F)

[0044] The attention map is calculated

[0045] Perform dual-branch cross-attention pooling on the feature map and the attention map. The dual-branch cross-attention pooling operation is as Figure 3 shown. Taking Branch 1 as an example, the feature map and the attention map of this branch are subjected to bilinear pooling. The calculation method is as follows

[0046] F k = A k ⊙ F(k = 1, 2, …, C)

[0047] For the attention map For each channel of the attention map where ⊙ represents element-wise multiplication. For each channel, calculate and output its corresponding feature map Furthermore, perform a pooling operation on this feature map

[0048] f k = g(F k )

[0049] where the g function represents the pooling operation, which can be an average pooling layer or a max pooling layer. Concatenate the obtained C feature vectors to obtain a feature matrix. Perform the above operations on the feature map of Branch 1 and the attention map of Branch 2 once again to obtain a feature matrix. Concatenate the two feature matrices on the channels with dimension 1 to obtain the total feature matrix 1 of Branch 1. Similarly, the total feature matrix 2 of Branch 2 can be obtained

[0050] The two feature matrices respectively pass through the Dropout layer and the fully connected layer to obtain the prediction results of fine-grained disease types. Further, a fully connected layer and a ReLU activation function are added after the feature matrix to obtain the prediction results of the disease area. A fully connected layer is added after the disease type results to obtain the results of the disease area. The two are connected by residual connection, and the results of the four disease areas are added element-wise to obtain the final results of the disease area. As Figure 1 shown in the process, the two results output by the network are the disease category and the disease area respectively.

[0051] S4. Use the attention map output in Branch 2 to perform attention-guided data augmentation on the pre-processed image of this branch, and input the enhanced image into the main network again. The obtained results are output as the results of the data augmentation module. The original image is x, and after passing through the data pre-processing of the input network, x' is obtained. After attention-guided data augmentation, x'' is obtained as follows:

[0052] x'' = Aug(x')

[0053] As Figure 4 shown, the number of channels of the attention map generated by the main network is C, that is, the number of disease categories. Select the 3 attention maps with the strongest attention scores as the enhanced attention maps. Respectively use these three attention maps and the pre-processed image for attention-guided cropping, cutmix, and mixup augmentation. The obtained new images are respectively input into the main network for calculation to obtain three groups of prediction results of the data augmentation module.

[0054] S5. Calculate the loss of the output results of the main module and the data augmentation module respectively with the label, and sum the obtained losses with weights. Among them, the main module selects the cross-entropy loss function. The data augmentation module uses the asymmetric loss function, and the loss calculation for negative samples is as follows:

[0055] L - = (p m ) γ- log(1 - p m )

[0056] where the calculation of p m is as follows. p represents the probability of negative sample prediction, and both m and γ are hyperparameters:

[0057] p m = max(p - m, 0)

[0058] Perform supervised training on the model. When the model is trained to the set number of iterations, stop training. During this period, save the model parameters with the optimal evaluation index results of the validation set. The optimal evaluation index of the validation set is set as the F1-score for disease type prediction, and the calculation method is as follows:

[0059]

[0060] The embodiments of the present invention described above do not constitute a limitation on the protection scope of the present invention. Any modifications, equivalent replacements, improvements, etc. made within the spirit and principle of the present invention shall be included within the protection scope of the claims of the present invention.

Claims

1. A multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling, characterized in that, It includes the following steps: S1. Obtain a slit lamp image dataset, and perform data preprocessing on the slit lamp images to unify the image sizes in the dataset; S2. Input the preprocessed images into a dual-branch network composed of branch 1 and branch 2. Through the operation of branch 1, a slice sequence is output, and through the operation of branch 2, a feature map F2 is output; In step S2, the backbone feature maps are respectively extracted by branch 1 and branch 2 composed of the ResNet50 network model. The backbone feature map extracted by branch 1 is projected to obtain a slice sequence, which is input into the Transformer module to output the slice sequence; the backbone feature map extracted by branch 2 is directly used as the feature map and input into the convolutional module to output the feature map. Feature coupling is performed between branch 1 and branch 2 through an interaction module. After feature coupling, branch 1 outputs a slice sequence, and the shape of this slice sequence is the same as the shape of the slice sequence before inputting into the interaction module. Branch 2 outputs a feature map, and the shape of this feature map is the same as the shape of the feature map before inputting into the interaction module; S3. Recombine the slice sequence to obtain a feature map F1. Convolution operations are respectively performed on feature map F1 and feature map F2 to obtain attention maps A1 and A2. Dual-branch cross-attention pooling is performed on feature maps F1, F2 and attention maps A1, A2 to output the disease type and disease area results; The method of the dual-branch cross-attention pooling is as follows: Denote the feature map and attention map of branch 1 as F1 and A1, and denote the feature map and attention map of branch 2 as F2 and A2. Bilinear attention pooling is respectively performed on F1 and A1, A2, and the obtained results are concatenated on the channel with dimension 1, denoted as feature matrix 1. Similarly, bilinear attention pooling is respectively performed on F2 and A1, A2, and the obtained results are concatenated on the channel with dimension 1, denoted as feature matrix 2; The method of outputting the disease type and disease area classification results is as follows: The fine-grained disease classification loss respectively connects fully-connected layers with the number of output channels equal to the number of disease categories to the feature matrix 1 and the feature matrix 2. The number of disease categories is C , and the obtained logits pass through the sigmoid activation function to output the disease category prediction 1 and the disease category prediction 2. The cross-entropy loss function is calculated with the true labels. The loss of coarse-grained disease area classification, as described above C These C categories belong to the fine-grained level of labels. At the same time, by dividing the diseases according to their regions, coarse-grained labels are obtained, that is, regional categories, and the number of regional categories is R， The calculation method of coarse-grained disease area prediction is as follows: After the feature matrix 1 and the feature matrix 2, fully connected layers with the number of categories R and ReLU activation functions are connected respectively to obtain the disease area prediction values; After the disease category prediction 1 and the disease category prediction 2, fully connected layers with the number of categories R are connected respectively to obtain the disease area prediction values. Add the above-obtained prediction values and finally input them into the sigmoid activation function to obtain the coarse-grained disease area prediction value, and calculate the cross-entropy loss function with the true label; S4. Use the attention map A2 output in branch 2 to perform attention-guided data augmentation on the images that have undergone data preprocessing through the data augmentation module. The data augmentation module includes cropping, cutmix, and mixup guided by the attention map to obtain three new pictures, denoted as enhanced picture 1, enhanced picture 2, and enhanced picture 3; S5. Respectively pass the enhanced pictures 1, 2, 3 obtained after data augmentation and the corresponding original pictures through the dual-branch network to output four groups of disease type and disease area results. The loss is calculated for each group of prediction results output and the label respectively, and the obtained losses are weighted and summed to supervise the training of the multi-label classification model. The multi-label classification model stops training when it reaches the set number of iterations. During this period, the model parameters are saved with the optimal evaluation index result of the validation set.

2. The multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling according to claim 1, wherein The preprocessing method in step S1 is: unify the image sizes in the dataset and divide the dataset into a training set and a validation set.

3. The multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling according to claim 1, characterized in that The branch 1 is composed of convolutional modules and is used to extract local features of images; branch 2 uses the global receptive field of the self-attention mechanism in Transformer to capture long-range dependencies of images. Feature coupling is performed between branch 1 and branch 2 through an interaction module to fuse local features based on CNN and global features based on Transformer.

4. The multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling according to claim 3, wherein, In the step S3, the method for obtaining the attention map is as follows: The slice sequence output by branch 1 is reshaped, and the number of channels is changed using a convolutional layer with a convolution kernel of 1 to obtain a feature map; the feature maps of the two branches are each convolved using a convolutional layer with a convolution kernel of 1 to obtain an attention map.

5. The multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling according to claim 1, wherein: In the step S4 The attention map output by branch 2 is upsampled once and subjected to attention-guided data augmentation with the pre-processed image. Cropping means cropping out the region with a higher attention score in the original pre-processed image and resizing it to the size of the original image; cutmix means cropping the region with a higher attention score in the original image and overlaying it on the upper left corner of the original image; mixup means cropping the region with a higher attention score in the original image, resizing it to the size of the original image, and mixing it with the original image in equal proportion.

6. The multi-label classification method for anterior segment diseases based on dual-branch cross-attention pooling according to claim 1, characterized in that: In the step S5, for the pre-processed image, it is input into the dual-branch network to generate a set of prediction results, and the loss is calculated with the label, denoted as loss 1. The loss function is the cross-entropy loss function. For enhanced image 1, enhanced image 2, and enhanced image 3, a set of disease types and disease region results are output by the dual-branch network respectively. For each set of prediction results of the data-augmented images, the loss is calculated with the label. The loss function is the asymmetry loss function. The average value of the 3 calculated loss results is taken, and the obtained result is denoted as loss 2. The weights of loss 1 and loss 2 are 0.5 respectively, and weighted summation is performed to obtain the final loss.

Images