Image classification model, model training method and application thereof

Abstract

The application relates to the technical field of deep learning image classification, and provides an image classification model, a model training method and application thereof, the image classification model training method specifically improves the quality and diversity of samples through data enhancement to obtain better training effect; a channel attention mechanism module is introduced into a ResNet50 model to form an SE-ResNet50 model, the channel attention mechanism is used for self-adaptive feature weight learning; then, a multi-scale feature interlayer information fusion module is added in the SE-ResNet50 model; finally, a multi-scale feature fusion self-attention module is introduced into a network structure to form an SEL-ResNet50 network, that is, the image classification model; the image classification model is applied to medical image picture classification, and classification efficiency and precision can be considered.

Description

TECHNICAL FIELD

[0001] The present application belongs to the field of deep learning image classification, and mainly relates to an image classification model, a model training method and application thereof, in particular to an image classification model based on multi-scale fusion attention, a model training method and application thereof. BACKGROUND

[0002] Image classification is the most basic task in computer vision. Today, the computer image classification level in the imagenet dataset has exceeded that of humans. However, there are still many difficulties and challenges in real-world image tasks. For cross-species semantic-level image classification, there is often a large inter-class variance between each class, so some simple models can achieve a high recognition accuracy. However, for some sub-class fine-grained image classification fields, simple models cannot achieve the desired results. Especially in the recognition and classification of medical disease images, there is a higher requirement for recognition accuracy, so a technology that can accurately recognize professional images is in great need.

[0003] In recent years, with the continuous maturity of deep learning technology, computer-aided diagnosis methods based on deep learning can help effectively improve the objectivity and effectiveness of clinical diagnosis. Since medical image graphs have the characteristic that images in the same region have a small inter-class variance, the recognition and classification accuracy of computer-aided medical image recognition in the prior art is not high, and classification models based on deep learning have the advantage of being able to effectively capture deep features of images. At present, the commonly used medical image classification model is a deep learning model, including AlexNet, VGG, GoogleNet, ResNet, DenseNet, U-Net, etc. These models have strong feature extraction capability and good performance in medical image analysis. Chest X-ray is one of the most common and cost-effective medical imaging technologies. However, due to the overlapping and hidden of deep and shallow tissues, multiple and multi-angle X-ray films often need to be taken, which causes the detection image range to be large, the detection area to be small, and the distribution interval to be extensive, making the clinical auxiliary diagnosis based on chest X-ray images quite challenging.

[0004] Recently, in the chest X-ray disease image screening, the literature "CheXNet: Radiologist-Level Pneumonia Detection on Chest X-Rays with Deep Learning" gives a lung disease image classification method based on CheXNet, which uses a convolutional neural network to train and test on the NIH (National Institutes of Health) chest X-ray image dataset, and gets good results of 88.31% for infiltration detection, 72.04% for pleural effusion detection, 92.60% for emphysema, 89.32% for pneumothorax, and 82.09% for atelectasis. The accuracy is higher than that of radiologists with years of experience. However, the method uses a 121-layer DenseNet convolutional neural network during training, which requires high hardware requirements for computer devices and long training time due to the use of deep neural networks. The trained model is large.

[0005] In addition, the Chinese patent application with publication number CN113592809A discloses a pneumonia image detection system and method based on a channel attention residual network. The pneumonia image detection method based on the channel attention residual network migrates the pre-trained ResNet model weights and parameters to the residual network model, introduces the ECA attention module into the residual structure from the channel dimension, and constructs the channel attention-based residual network model ECA-XNet for detecting pneumonia from chest X-rays. However, this method only addresses the problem of pneumonia, and may not perform well in image recognition of multiple disease combinations. SUMMARY

[0006] In view of the above-mentioned defects or deficiencies in the prior art, in view of the limitations of existing image classification models in efficiency and accuracy, and in order to meet the higher classification accuracy requirements of medical image images, the present application provides an image classification model based on an improved ResNet50 attention mechanism, hereinafter referred to as SEL-ResNet50 network, and also provides a training method and application of the above-mentioned image classification model. The above-mentioned image classification model provided by the present application has better discrimination ability when applied to medical images, especially medical X-ray images, which is beneficial to assisting the diagnosis of related diseases.

[0007] The technical solution of the present application is:

[0008] In a first aspect, the present application provides an image classification model training method, specifically a multi-scale fusion attention-based image classification model training method, as shown in Figure 2 The model training method comprises:

[0009] Step one, input the training image in the training sample, data augmentation is performed on the training image, the data augmentation process includes: random probability of flipping and rotating, scaling and cropping from the center area of the image, through the data augmentation, the diversity of the training sample can be increased, and the generalization performance of the model can be improved; then the training image after data augmentation is normalized, and the preprocessed image is output;

[0010] Step two, the preprocessed image obtained in step one is input into a first scale reduction layer, and a first stage image feature X1 is output, then X1 is input into a second scale reduction layer, and X2 is output; X2 is input into a third scale reduction layer, and X3 is output.

[0011] The first scale reduction layer, the second scale reduction layer and the third scale reduction layer are the first three stages of scale reduction layers of each feature map size reduction network layer in the SE-ResNet50 network added with an attention mechanism module.

[0012] The SE-ResNet50 network comprises four scale reduction layers; each scale reduction layer is part of the original SE-ResNet50 network structure, and the size of the feature map is reduced after passing through the structure. Each scale reduction layer is stacked by n SE-ResNet modules added with a channel attention mechanism SE-Block (Squeeze-and-Excitation) module, and n is a positive integer greater than 1; preferably, the first scale reduction layer is stacked by three SE-ResNet modules; the second scale reduction layer is stacked by four SE-Resnet modules; the third scale reduction layer is stacked by six SE-ResNet modules; and the fourth scale reduction layer is stacked by three SE-ResNet modules.

[0013] The operation steps of the channel attention mechanism SE-Block include: performing global pooling on the input feature, then inputting the pooled feature into two perception machine modules and obtaining a weight by using a sigmoid, then multiplying the feature of the original residual block by the weight to obtain a feature map with more channel feature information than the original feature map.

[0014] Step three, the outputs X1, X2 and X3 of the first scale reduction layer, the second scale reduction layer and the third scale reduction layer are input into a hierarchical fusion network at the same time, and are spliced in a new dimension layer after down-sampling transformation, to obtain a feature map Y, and the feature map Y is a fused feature map.

[0015] Step four, after the hierarchical fusion network, a multi-scale feature fusion attention module based on self-attention is adopted to splice the input feature map Y into a matrix and multiply it with its transpose to obtain a weight matrix, the weight matrix is multiplied with the spliced matrix of the input feature map Y after adjusting the dimension to obtain the output feature map X' of the attention network, multiplied by the weight coefficient a and added to X3 as the overall output Z.

[0016] Step five, input the network model output Z into the fourth scale reduction layer, output the feature with the size of 7x7x2048, and directly input the classifier. The processing process of the classifier includes: using adaptive pooling to adjust the output size to 1x1x2048, then changing it to 1x1x15 in size through a linear layer, and finally connecting a sigmoid activation function to adjust the output range to (0, 1).

[0017] Preferably, the flipping and rotating operation of the random probability is specifically that the original training image is flipped with a probability of 50%, and rotated with a random probability in the interval [-10°, 10°];

[0018] Preferably, the scaling and cropping operation from the center region of the image is specifically that the short side size of the image is scaled to 224 and the aspect ratio is kept unchanged, and then a 224x224 square region is cropped from the center of the image.

[0019] Preferably, the specific operation of the hierarchical fusion network is that the feature map output by the first stage is Convolution operation is performed using a convolution kernel with a size of 7x7 and a step of 4 to obtain the down-sampled feature The feature map output by the second stage is Convolution operation is performed using a convolution kernel with a size of 3x3 and a step of 2 to obtain the down-sampled feature Let the feature map output by the third stage be X3, and the fused feature map be Y, then

[0020]

[0021] Where, CAT represents the splicing operation, Conv 7×7 and Conv 3×3 are 7x7 and 3x3 convolution.

[0022] Preferably, the multi-scale feature fusion attention operation of self-attention is represented by the following formula:

[0023]

[0024]

[0025] wherein, sigma represents a sigmoid activation function, represents an operation of adjusting dimensions, which is its corresponding inverse operation.

[0026] the feature map output by the attention network convolution is performed using a convolution kernel with a size of 1*1 and a step of 1 to obtain output Then, it is multiplied by a weight coefficient alpha, and added to the output of the third stage to obtain the entire output Z; the weight of the attention network is adaptively adjusted according to the effect of the attention output, so that higher prediction accuracy is obtained.

[0027] In a second aspect, the present application provides an image classification model obtained by the above training method, that is, a multi-scale fusion attention-based image classification model.

[0028] In a third aspect, the present application provides an application of the above image classification model in medical image classification, that is, an automatic medical image classification method is provided, which comprises: acquiring a clinical medical image, performing scaling and cropping processing on the image to obtain a data-enhanced image; using a normalization method to process the data-enhanced image to generate a preprocessed image for inputting into a network; then inputting the above trained multi-scale fusion attention-based image classification model, and outputting a result as a prediction probability of the image for each disease. In combination with the threshold selected in the training process, the recognition result of the image for the disease is further obtained.

[0029] In a fourth aspect, the present application provides an automatic medical image classification device, which comprises a memory, a processor, and a computer program stored on the memory and executable on the processor, wherein the computer program, when executed by the processor, implements the steps of the above image classification method.

[0030] In a fifth aspect, the present application provides a computer readable storage medium, wherein an information transmission implementation program is stored on the computer readable storage medium, and the program, when executed by a processor, implements the steps of the above image classification method.

[0031] The present application has the beneficial effects that: the present application provides a multi-scale fusion attention-based image classification method, and tests it in the field of medical images, which can judge whether a disease is present and the specific category of the disease based on the NIH dataset image, and can assist doctors in clinical pathological diagnosis in actual application.

[0032] Compared with a CheXNet-based lung disease image classification model training method disclosed in the literature "CheXNet: Radiologist-Level Pneumonia Detection on Chest X-Rays with Deep Learning", the image classification model disclosed in the application adds an attention mechanism in a simpler ResNet network, has smaller complexity and fewer parameters; at the same time, the image classification model has stronger extraction ability for key information regions in feature maps, can filter out more information conducive to judgment from feature information, has lower requirements for computer device hardware in the training process, has shorter training time, and has smaller model size.

[0033] The image classification model disclosed in the Chinese patent application with the publication number CN113592809A is more practical in medical image applications, can not only identify pneumonia as one feature, but also identify images containing multiple disease combinations; the image classification model enhances the feature extraction capability of the network, can simultaneously perform multi-type output on the basis of ensuring image classification accuracy; test results show that the image classification model described in the application can obtain good classification results for a total of 14 diseases in the lungs, and provides doctors with more comprehensive reference opinions. BRIEF DESCRIPTION OF DRAWINGS

[0034] Figure 1 Loss function value change in the training process;

[0035] Figure 2 The image classification model training method described in the application is a whole flowchart;

[0036] Figure 3 The SE-ResNet module structure in the image classification model described in the application;

[0037] Figure 4 The hierarchical fusion network structure in the image classification model described in the application;

[0038] Figure 5 The multi-scale feature fusion attention mechanism network structure in the image classification model described in the application;

[0039] Figure 6 Part of the prediction results of the image classification model described in the application when applied to lung inflammation data set image classification. DETAILED DESCRIPTION

[0040] Step 1, in the preprocessing process, first, the data is enhanced, the original training image is flipped with a probability of 50%, and the angle is randomly selected in the interval [-10°, 10°] for rotation, then the short side size of the image is scaled to 224 and the aspect ratio is kept unchanged, and then a 224x224 square region is cropped from the center of the image; finally, the region image is normalized. After preprocessing, the diversity of the training sample is increased, the robustness of the trained model is improved, and the disease image information is preserved to the greatest extent and the training speed of the network is improved.

[0041] Step 2, as shown in Figure 3 , define each network layer in which the feature map size is reduced in the SE-ResNet50 network with an attention mechanism module as a stage, called a scale reduction layer, so as to divide the SE-ResNet50 into four scale reduction layers. The image obtained in step one is input into the first three scale reduction layers in turn to obtain three stage image features, X1, X2 and X3.

[0042] Based on the above, each scale reduction layer is part of the original SE-ResNet50 network structure, and the size of the feature map is reduced after passing through the structure. The scale reduction layer is stacked by several SE-ResNet modules with channel attention mechanism SE-Block, and the number of stacked four scale reduction layers is 3, 4, 6, and 3.

[0043] Based on the above, the operation process of the channel attention mechanism SE-Block is to first perform global pooling on the feature, then input the pooled feature into two perception modules and obtain the weight using sigmoid, then multiply the original residual block feature and assign the weight to obtain the feature map with more channel feature information than the original.

[0044] Specifically, when the SE-ResNet module input is X in , let its output after the residual module be At the same time, the feature map as the input of the SE-Block; the global pooling layer uses global average pooling, and let the pooling operation output be The global average pooling operation is represented as GAP, and has:

[0045]

[0046] Where subscript c represents the cth channel, H and W represent the length and width of U, and i and j represent enumeration from 1 to H and from 1 to W, respectively.

[0047] The perceptron module consists of two fully connected layers, which then parameterize the channel selection mechanism. The two fully connected layers are a fully connected layer W1 with a dimensionality reduction ratio of 16 and a fully connected layer W2 with increased dimensionality. Both fully connected layers use the ReLU activation function. Then, the sigmoid activation function is used to obtain the weights. The corresponding formula is expressed as:

[0048] s=σ(W2(δ(W1(z))))

[0049] Where δ represents the ReLU function and σ represents the sigmoid activation function;

[0050] Finally, a scale operation is performed, which involves multiplying the output of the SE-Block with the output of the original residual block module to obtain the final weighted output value:

[0051]

[0052] Furthermore, in a separate SE-Resnet module, the output X is... out The expression satisfies the following:

[0053]

[0054] Furthermore, the first scale-down layer consists of three such SE-ResNet modules cascaded together; the second scale-down layer consists of four such SE-ResNet modules cascaded together; the third scale-down layer consists of six such SE-ResNet modules cascaded together; and the fourth scale-down layer consists of three such SE-ResNet modules cascaded together.

[0055] Step 3, as follows Figure 4 As shown, a hierarchical fusion network is added to the overall structure of the improved model to fuse multi-scale features of the image, further enhancing the model's ability to extract image features and thus improving accuracy. The outputs X1, X2, and X3 of the three scale reduction layers are fed into the hierarchical fusion network, and after downsampling transformation, they are concatenated at a new dimension to obtain the feature map Y.

[0056] Based on the above, the specific process of the hierarchical fusion network is as follows: The feature map output from the first stage... A convolution operation is performed using a 7×7 kernel with a stride of 4 to obtain the downsampled features. Feature map output from the second stage A convolution operation is performed using a 3×3 kernel with a stride of 2 to obtain the downsampled features. Let the feature map output of the third stage be X3, and the fused feature map be Y, then

[0057]

[0058] where CAT represents the concatenation operation, Conv 7×7 and Conv 3×3 are 7x7 and 3x3 convolution.

[0059] Step 4, as shown in Figure 5 , a multi-scale feature fusion attention network is constructed to further enhance the model's ability to extract information of different scales. The input feature map Y is concatenated into a matrix and multiplied by its transpose to obtain a weight matrix. The weight matrix is multiplied by the adjusted dimension to serve as the output feature map X' of the attention network, multiplied by the weight coefficient a and added to X3 as the overall output Z.

[0060] Based on the above, the multi-scale feature fusion attention network adopts a self-attention mechanism. The input fused feature map is concatenated and adjusted in dimension to a matrix F with a size of NxHWC, where N is the number of sub-maps of the concatenated feature map. F is multiplied by its transpose to serve as its weight matrix, and the matrix F is multiplied by the adjusted dimension to obtain the output of the attention network. Let the input multi-scale attention network be the feature map concatenation Y, and the attention mechanism is represented by the following formula:

[0061]

[0062]

[0063] where σ represents the sigmoid activation function, represents the dimension adjustment operation, and it is the inverse operation thereof.

[0064] Further, the feature map output of the attention network is convolved using a 1x1 convolution kernel with a step size of 1 to obtain the output which is multiplied by the weight coefficient a and added to the third stage output to serve as the overall output Z. The corresponding formula is represented as:

[0065]

[0066]

[0067] The self-attention structure can adaptively adjust the weight of the attention network according to the effect of the attention output, thereby obtaining higher prediction accuracy.

[0068] Step 5, input the network model output Z into the fourth scale reduction layer to obtain features with a size of 7x7x2048, and input the features into a classifier. The processing process of the classifier includes: adjusting the size of the output to 1x1x2048 by using adaptive pooling, changing the size to 1x1x15 through a linear layer, and finally inputting a sigmoid activation function to adjust the range of the output to (0, 1).

[0069] According to the above implementation process, the working principle of the present application is summarized as follows:

[0070] For images, first perform data preprocessing in Step 1, including data augmentation, normalization, etc.

[0071] Under the improved model, the image data is trained and predicted, and the improvement includes the addition of SE-Block channel attention mechanism, the introduction of hierarchical fusion network and the addition of multi-scale feature fusion attention network, covering the contents of Step 2, Step 3 and Step 4.

[0072] Thus, an effective classification model for images of the same type as the training images is obtained, and when applied to X-ray images, it can assist doctors in judging the disease type.

[0073] In order to further illustrate the application effect of the present application, the experimental results of the present application are shown:

[0074] The experimental environment of the present application is: processor 4-core Intel(R) Xeon(R) Silver 4110 CPU@2.10GHz, memory: 15GB. The graphics card is NVDIA GeForce RTX 2080 Ti, and the video memory is 11GB.

[0075] The present application trains and tests the disease classification on the Chest X-ray 14 data set provided by NIH. The data set is divided into a training set (76241 images), a validation set (13455 images) and a test set (22424 images). The initial epoch is set to 50, and the batch size is set to 64. The Adam optimizer is used. Since each disease may occur in the same person, the present application uses the binary cross-entropy BCELoss function to convert the problem into 15 binary classification problems for processing.

[0076] The learning rate decay strategy used in training is: if there is no lower loss function value calculated in the validation set for three consecutive epochs, then update the learning rate by a multiplication factor of 0.5; if there is no lower loss function value calculated in the validation set for 10 consecutive epochs, it means that the model is basically stable, and the training can be exited.

[0077] As shown in Figure 1 the loss value of the validation set after 35 epochs, the downward trend is not obvious, which indicates that the model is basically stable, and further training may produce overfitting phenomenon. At this time, the model corresponding to the epoch with the lowest loss value on the validation set in all epochs is selected as the final model.

[0078] As shown in Figure 6 the model is trained and tested on the Chest X-ray 14 data set provided by NIH for disease classification. After training, the classification test is performed on the 12 X-ray images in the figure, and a good classification effect is obtained.

[0079] Table 1, the average accuracy and average AUC of three networks on the data set

[0080] Network model Accuracy AUC Resnet 0.895 0.810 SE-Resnet 0.897 0.836 SEL-Resnet 0.907 0.836

[0081] Table 1 shows the experimental results of three networks on the NIH data set, where the accuracy and AUC area are the average of 15 classifications. Compared with SE-ResNet, the accuracy of SEL-ResNet designed by the patent increases by 1%, and the AUC value does not decrease.

Claims

1. An image classification model training method, characterized in that, The model training method comprises: Step one, input the training image in the training sample, perform data enhancement on the training image, the data enhancement process comprises: random probability flip and rotation, zooming and cropping processing from the image center area, through the data enhancement, the diversity of the training sample can be increased, and the generalization performance of the model can be improved; then, the training image after the data enhancement is normalized, and a preprocessed image is output; Step two, the preprocessed image obtained in step one is input into a first scale reduction layer, and a first stage image feature X1 is output, then X1 is input into a second scale reduction layer, and X2 is output; X2 is input into a third scale reduction layer, and X3 is output; the first scale reduction layer, the second scale reduction layer and the third scale reduction layer are the first three stages of scale reduction layers of each feature map size reduction network layer in an SE-ResNet50 network to which an attention mechanism module is added; the SE-ResNet50 network comprises four scale reduction layers; each scale reduction layer is part of the original SE-ResNet50 network, and the size of the feature map is reduced after passing through the structure; each scale reduction layer is stacked by n SE-ResNet modules to which a channel attention mechanism SE-Block is added, and n is a positive integer greater than 1; The operation steps of the channel attention mechanism SE-Block comprise: performing a global pooling operation on the input feature, then inputting the pooled feature into two perception machine modules and adopting a sigmoid activation function to obtain a weight, then multiplying the input feature of the original residual block and assigning the weight to obtain a feature map with more channel feature information than the original; Step three, X1, X2 and X3 output by the first scale reduction layer, the second scale reduction layer and the third scale reduction layer are input into a hierarchical fusion network at the same time, and after down-sampling conversion, they are spliced in a new dimension layer to obtain a feature map Y, and the feature map Y is taken as a fused feature map; Step four, after the hierarchical fusion network, a multi-scale feature fusion attention module based on self-attention is adopted, the input feature map Y is spliced into a matrix and multiplied by its transpose to obtain a weight matrix, the weight matrix is multiplied by the spliced matrix of the input feature map Y after adjusting the dimension, and the output feature map X' of the attention network is obtained, multiplied by a weight coefficient alpha and added to X3 as the overall output Z; Step five, the output Z of the network model is input into a fourth scale reduction layer, and a feature with a size of 7x7x2048 is output, and then directly input into a classifier; the processing process of the classifier comprises: adopting adaptive pooling to adjust the output size to 1x1x2048, then changing it to 1x1x15 through a linear layer, and finally connecting a sigmoid activation function to adjust the output range to (0, 1).

2. The method of claim 1, wherein, The first scale reduction layer is formed by 3 SE-ResNet module stacks in series; the second scale reduction layer is formed by 4 SE-ResNet module stacks in series; the third scale reduction layer is formed by 6 SE-ResNet module stacks in series; and the fourth scale reduction layer is formed by 3 SE-ResNet module stacks in series.

3. The method of claim 1, wherein, The flipping and rotating operation of the random probability is specifically that the original training image is flipped with a probability of 50%, and rotated with a random angle in the interval [-10°, 10°].

4. The method of claim 1, wherein, The scaling and cropping operation from the center region of the image is specifically that the short side size of the image is scaled to 224 and the aspect ratio is kept unchanged, and then a square region of 224x224 size is cropped from the center of the image.

5. The method of claim 1, wherein, The specific operation of the hierarchical fusion network is: performing convolution operation on the feature map output by the first stage Convolution operation is performed by using a convolution kernel with a size of 7*7 and a step of 4 to obtain down-sampled features The feature map output by the second stage is Convolution operation is performed by using a convolution kernel with a size of 3*3 and a step of 2 to obtain down-sampled features Supposing that the feature map output by the third stage is X3 and the fused feature map is Y, then where CAT denotes a concatenation operation, Conv 7×7 and Conv 3×3 are 7x7 and 3x3 convolutions.

6. The method of claim 1, wherein, The multi-scale feature fusion attention operation of the self-attention is expressed by the following formula: where σ denotes a sigmoid activation function, denotes an operation of adjusting dimensions, is its corresponding inverse operation; Feature map output by attention network Convolution is performed by using a convolution kernel with a size of 1*1 and a step of 1 to obtain output Then, it is multiplied by a weight coefficient a, and added to the output of the third stage to obtain the whole output Z; the weight of the attention network is adaptively adjusted according to the effect of the attention output, so as to obtain higher prediction accuracy.

7. The image classification model obtained by the image classification model training method according to any one of claims 1-6.

8. A method of automatic classification of medical image images, characterized in that, The method comprises: acquiring a clinical medical image, performing scaling and cropping processing on the image to obtain a data-enhanced image; using a normalization method to process the data-enhanced image to generate a preprocessed image; then inputting the preprocessed image into the image classification model which has been trained as claimed in claim 7, and outputting a prediction probability of suffering from each disease corresponding to the clinical medical image.

9. A medical image automatic classification apparatus characterized by, Comprise: A memory, a processor, and a computer program stored on the memory and executable on the processor, wherein the computer program, when executed by the processor, implements the steps of the medical image automatic classification method of claim 8.

10. A computer-readable storage medium, characterized in that, The computer readable storage medium stores an information transmission implementation program, and the program, when executed by the processor, implements the steps of the medical image automatic classification method of claim 8.

Images