An image classification model training method, an image classification method, and a system

Abstract

The present invention provides an image classification model training method and an image classification method and system. The image classification model training method includes: obtaining an initial image sample and performing data enhancement on the initial image sample to obtain an enhanced image sample; inputting the initial image sample and the enhanced image sample into a preset image classification model to obtain an initial output feature and an enhanced output feature; fitting the enhanced output feature with the initial output feature to obtain a feature fitting result; splicing the initial output feature and the enhanced output feature to obtain a fused feature, and fitting the fused feature with a preset label vector to obtain a label fitting result; updating parameters in the model according to the feature fitting result and the label fitting result until a trained image classification model is obtained. The present invention improves the generalization ability of the image classification model by performing data enhancement on the sample, and ensures the classification performance of the model trained by the sample, thereby improving the performance of the model.

Description

Technical Field

[0001] The present invention relates to the field of data processing technology, and in particular to an image classification model training method and an image classification method and system. Background Art

[0002] With the development and application of new technologies such as cloud services, big data, and artificial intelligence, deep learning networks have gradually been introduced into the fields of image classification and image recognition, avoiding the time-consuming and high-error rate problems of human eye observation. Deep learning has certain requirements on the size of the data set. If the original data set is relatively small, it cannot meet the training needs of the network model well. Image enhancement is to process the original image to expand the data set, which can improve the performance of the model to a certain extent. However, there will be an imbalance between the original data and the enhanced data. In addition, the existing classification methods focus more on data enhancement of a single channel of the model, ignoring the common knowledge of the original image and the image after data enhancement, which in turn affects the performance of the entire image classification model. Summary of the Invention

[0003] In view of this, an embodiment of the present invention provides an image classification model training method and an image classification method and system, which solves the problem that existing classification methods ignore the common knowledge of the original image and the image after data enhancement, thereby affecting the performance of the entire image classification model.

[0004] According to the first aspect, an embodiment of the present invention provides an image classification model training method, including: obtaining an initial image sample, and performing data enhancement on the initial image sample to obtain an enhanced image sample; inputting the initial image sample and the enhanced image sample into a preset image classification model to obtain corresponding initial output features and enhanced output features; using a feature-level loss function to fit the enhanced output features with the initial output features to obtain a feature fitting result; splicing the initial output features and the enhanced output features to obtain a fusion feature, and using a label-level loss function to fit the fusion feature with a preset label vector to obtain a label fitting result; updating the parameters in the preset image classification model according to the feature fitting result and the label fitting result until a trained image classification model is obtained.

[0005] Optionally, the use of a label-level loss function to fit the fused features and a preset label vector to obtain a label fitting result includes: classifying the fused features using a preset classification function to obtain a probability distribution of the classification results; converting the probability distribution of the classification results into a probability vector, and fitting the probability vector and the preset label vector using a label-level loss function to obtain a label fitting result.

[0006] Optionally, the feature-level loss function is expressed as follows:

[0007]

[0008] Among them, Loss1 represents the feature-level loss function, MSE represents the Mean Square Error loss function, n represents the vector dimension of the output feature, r represents the initial output feature, and r′ represents the enhanced output feature.

[0009] Optionally, the label-level loss function is expressed as follows:

[0010]

[0011] Among them, Loss2 represents the label level loss function, CE represents the Cross Entropy loss function, and y true (x i ) is the picture x i The output probability, y pred (x i ) is the picture x i The corresponding sample true label.

[0012] Optionally, after obtaining the enhanced image samples, the image classification model training method further includes: performing standardization processing on the initial image samples and the enhanced image samples to obtain standard initial image samples and standard enhanced image samples.

[0013] According to the second aspect, an embodiment of the present invention provides an image classification method, including: obtaining an image to be classified, performing data enhancement on the image to be classified, and obtaining an enhanced image; using the image classification model training method described in the first aspect of the present invention and any optional method to train a preset image classification model to obtain a trained image classification model; inputting the image to be classified and the enhanced image into the trained image classification model respectively to obtain corresponding initial output features and enhanced output features, and splicing the initial output features and the enhanced output features to obtain fused features; using a preset classification function to classify the fused features to obtain a classification result.

[0014] According to the third aspect, an embodiment of the present invention provides an image classification model training system, including: a first processing module, used to obtain an initial image sample, and perform data enhancement on the initial image sample to obtain an enhanced image sample; a second processing module, used to input the initial image sample and the enhanced image sample into a preset image classification model to obtain corresponding initial output features and enhanced output features; a third processing module, used to fit the enhanced output features with the initial output features using a feature-level loss function to obtain a feature fitting result; a fourth processing module, used to splice the initial output features and the enhanced output features to obtain a fused feature, and use a label-level loss function to fit the fused feature with a preset label vector to obtain a label fitting result; a fifth processing module, used to update the parameters in the classification model according to the feature fitting result and the label fitting result, until a trained image classification model is obtained.

[0015] According to the fourth aspect, an embodiment of the present invention provides an image classification system, including: an acquisition module, used to acquire an image to be classified, perform data enhancement on the image to be classified, and obtain an enhanced image; a model training module, used to train a preset image classification model using the image classification model training system described in the third aspect of the present invention and any optional method, to obtain a trained image classification model; a feature extraction module, used to input the image to be classified and the enhanced image into the trained image classification model respectively, to obtain corresponding initial output features and enhanced output features respectively, and to splice the initial output features and the enhanced output features to obtain a fused feature; a classification module, used to classify the fused feature using a preset classification function to obtain a classification result.

[0016] An embodiment of the present invention provides a non-transitory computer-readable storage medium, which stores computer instructions. When the computer instructions are executed by a processor, they implement the image classification model training method described in the first aspect of the present invention and any optional method, or implement the image classification method described in the second aspect of the present invention and any optional method.

[0017] An embodiment of the present invention provides an electronic device, comprising: a memory and a processor, wherein the memory and the processor are communicatively connected to each other, the memory storing computer instructions, and the processor executing the computer instructions to thereby execute the image classification model training method described in the first aspect of the present invention and any optional embodiment, or to execute the image classification method described in the second aspect of the present invention and any optional embodiment.

[0018] The technical solution of the present invention has the following advantages:

[0019] 1. An embodiment of the present invention provides an image classification model training method, which performs data enhancement on initial image samples with known labels, inputs the initial image samples and the enhanced image samples into a preset image classification model to obtain corresponding initial output features and enhanced output features, and updates the parameters in the preset image classification model by performing feature-level fitting and label-level fitting on the output features, until a trained image classification model is obtained; the generalization ability of the image classification model is improved by performing data enhancement on the samples, so that balance can be maintained when the sample data set is severely unbalanced, and the classification performance of the model trained by the samples is guaranteed; the dual-channel data enhancement technology is used to learn not only the information of the data itself, but also the information of the corresponding data after data enhancement, so that the classification model can obtain common knowledge of the original image and the image after data enhancement, thereby reducing the redundancy disadvantage caused by the oversampling technology and improving the performance of the model.

[0020] 2. An embodiment of the present invention provides an image classification method, which trains a preset image classification model through an image classification model training method to obtain a trained image classification model, and inputs the image to be classified and the image after data enhancement into the trained image classification model to obtain the final image classification result; thereby achieving accurate classification of images in the case of unbalanced data samples, and no human participation is required in the classification detection process, which reduces human factors and enables better application of detection results and judgment of image classification results, providing an efficient judgment reference. BRIEF DESCRIPTION OF THE DRAWINGS

[0021] In order to more clearly illustrate the specific embodiments of the present invention or the technical solutions in the prior art, the following briefly introduces the drawings required for use in the specific embodiments or the description of the prior art. Obviously, the drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0022] Figure 1 Flowchart of the image classification model training method according to an embodiment of the present invention;

[0023] Figure 2 This is another specific flow chart of the image classification model training method according to an embodiment of the present invention;

[0024] Figure 3 Schematic diagram of the network structure of the image classification model training method in an embodiment of the present invention;

[0025] Figure 4 Schematic diagram of a preset image classification model in an embodiment of the present invention;

[0026] Figure 5Schematic diagram of an image classification model training system according to an embodiment of the present invention;

[0027] Figure 6 is a schematic diagram of an image classification system according to an embodiment of the present invention;

[0028] Figure 7 Schematic diagram of the structure of an electronic device in an embodiment of the present invention. DETAILED DESCRIPTION

[0029] The technical solution of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the embodiments described are only some embodiments of the present invention, not all embodiments. All other embodiments obtained by ordinary technicians in this field based on the embodiments of the present invention without making any creative efforts shall fall within the scope of protection of the present invention.

[0030] In addition, the technical features involved in the different embodiments of the present invention described below can be combined with each other as long as they do not conflict with each other.

[0031] For the situation where image samples with large differences in sample numbers cannot adapt to the network model for training, take gastroscopy images to assist in gastric cancer judgment as an example. Although gastroscopy can directly detect the diseased tissue area in the stomach and make corresponding judgments, tissue biopsy can be performed under gastroscopy. It is important for the diagnosis of early gastric cancer precancerous diseases or precancerous lesions and the differentiation of benign and malignant ulcers. However, due to human factors such as inconsistent doctor's experience level or negligence and other special circumstances, it will directly affect the final gastric cancer result judgment. At the same time, human eyes observing gastroscopy images will also consume a lot of time.

[0032] Deep learning networks are gradually being used in the fields of image classification and image recognition, avoiding the time-consuming and high-error rate problems of human visual observation. Deep learning has certain requirements for the size of the data set. For example, when judging gastric cancer results from gastroscopy images, there will be a serious imbalance in samples. Existing classification methods focus more on data enhancement of a single channel of the model, which in turn affects the performance of the entire image classification model.

[0033] The embodiment of the present invention provides an image classification model training method, such as Figure 1 As shown, the image classification model training method specifically includes:

[0034] Step S1: Obtain an initial image sample, and perform data enhancement on the initial image sample to obtain an enhanced image sample.

[0035] In an embodiment of the present invention, an initial image sample is first obtained. Taking the above-mentioned gastroscopy image as an example, the label categories in this initial sample are unknown. It contains a very small number of images that are confirmed to be gastric cancer, and also contains a large number of images that are not gastric cancer. For subsequent model training, the initial image sample is subjected to data enhancement to obtain an enhanced image sample. It should be noted that new data can be created by translating, rotating, compressing, adjusting the color of the image, etc., specifically including horizontal or vertical flipping, 90-degree rotation, increasing the brightness by one or two levels, increasing the contrast, and changing the hue. Although the new enhanced image sample changes the appearance to a certain extent, the label category of the sample remains unchanged, and by using limited data with known labeled categories, more training data is obtained, which reduces the overfitting phenomenon in the network and trains a model with stronger generalization ability.

[0036] It should be noted that the embodiments of the present invention only illustrate a method for performing data enhancement on image samples. In actual applications, other methods may be selected for data enhancement. The present invention only takes this as an example and is not limited thereto.

[0037] Step S2: Input the initial image samples and the enhanced image samples into a preset image classification model to obtain corresponding initial output features and enhanced output features.

[0038] In an embodiment of the present invention, the initial image samples and the enhanced image samples are input into a preset image classification model. This model is mainly composed of ResNet-50, which consists of 50 weighted layers. The convolution kernels used are 1*1, 3*3, and 1*1. ResNet-50 is mainly divided into 6 parts, namely conv1, conv2_x, conv3_x, conv4_x, conv5_x, and the final global average pooling layer and fully connected layer. The global average pooling layer maps the two-dimensional feature map into a vector and inputs it into the fully connected layer. The fully connected layer includes two layers. The first fully connected layer contains 1000 neurons and the second fully connected layer contains 100 neurons to obtain the final output features corresponding to the input image. That is, the initial image samples are input into the preset image classification model to obtain the corresponding initial output features, and the enhanced image samples are input into the preset image classification model to obtain the corresponding enhanced output features.

[0039] It should be noted that in the embodiments of the present invention, only the BiT-ResNet model is selected as the preset image classification model. In actual applications, other network models can be selected according to sample data and actual needs, such as deep residual shrinkage network, ultra-deep convolutional network (VGG-16) and other mainstream image classification training models. The present invention is not limited to this.

[0040] Step S3: Use the feature-level loss function to fit the enhanced output features with the initial output features to obtain feature fitting results.

[0041] In an embodiment of the present invention, the initial image sample x and the enhanced image sample x′ obtained after data enhancement are input into the same preset image classification model, and the corresponding initial output features r and enhanced output features r′ are respectively obtained. The enhanced output features are fitted with the initial output features using the feature-level loss function to obtain a feature fitting result. Since the enhanced image sample is obtained after data enhancement of the initial image sample, the final fitting result should ensure that the enhanced output features are infinitely close to the initial output features.

[0042] Among them, the feature-level loss function is expressed by the following formula:

[0043]

[0044] Wherein, Loss1 represents the feature-level loss function, MSE represents the Mean Square Error loss function, n represents the vector dimension of the output feature, r represents the initial output feature, and r′ represents the enhanced output feature. The embodiment of the present invention only illustrates the use of the Mean Square Error loss function as an example. In actual applications, other loss functions can be selected based on actual needs, and the present invention is not limited thereto.

[0045] Step S4: The initial output features and the enhanced output features are concatenated to obtain fused features, and the fused features are fitted to the preset label vector using a label-level loss function to obtain a label fitting result. In an embodiment of the present invention, a fused feature is obtained, and the fused feature is fitted to the preset label vector using a label-level loss function to obtain a label fitting result. The label predicted by the fused feature should be infinitely close to the preset label vector, where the preset label vector is a vector composed of the known labels of each initial image.

[0046] Step S5: According to the feature fitting results and the label fitting results, the parameters in the preset image classification model are updated until a trained image classification model is obtained.

[0047] The image classification model training method provided by the present invention performs data enhancement on initial image samples with known labels, inputs the initial image samples and the enhanced image samples into a preset image classification model to obtain corresponding initial output features and enhanced output features, and updates the parameters in the preset image classification model by performing feature-level fitting and label-level fitting on the output features, until a trained image classification model is obtained; the generalization ability of the image classification model is improved by performing data enhancement on the samples, so that balance can be maintained when the sample data set is severely unbalanced, and the classification performance of the model trained by the samples is guaranteed; the dual-channel data enhancement technology is used to learn not only the information of the data itself but also the information of the corresponding data after data enhancement, so that the classification model can obtain common knowledge of the original image and the image after data enhancement, thereby reducing the redundancy disadvantage caused by the oversampling technology and improving the performance of the model.

[0048] Specifically, in one embodiment, the above step S4 specifically includes the following steps:

[0049] Step S41: Classify the fused features using a preset classification function to obtain a classification result probability distribution. In this embodiment of the present invention, the fused features are classified using a preset classification function to obtain a classification result probability distribution. The final classification result probability distribution P is obtained through a linear classification layer, where P = [p1, p2], p1 is the probability that the image result reference is gastric cancer, p2 is the probability that the image result reference is not gastric cancer, p1 + p2 = 1, when p1 > = 0.5 and p2 = < 0.5, it indicates that the network determines that the image is a gastric cancer image, and when p1 < 0.5 and p2 > 0.5, it indicates that the network determines that the image is not a gastric cancer image.

[0050] Step S42: Convert the probability distribution of the classification result into a probability vector, and use the label-level loss function to fit the probability vector with the preset label vector to obtain a label fitting result.

[0051] In this embodiment of the present invention, the probability distribution of the classification result is converted into a probability vector. Before calculating the loss function at the label level, for each image, the output probability distribution P is converted into a probability scalar p to represent the probability of whether the image result reference is gastric cancer. If p1 in P>=0.5, the image result reference is determined to be gastric cancer, and p=1 is set; if p1 in P<0.5, it is determined not to be gastric cancer, and p=0 is set. Since each image has a classification result p, we combine them to obtain the probability vector y pred , the corresponding sample formal label vector is y true , the label-level loss function is expressed by the following formula:

[0052]

[0053] Among them, Loss2 represents the label level loss function, CE represents the Cross Entropy loss function, and y true (x i ) is the picture x i The output probability, y pred (x i ) is the picture x i It should be noted that the embodiment of the present invention only illustrates the use of the Cross Entropy loss function as an example loss function. In practical applications, other loss functions can be selected according to actual needs, and the present invention is not limited thereto.

[0054] Specifically, in one embodiment, after the above step S1, the image classification model training method provided by the embodiment of the present invention further includes the following steps:

[0055] Step S6: performing standardization processing on the initial image samples and the enhanced image samples to obtain standard initial image samples and standard enhanced image samples.

[0056] In an embodiment of the present invention, the initial image samples and enhanced image samples are standardized. The standardization process includes scaling, cropping, random horizontal flipping, and normalization. Scaling involves scaling all different input images to 256*256*3, followed by cropping to 224*224*3 to remove the redundant black portions at the four corners of the gastroscopic image. Random flipping is controlled by a parameter P, which represents the probability of image flipping. Finally, data normalization is performed to obtain the final standard initial image samples and standard enhanced image samples.

[0057] Through the above steps S1 to S6, an embodiment of the present invention provides an image classification model training method, which performs data enhancement on initial image samples with known labels, inputs the initial image samples and the enhanced image samples into a preset image classification model to obtain corresponding initial output features and enhanced output features, and updates the parameters in the preset image classification model by performing feature-level fitting and label-level fitting on the output features, until a trained image classification model is obtained; the generalization ability of the image classification model is improved by performing data enhancement on the samples, so that balance can be maintained when the sample data set is severely unbalanced, and the classification performance of the model trained by the samples is guaranteed; through the dual-channel data enhancement technology, not only the information of the data itself is learned, but also the information of the corresponding data after data enhancement is learned, so that the classification model can obtain common knowledge of the original image and the image after data enhancement, reducing the redundancy disadvantages brought by the oversampling technology and improving the performance of the model.

[0058] The embodiment of the present invention also provides an image classification method, such as Figure 2 As shown, the image classification method specifically includes:

[0059] Step S01: Obtain an image to be classified, perform data enhancement on the image to be classified, and obtain an enhanced image. For details, please refer to the description of data enhancement in the above-mentioned image classification model training method embodiment.

[0060] Step S02: Use the image classification model training method to train the preset image classification model to obtain a trained image classification model. For details, please refer to the description of each step in the above-mentioned image classification model training method embodiment.

[0061] Step S03: Input the image to be classified and the enhanced image into the trained image classification model, respectively, to obtain corresponding initial output features and enhanced output features, respectively. The initial output features and enhanced output features are then concatenated to obtain fused features. For details, see the description of the step of obtaining fused features in the above-mentioned embodiment of the image classification model training method.

[0062] Step S04: Classify the fused features using a preset classification function to obtain a classification result. In an embodiment of the present invention, the preset classification function may be a linear classification function that performs linear classification on the fused features to obtain a classification result. For example, a Softmax classification function may be selected for classification to obtain the probability of the model. Other classification functions may also be selected, and the present invention is not limited thereto.

[0063] In practical applications, the image classification method provided by the embodiment of the present invention can be used to assist in the judgment process of gastric cancer results. The training image samples include 140 images with actual gastric cancer results and 560 other images. Therefore, the ratio of positive samples to negative samples is 1:4, and the samples are greatly unbalanced. Figure 3 The network structure used in the embodiment of the present invention is shown, where x is the input image of the network model. This input image and the image x′ after data enhancement are input into two identical image classification models (BiT-ResNet models), and the two corresponding output features r and r′ are output respectively. The two output features are fitted using the feature-level loss function (Loss1), and then the two output features are spliced. The spliced fusion features are classified in the linear classification layer to obtain the probability (P) distribution of the image. The result corresponding to this probability is fitted with the known image label result (label) through the label-level fitting function (Loss2). Finally, the model training is completed to obtain a more accurate image classification result. Figure 4The BiT-ResNet model is the default image classification model shown. This model is primarily composed of ResNet-50, consisting of 50 weighted layers and using convolution kernels of 1*1, 3*3, and 1*1. ResNet-50 consists of six parts: conv1, conv2_x, conv3_x, conv4_x, conv5_x, and finally the global average pooling layer (fc1) and the fully connected layer (fc2).

[0064] The initial image x is input and its corresponding data-augmented image x′ is generated in real time. Two features, r and r′, are derived from x and x′, respectively, using a BiT-ResNet model trained on a gastroscopic image dataset. These features are then concatenated and passed through a linear classification layer to produce the final classification distribution P = [p1, p2], where p1 is the probability that the reference gastroscopic image is gastric cancer, and p2 is the probability that the image is otherwise classified. For p1 + p2 = 1, when p1 > = 0.5 and p2 < 0.5, the network determines that the reference image is gastric cancer. When p1 < 0.5 and p2 > 0.5, the network determines that the image is not gastric cancer. For the test image (initial image sample) x, only one forward pass through the network is required to obtain the classification result. During training, the image is compressed to 256 × 256 × 3 pixels, then cropped to 224 × 224 × 3 pixels, the four redundant black edges of the gastroscopic image are trimmed, and then randomly horizontally flipped. Finally, the image is normalized to obtain the input features. The initial learning rate is set to 0.0001. The Adam (Adaptive Moment Estimation) optimizer is used to optimize the network parameters. The feature-level loss function uses the mean square error loss function, and the label-level loss function uses the cross entropy loss function as the model loss function. The batch size is set to 64. To prevent overfitting, dropout is used to randomly ignore some neurons in the fully connected layer.

[0065] After obtaining the image to be classified, the image to be classified is resized to 224*224*3. The data-augmented image is passed through the same pre-trained BiT-ResNet model that has been trained on gastroscopy images to obtain two output features. The two output features are spliced and passed through the final linear classification layer to obtain the final classification result.

[0066] Comparisons were made using undersampling, oversampling, ordinary dual-channel data enhancement, and the image classification method provided by the present invention. The undersampling technique randomly removes negative samples so that the number of negative samples is consistent with the number of positive samples. The oversampling technique replicates positive samples multiple times so that the number of positive samples is the same as the number of negative samples. The ordinary dual-channel data enhancement technique simultaneously inputs a gastroscopy image and its real-time data-enhanced image, and splices the results to obtain a classification result. The image classification method of an embodiment of the present invention fits the output features based on the original dual-channel data enhancement technique. According to the results shown in Table 1, although the undersampling technology makes the positive and negative samples consistent and the training samples are not repeated, a large number of negative samples are lost, which greatly reduces the information learned by the model. Therefore, the performance of the undersampling technology is poor, with an accuracy of only 72.2%. The oversampling technology retains more negative sample information than undersampling, but also adds a lot of redundant information, which affects the performance of the model. Its classification results are better than undersampling, and the accuracy is improved by 10%. On the basis of oversampling, the ordinary dual-channel data enhancement technology not only learns the information of the data itself, but also learns the information of the corresponding data after data enhancement, reducing the redundant shortcomings brought by the oversampling technology, improving the performance of the model, and reaching an accuracy of 84.3%. The image classification method of the embodiment of the present invention further fits the output feature distribution of the original image and the enhanced image, achieving an accuracy of nearly 90%.

[0067] Table 1

[0068] Using the Model Accuracy Avg-F value Undersampling method 0.722 0.704 Oversampling method 0.822 0.763 Ordinary dual-channel data enhancement 0.843 0.783 Image classification methods 0.892 0.81

[0069] Among them, the Avg-F value represents the micro-average F value, which is a global confusion matrix established by statistically analyzing each instance in the data set regardless of category, and then calculating the corresponding indicators. This value can more accurately reflect the performance of the experiment when the samples are unbalanced. Therefore, the image classification method of the embodiment of the present invention can alleviate the problem of severe data imbalance and achieve the effect of matching balanced samples or even higher performance.

[0070] Through the above steps S01 to S04, an embodiment of the present invention provides an image classification method, in which a preset image classification model is trained through an image classification model training method to obtain a trained image classification model, and the image to be classified and the image after data enhancement are input into the trained image classification model to obtain a final image classification result; thereby, accurate classification of images is achieved in the case of unbalanced data samples, and no human participation is required in the classification detection process, which reduces human factors and enables better application of detection results and judgment of image classification results, providing a highly efficient judgment reference.

[0071] The embodiment of the present invention also provides an image classification model training system, such as Figure 5 As shown, including:

[0072] The first processing module 1 is used to obtain an initial image sample and perform data enhancement on the initial image sample to obtain an enhanced image sample. For details, please refer to the relevant description of step S1 in the above method embodiment.

[0073] The second processing module 2 is used to input the initial image sample and the enhanced image sample into a preset image classification model to obtain the corresponding initial output features and enhanced output features. For details, please refer to the relevant description of step S2 in the above method embodiment.

[0074] The third processing module 3 is used to fit the enhanced output features with the initial output features using a feature-level loss function to obtain a feature fitting result. For details, please refer to the relevant description of step S3 in the above method embodiment.

[0075] The fourth processing module 4 is used to concatenate the initial output features and the enhanced output features to obtain a fused feature, and fit the fused feature to the preset label vector using a label-level loss function to obtain a label fitting result. For details, see the description of step S4 in the above method embodiment.

[0076] The fifth processing module 5 is used to update the parameters in the classification model according to the feature fitting results and the label fitting results until a trained image classification model is obtained. For details, please refer to the relevant description of step S5 in the above method embodiment.

[0077] Through the collaborative cooperation of the above-mentioned components, an embodiment of the present invention provides an image classification model training system, which performs data enhancement on initial image samples with known labels, inputs the initial image samples and the enhanced image samples into a preset image classification model to obtain corresponding initial output features and enhanced output features, and updates the parameters in the preset image classification model by performing feature-level fitting and label-level fitting on the output features, until a trained image classification model is obtained; the generalization ability of the image classification model is improved by performing data enhancement on the samples, so that balance can be maintained when the sample data set is severely unbalanced, and the classification performance of the model trained by the samples is guaranteed; through the dual-channel data enhancement technology, not only the information of the data itself is learned, but also the information of the corresponding data after data enhancement is learned, so that the classification model can obtain common knowledge of the original image and the image after data enhancement, reducing the redundancy disadvantages brought by the oversampling technology and improving the performance of the model.

[0078] The embodiment of the present invention also provides an image classification system, such as Figure 6 As shown, including:

[0079] The acquisition module 01 is used to acquire an image to be classified, perform data enhancement on the image to be classified, and obtain an enhanced image. For details, please refer to the description of step S01 in the above method embodiment.

[0080] The model training module 02 is used to train the preset image classification model using the image classification model training system to obtain a trained image classification model. For details, please refer to the relevant description of step S02 in the above method embodiment.

[0081] Feature extraction module 03 is used to input the image to be classified and the enhanced image into the trained image classification model, respectively, to obtain the corresponding initial output features and enhanced output features, and then to concatenate the initial output features and enhanced output features to obtain fused features. For details, see the description of step S03 in the above method embodiment.

[0082] The classification module 04 is used to classify the fused features using a preset classification function to obtain a classification result. For details, please refer to the description of step S04 in the above method embodiment.

[0083] Through the collaborative cooperation of the above-mentioned components, an embodiment of the present invention provides an image classification system, which trains a preset image classification model through an image classification model training method to obtain a trained image classification model, and inputs the image to be classified and the image after data enhancement into the trained image classification model to obtain the final image classification result; thereby achieving accurate classification of images in the case of unbalanced data samples, and no human participation is required in the classification detection process, which reduces human factors and can better apply the detection results and judge the image classification results, providing a highly efficient judgment reference.

[0084] The embodiment of the present invention further provides an electronic device, such as Figure 7 As shown, the electronic device may include a processor 901 and a memory 902, wherein the processor 901 and the memory 902 may be connected via a bus or other means. Figure 7 The bus connection is taken as an example.

[0085] The processor 901 may be a central processing unit (CPU). The processor 901 may also be other general-purpose processors, digital signal processors (DSP), application-specific integrated circuits (ASIC), field-programmable gate arrays (FPGA), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, or a combination of the above chips.

[0086] Memory 902, as a non-transitory computer-readable storage medium, can be used to store non-transitory software programs, non-transitory computer executable programs, and modules, such as the program instructions / modules corresponding to the methods in the embodiments of the present invention. Processor 901 executes the non-transitory software programs, instructions, and modules stored in memory 902 to perform various processor functions and data processing, thereby implementing the aforementioned methods.

[0087] The memory 902 may include a program storage area and a data storage area, wherein the program storage area may store an operating system and application programs required for at least one function; the data storage area may store data created by the processor 901, etc. In addition, the memory 902 may include a high-speed random access memory, and may also include a non-volatile memory, such as at least one disk storage device, a flash memory device, or other non-volatile solid-state storage device. In some embodiments, the memory 902 may optionally include a memory remotely located relative to the processor 901, and these remote memories may be connected to the processor 901 via a network. Examples of the above-mentioned network include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0088] One or more modules are stored in the memory 902 and, when executed by the processor 901 , perform the above method.

[0089] The specific details of the above electronic device can be understood by referring to the corresponding descriptions and effects in the above method embodiments, and will not be repeated here.

[0090] Those skilled in the art will appreciate that all or part of the processes in the above-mentioned embodiments can be implemented by instructing the relevant hardware through a computer program. The program can be stored in a computer-readable storage medium. When the program is executed, it can include the processes of the above-mentioned embodiments. The storage medium can be a magnetic disk, an optical disk, a read-only memory (ROM), a random access memory (RAM), a flash memory, a hard disk drive (HDD), or a solid-state drive (SSD). The storage medium can also include a combination of the above-mentioned types of memory.

[0091] The GPU used to implement the above-described method is a Tesla K40m graphics card. The GPU is one, Kepler GK110, with 2880 CUDA cores, 1.43 Tflops double-precision floating-point performance, 4.29 Tflops single-precision floating-point performance, 12 GB dedicated memory capacity, 288 GB / s memory bandwidth, and PCI-E 3.0 support. The power consumption is 235 W with passive heat dissipation. It should be noted that while the present embodiment illustrates the Tesla K40m as a GPU, other types of graphics cards can be used in practice, and the present invention is not limited thereto.

[0092] The above embodiments are only used to illustrate the technical solutions of the present invention rather than to limit the same. Although the present invention has been described in detail with reference to the above embodiments, those skilled in the art should understand that the specific implementation methods of the present invention can still be modified or replaced by equivalents. Any modification or equivalent replacement that does not depart from the spirit and scope of the present invention should be included in the scope of the claims of the present invention.

Claims

1. A method for training an image classification model, characterized in that: include: Acquire an initial image sample, and perform data enhancement on the initial image sample to obtain an enhanced image sample; Inputting the initial image sample and the enhanced image sample into a preset image classification model to obtain corresponding initial output features and enhanced output features; Fitting the enhanced output features to the initial output features using a feature-level loss function to obtain a feature fitting result; The initial output features and the enhanced output features are spliced to obtain fused features, and the fused features are fitted with a preset label vector using a label-level loss function to obtain a label fitting result; According to the feature fitting results and the label fitting results, the parameters in the preset image classification model are updated until a trained image classification model is obtained; Among them, the feature-level loss function is expressed by the following formula: Among them, Loss1 represents the feature-level loss function, MSE represents the Mean Square Error loss function, n represents the vector dimension of the output feature, r represents the initial output feature, r′ represents the enhanced output feature; r t represents the feature representation after the tth iteration; r′ t represents the enhanced output features after the tth iteration; The label-level loss function is expressed by the following formula: Among them, Loss2 represents the label level loss function, CE represents the Cross Entropy loss function, and y true (x i ) is the picture x i The output probability, y pred (x i ) is the picture x i The corresponding sample true label; y pred is the probability vector, y true is the sample formal label vector; The step of fitting the fusion feature with the preset label vector using the label-level loss function to obtain a label fitting result includes: Classify the fused features using a preset classification function to obtain a probability distribution of the classification results; Converting the probability distribution of the classification result into a probability vector, and fitting the probability vector with the preset label vector using a label-level loss function to obtain a label fitting result; The performing data enhancement on the initial image sample includes: The initial image is flipped horizontally or vertically, rotated 90 degrees, has one or two levels of brightness increased, has contrast increased, and has a changed hue.

2. The image classification model training method according to claim 1, characterized in that After obtaining the enhanced image samples, the image classification model training method further includes: performing standardization processing on the initial image samples and the enhanced image samples to obtain standard initial image samples and standard enhanced image samples.

3. An image classification method, characterized in that: include: Acquire an image to be classified, and perform data enhancement on the image to be classified to obtain an enhanced image; The image classification model training method according to any one of claims 1 to 2 is used to train a preset image classification model to obtain a trained image classification model; Inputting the image to be classified and the enhanced image into the trained image classification model respectively to obtain corresponding initial output features and enhanced output features respectively, and splicing the initial output features and the enhanced output features to obtain fusion features; The fusion features are classified using a preset classification function to obtain a classification result.

4. An image classification model training system, characterized in that: include: A first processing module is used to obtain an initial image sample and perform data enhancement on the initial image sample to obtain an enhanced image sample; A second processing module, configured to input the initial image sample and the enhanced image sample into a preset image classification model to obtain corresponding initial output features and enhanced output features; A third processing module is used to fit the enhanced output feature with the initial output feature using a feature-level loss function to obtain a feature fitting result; a fourth processing module, configured to concatenate the initial output features and the enhanced output features to obtain a fused feature, and fit the fused feature to a preset label vector using a label-level loss function to obtain a label fitting result; a fifth processing module, configured to update parameters in the classification model according to the feature fitting result and the label fitting result, until a trained image classification model is obtained; Among them, the feature-level loss function is expressed by the following formula: Among them, Loss1 represents the feature-level loss function, MSE represents the Mean Square Error loss function, n represents the vector dimension of the output feature, r represents the initial output feature, r′ represents the enhanced output feature; r t represents the feature representation after the tth iteration; r′ t represents the enhanced output features after the tth iteration; The label-level loss function is expressed by the following formula: Among them, Loss2 represents the label level loss function, CE represents the Cross Entropy loss function, and y true (x i ) is the picture x i The output probability, y pred (x i ) is the picture x i The corresponding sample true label; y pred is the probability vector, y true is the sample formal label vector; The step of fitting the fusion feature with the preset label vector using the label-level loss function to obtain a label fitting result includes: Classify the fused features using a preset classification function to obtain a probability distribution of the classification results; Converting the probability distribution of the classification result into a probability vector, and fitting the probability vector with the preset label vector using a label-level loss function to obtain a label fitting result; The performing data enhancement on the initial image sample includes: The initial image is flipped horizontally or vertically, rotated 90 degrees, has one or two levels of brightness increased, has contrast increased, and has a changed hue.

5. An image classification system, characterized in that: include: An acquisition module is used to acquire an image to be classified, perform data enhancement on the image to be classified, and obtain an enhanced image; A model training module, configured to train a preset image classification model using the image classification model training system according to claim 4 to obtain a trained image classification model; A feature extraction module is used to input the image to be classified and the enhanced image into the trained image classification model respectively, obtain corresponding initial output features and enhanced output features respectively, and splice the initial output features and the enhanced output features to obtain fusion features; The classification module is used to classify the fusion features using a preset classification function to obtain a classification result.

6. A computer-readable storage medium, characterized in that The computer-readable storage medium stores computer instructions, which, when executed by a processor, implement the image classification model training method according to any one of claims 1 to 2, or the image classification method according to claim 3.

7. An electronic device, characterized in that: include: A memory and a processor, wherein the memory and the processor are communicatively connected to each other, the memory stores computer instructions, and the processor executes the image classification model training method according to any one of claims 1 to 2, or executes the image classification method according to claim 3 by executing the computer instructions.

Images