A Small-Sample Bearing Surface Defect Classification Method and System Based on Feature Enhancement

Abstract

The present invention provides a method and system for small-sample bearing surface defect classification based on feature enhancement, belonging to the field of bearing surface defect classification. To solve the problems of poor robustness and weak self-adaptability of traditional image processing methods; and the need for a large amount of data for training in deep learning methods. The present invention designs a global and local feature fusion layer. The global feature is obtained by a feature embedding network layer. On this basis, a local feature extraction layer is designed. This layer calculates the correlation matrix of each region on the global feature and its neighborhood, and at the same time fuses the correlation matrix in the channel and spatial directions to obtain local feature information to enrich and enhance the detailed features. In addition, a similar feature attention layer is designed to calculate the cosine similarity between different class prototype centers and different query sets to assist in obtaining the similar regions between the two, and assigns higher weights, and classifies by comparing the similar regions between the query set and the class prototype center, effectively overcoming the influence of image background and noise on the classification result.

Description

Technical Field

[0001] The present invention relates to the technical field of bearing surface defect classification. Specifically, it relates to a small-sample bearing surface defect classification method and system based on feature enhancement. Background Art

[0002] As a basic component widely used in mechanical equipment, the quality of bearings will seriously affect the stability of equipment operation. In recent years, with the vigorous development of China's manufacturing industry, the demand for bearing products in various industries has been increasing continuously, and at the same time, the requirements for bearing quality have also been rising day by day. Although China's machining technology has reached a relatively high level, certain damages are inevitable in the mass production of bearings. Therefore, before leaving the factory, bearings must undergo strict inspection to ensure that they can be used normally after leaving the factory. At present, domestic bearing manufacturers have basically achieved automated production and assembly line assembly, but the surface quality inspection after bearing assembly still relies on human visual inspection. This inspection method has a large labor intensity, low work efficiency, high labor costs, and it is difficult to achieve standardization in terms of speed and accuracy. Moreover, personal experience and subjective factors are strong, which will have a certain impact on the inspection results and seriously affect the efficiency of the entire production line. To save labor costs and improve production efficiency, more and more enterprises use automated bearing surface defect methods to replace manual labor. At present, using machine vision technology to achieve bearing surface defect detection tasks mainly includes two aspects: traditional image processing methods and deep learning-based methods.

[0003] Traditional image processing methods analyze images based on a series of technologies of traditional mathematics and engineering principles, and then use classifiers for detection. Zhang Chuangchuang used gray-scale transformation, bilateral filtering, least squares method, and ROI extraction and positioning to achieve defect detection on the end face and side face of bearings; Bin Liu proposed an innovative bearing surface defect detection method. This method collects images under different lighting conditions at different angles and combines polar coordinate transformation and improved OTSU threshold method to achieve accurate detection of pits and scratches on the bearing dust cover surface. However, traditional methods have insufficient robustness, lack adaptability, are only effective for specific categories, cannot automatically adjust parameters to adapt to different detection tasks, and at the same time have a high computational complexity and poor real-time performance.

[0004] With the rapid development of artificial intelligence, the research on intelligent detection algorithms for bearing surface defects has gradually received attention, and more and more deep learning algorithms have begun to be applied in the field of bearing surface defect detection. Lan Yeshen proposed an improved Faster R-CNN algorithm, combined with a feature pyramid network and introduced deformable convolution to adaptively adjust the receptive field for bearing surface defect detection; Yuan Tianle proposed automatic extraction of the detection area preprocessing based on the YOLOv5 network structure, and improved the multi-head self-attention mechanism in the Transformer, enhancing the extraction ability of small and medium-sized surface defects, with better defect localization ability and higher accuracy. Li Yadong et al. proposed a small target defect detection algorithm for bearing surfaces with multi-attention feature weighted fusion, introduced the Res2Block module that increases image fine-grained features in the backbone network, and designed embedded coordinate attention to enhance the cross-fusion ability of shallow detail features and deep high-level semantic features, significantly improving the recognition rate of small targets. The surface defect detection method based on deep learning extracts the common features of bearing images by training a convolutional neural network on the basis of a large amount of training experimental data to achieve the detection task of bearing surface defects before leaving the factory. Although the defect detection method based on deep learning has a certain generalization ability and can adapt to detection tasks in different scenarios, the training of convolutional neural networks requires a large amount of defect data and is not suitable for the bearing surface defect detection task in the actual industrial scenario.

[0005] Under actual industrial conditions, with the continuous improvement of automation production technology and manufacturing processes, it is not easy to generate defects during the production and assembly of bearings, resulting in difficulties in collecting defect data. At the same time, the location, size, and shape of the generated defects are random, and the distribution of the generated defect data is unbalanced. The scarcity of bearing defect samples poses challenges to the training of intelligent detection networks. Summary of the Invention

[0006] The technical problem to be solved by the present invention is:

[0007] To solve the problems in the existing bearing surface defect classification methods, such as the poor robustness and weak self-adaptability of traditional image processing methods; and the need for a large amount of data for training in deep learning methods.

[0008] The technical solution adopted by the present invention to solve the above technical problems:

[0009] The present invention provides a small-sample bearing surface defect classification method based on feature enhancement, specifically including the following steps: global feature and local feature fusion, design of similar feature attention for the support set and query set, and adaptive metric.

[0010] S100. Design a global and local feature fusion layer. The global features are obtained from the bearing surface defect image dataset through a feature embedding network layer. On this basis, design a local feature extraction layer, which calculates the correlation matrix of each region on the global features with its neighborhood, and fuses the correlation matrix in both the channel and spatial directions to obtain local feature information;

[0011] S200. Utilize the correlation between the support set and the query set to design a similar feature attention module, obtain the class prototype center of the support set and the common similar feature regions between the two, and compare the similarity scores between the two to avoid the interference of background information existing in the global features on the classification accuracy;

[0012] S300. Design an adaptive metric network, including predicting the category of the similar features of the query set, and predicting the similarity scores of the similar features of the query set and the similar features of each class prototype, and then calculating the loss of the similar features of the class prototype and the predicted similar features of the query set and the category prediction loss to update the network, and complete the classification according to the similarity metric.

[0013] Further, in step S100, specifically, map the preprocessed image data at the input end to the same feature space through a feature embedding network to obtain the global feature representations of the input support set and query set in the feature space;

[0014] S110. Send the support set and the query set into the feature embedding module after data preprocessing to obtain the global feature representation matrices of the support set and the query set of the input image data where C represents the number of global feature channels, and H and W are the height and width of the global features respectively;

[0015] S120. Design a local feature extraction module to enhance the feature information of the input data, including the local feature extraction module processing the global features of the obtained support set and query set, and extracting local features through a sliding window on each channel; set the sliding window size to M×N, slide it sequentially on the global features, and each pixel point in the spatial dimension obtains a local block of size M×N around it. Therefore, there are C local blocks of the same size M×N on the same channel, and the local feature block obtained by one pixel point is The local features extracted for each pixel point on all channels together form a neighborhood feature matrix

[0016] Dimensionality increase the global features G of the support set and the query set, and the size after dimensionality increase is Make its dimension consistent with that of the neighborhood feature matrix D;

[0017] Perform the Hadamard product calculation on each pixel point in the upsampled global feature G and the neighborhood feature matrix D composed of local feature blocks around each pixel point to obtain the feature correlation matrix R. The calculation process for a certain pixel point is shown in Equation (1), and the calculation formulas for the global feature G and the neighborhood feature matrix D are shown in Equation (2):

[0018] R(x) = G(x) ⊙ D(x) (1)

[0019] R = G ⊙ D (2)

[0020] Where G(x) ∈ G, with a size of 1*1, x represents a certain pixel point in the global feature G, and x ∈ [1, H] × [1, W]; D(x) ∈ D, with a size of M × N, representing the local feature block around the pixel point x; the dot product represents the Hadamard product, R(x) ∈ R, with a size of M × N, representing the correlation between the pixel point and its surrounding features; the Hadamard products of pixel points G and the neighborhood feature matrix D in all dimensions together form the final correlation matrix R, and

[0021] Design a four-layer convolutional neural network to extract features from the correlation matrix R to obtain the local features of the support set and the query set;

[0022] S130. Combine the global and local information to obtain the fused feature, as shown in Formulas (3) and (4); the design of the local feature extraction module is based on the global feature, calculates the correlation matrix, and extracts the local structure, texture information, and color and intensity relationships of the input features; combining the global and local feature information is more helpful for distinguishing between different categories in the small-sample image classification task;

[0023]

[0024]

[0025] Where S represents the fused feature obtained by combining the global and local features of the support set, and Q represents the fused feature obtained by combining the global and local features of the query set; G s represents the global feature of the support set output by the feature embedding network; G q represents the global feature of the query set output by the feature embedding network; L s represents the local graph feature of the support set output by the feature embedding network; L q represents the local feature of the query set output by the feature embedding network.

[0026] Furthermore, after the support set data belonging to the same category is processed by the feature embedding network and the local feature extraction network, the class prototype center P of each support set category is obtained using the average value calculation method, as shown in Equation (5):

[0027]

[0028] Among them, P Ι represents the class prototype center belonging to the support set of category Ι, K represents the number contained in each support set, and S Ι represents the global and local fusion features of the support set of category Ι.

[0029] Furthermore, in step S200, it specifically includes

[0030] S210. Similarity calculation: By calculating the cosine distance between the support set class prototype center feature and the query set feature, the correlation score between different feature vectors is obtained, as shown in formula (6):

[0031]

[0032] Among them, Z ij represents the cosine similarity between two feature vectors. i and j respectively represent the local feature vectors of the class prototype center and the local feature vectors of the query set, and both have H×W local feature vectors; s i represents the i-th local feature vector of the class prototype center; q j represents the j-th local feature vector in the query set feature;

[0033] Use cosine similarity calculation to find the cosine similarity between each local vector s i of the class prototype and each local vector q j in the query set feature, and obtain the similarity matrix Z s of this class prototype with respect to this query set; Similarly, use the transposed matrix of the query set feature and the class prototype center feature representation for cosine calculation to obtain the similarity matrix Z q of the query set with respect to each class prototype center, where and M = H×W;

[0034] S220. Design a similar feature attention layer to find the similar regions between the similarity matrix Z s of the class prototype with respect to the query set and the similarity matrix Z q of the query set with respect to the class prototype; Z s can generate a class prototype attention map s a for a specific query set, and Z q can generate a query set attention map q a for a specific class prototype;

[0035] Obtain the class prototype attention map s s from the class prototype similarity matrix Z a , and obtain the query set attention map q qObtain the query set attention map q a 。

[0036] Furthermore, in step S220, it specifically includes

[0037] S221. For the class prototype similarity matrix Z s Adopt global average pooling operation in each channel to calculate the average similarity R between each local area in the class prototype and a single local area in the query set s , and convert the class prototype similarity matrix into a region similarity of size M×1, as shown in formula (7):

[0038]

[0039] Where represents the average similarity between the overall class prototype and the i-th local feature in the query set;

[0040] S222. After obtaining the region similarity through step S221, obtain the similarity fusion coefficient through the similar feature attention layer, multiply the correlation scores in the class prototype similarity matrix by the average similarity of the class prototype with respect to a single region in the query set in the similarity fusion coefficient in turn, and sum to obtain the average similarity of each region in the class prototype with respect to the global feature of the query set;

[0041] S223. Use the classifier to process the average similarity between each region in the class prototype and the global feature of the query set, assign a larger weight to the region with higher similarity to highlight the region feature, and obtain s a ; Based on the class prototype fusion feature, enhance the similar region feature through similar feature attention, and use the final class prototype similar feature for classification; The method for determining the query set similar feature is the same, and finally obtain the similar features F s and F q , as shown in formula (8) and formula (9):

[0042] F s =S + S×s a (8)

[0043] F q =Q + Q×q a (9)

[0044] Through the above operations, obtain the common similar feature region between the class prototype center and the query set.

[0045] Furthermore, in step S300, it specifically includes

[0046] Use the adaptive metric network for the query set similar feature F qPerform class prediction, and its prediction loss is calculated as shown in formula (10):

[0047]

[0048] where U is the number of test query sets, and C is the number of classes; y ij is the true label that the i-th input data belongs to the j-th class. If the true label of the i-th input is the j-th class, then y ij = 1, otherwise y ij = 0; p ij is the probability that the model predicts the i-th input as the j-th class, and log is the natural logarithm;

[0049] Calculate the similarity probability between the similar features of the query set and the similar features of the class prototype through cosine similarity. The class with the highest score is the predicted class; the cosine similarity calculation formula is shown in formula (11), and the similarity probability calculation formula is shown in formula (12):

[0050]

[0051] where d() represents the cosine similarity calculation function, and A and B are two vectors in the n-dimensional space. A = (x 11 , x 12 ,..., x 1n ), B = (x 21 , x 22 ,..., x 2n );

[0052]

[0053] where a represents the a-th query set, k represents the class of the class prototype, N represents the number of classes in the support set, represents the similar features of the k-th class prototype with respect to the a-th query set, represents the similar features of the a-th query set with respect to the k-th class prototype, and at the same time, the class label of this query set is consistent with the class prototype; through the above formula, calculate the probability of the similar features of the class prototype that is consistent with the class label of the query set a in the class prototype similarity matrix obtained for the query set a;

[0054] At the same time, calculate the loss predicted by the similar features of the class prototype and the similar features of the query set, as shown in formula (13):

[0055]

[0056] where p(y = kQ a ∈ k) is the query set a in the class prototype similarity matrix Among them, the similar feature of the class prototype that is consistent with the query set a category probability; M is the total number of query sets; this loss function can ensure that the similar feature attention module strengthens the attention to the similar features shared between the class prototype and the query set; the final loss function is calculated according to formula (14):

[0057] loss = loss1 + loss2 (14).

[0058] Furthermore, the feature embedding network is a Conv64 structure, including 4 convolutional blocks; each convolutional block includes an activation function layer, a batch normalization layer, and a convolutional layer; a max pooling layer is introduced after the first two convolutional blocks.

[0059] Furthermore, the input data set is the small sample image classification public data set mini-ImageNet and the bearing surface defect image data set under actual factory production conditions; the bearing surface defect image data set includes bearing images with grooves, red rust, scratches, indentations, qualified, mill scale, and pitting.

[0060] A small sample bearing surface defect classification system based on feature enhancement according to the present invention, the system has program modules corresponding to the above steps, and executes the steps in the above-mentioned small sample bearing surface defect classification method based on feature enhancement when running.

[0061] A computer-readable storage medium according to the present invention, the computer-readable storage medium stores a computer program, and the computer program is configured to implement the steps of the small sample bearing surface defect classification method based on feature enhancement when called by a processor.

[0062] Compared with the prior art, the beneficial effects of the present invention are:

[0063] A small sample bearing surface defect classification method and system based on feature enhancement according to the present invention, the design of local feature extraction further extracts local feature information on the basis of global features, and performs a more detailed feature representation on this feature; similar feature attention determines the similar feature region between the two through the similarity calculation between the class prototype and the query set, avoiding the interference of background information, and is more helpful for classification between similar categories; the adaptive metric network has better generalization than the fixed metric distance function. Compared with the existing small sample image classification methods, certain improvements have been achieved on both the public data set and the self-made bearing defect data set. At the same time, the present invention first applies the small sample learning method to solve the bearing surface defect classification task, solves the challenging task of few-sample bearing surface defect classification, overcomes the limited data availability, and provides an innovative method for classification tasks with scarce or difficult-to-obtain data. Description of the Drawings

[0064] Figure 1 It is the structural block diagram of a small-sample bearing surface defect classification method based on feature enhancement in an embodiment of the present invention;

[0065] Figure 2 It is the schematic diagram of the feature embedding network structure in an embodiment of the present invention;

[0066] Figure 3 It is the schematic diagram of the local feature extraction module structure in an embodiment of the present invention;

[0067] Figure 4 It is the schematic diagram of the similar feature attention structure in an embodiment of the present invention Figure 1 ;

[0068] Figure 5 It is the schematic diagram of the similar feature attention structure in an embodiment of the present invention Figure 2 ;

[0069] Figure 6 It is the structure diagram of the similarity measurement network in an embodiment of the present invention;

[0070] Figure 7 It is the physical diagram of the self-made bearing surface defect data set in the simulation experiment of the present invention;

[0071] Figure 8 It is the comparison chart of the classification accuracy curve of the small-sample image classification method in the simulation experiment of the present invention on the benchmark mini-ImageNet data set after each round of training;

[0072] Figure 9 It is the bar chart of the classification accuracy of the small-sample image classification method in the simulation experiment of the present invention on the benchmark mini-ImageNet data set after each round of training;

[0073] Figure 10 It is the confusion matrix of the classification of the self-made bearing surface defect data set in the simulation experiment of the present invention;

[0074] Figure 11 It is the heat map comparison of the output features of the publicly available mini-ImageNet data set in the ablation experiment of the present invention after the benchmark network and the final method;

[0075] Figure 12 It is the heat map comparison of the output features of the self-made bearing surface defect data set in the ablation experiment of the present invention after the benchmark network and the final method. Detailed implementation manners

[0076] To make the above objects, features and advantages of the present invention more obvious and understandable, the following detailed description of the specific embodiments of the present invention will be given with reference to the accompanying drawings.

[0077] Specific implementation method 1: Combine Figures 1 to 6 As shown in Figures 1 to 6 , the present invention provides a small-sample bearing surface defect classification method based on feature enhancement, including global feature and local feature fusion, similar feature attention of support set and query set, and adaptive metric;

[0078] The present invention is an improvement based on the prototype network in the metric learning framework. A global and local feature fusion layer is designed. The global feature is obtained by the feature embedding network layer. On this basis, a local feature extraction layer is designed. This layer calculates the correlation matrix of each region on the global feature and its neighborhood, and at the same time fuses the correlation matrix in the channel and spatial directions to obtain local feature information to enrich and enhance the detailed features; A similar feature attention layer is designed to calculate the cosine similarity between different class prototype centers and different query sets to assist in obtaining the similar regions between the two, and assign higher weights, and classify by comparing the similar regions between the query set and the class prototype center, effectively overcoming the influence of image background and noise on the classification result;

[0079] Specifically, it includes:

[0080] S100. Map the image data at the input end to the same feature space through the feature embedding network, including passing the preprocessed input end image through the feature embedding network to obtain the global feature representations of the input support set and query set in the feature space; wherein, the data set is the public data set mini-Imagenet, and the data set preprocessing is to uniformly resize the input data size to 84*84;

[0081] Then, on the basis of extracting the input data features by using the common feature embedding module, a local feature extraction module is designed to further extract the local features of the samples, so as to extract as much feature information as possible under the limitations of few samples and shallow feature embedding networks;

[0082] Combine Figure 2 As shown in Figure 2 , the feature embedding network in this step can select the common Conv64 structure, including 4 convolutional blocks; each convolutional block consists of an activation function layer, a batch normalization layer and a convolutional layer; at the same time, to reduce overfitting, a max pooling layer is introduced after the first two convolutional blocks;

[0083] After preprocessing the support set and query set, send them into the feature embedding module, and the global feature representation matrices of the support set and query set of the input image data are where C represents the number of global feature channels, and H and W are the height and width of the global feature respectively;

[0084] During the meta-training and meta-testing processes of few-shot image classification, the number of support sets for each category is small. Generally, 1 or 5 data are used to represent the features of that category, which is not conducive to the distinction between different categories. To improve the accuracy of few-shot image classification, more information that can represent the features of that category needs to be obtained under the condition of limited data. The local self-similarity descriptor can represent the local feature information of an image, including the object skeleton, color, edges, or other common local feature information. Inspired by the local self-similarity descriptor, combined with Figure 3 as shown in, the present invention designs a local feature extraction module to enhance the feature information of the input data and increase the distinguishability between different categories, including

[0085] First, the local feature extraction module processes the obtained class prototypes and the global features of the query set. The unfold function is used to extract local patches from the multi-dimensional global features through a sliding window. The size of the sliding window is set to M×N, and it slides sequentially on the global features. In the spatial dimension, each pixel point obtains local patches of size M×N around it. Therefore, C local patches of the same size M×N are obtained on the same channel, and the local feature patches obtained by one pixel point are The local features extracted for each pixel point on all channels together form a neighborhood feature matrix where C represents the number of channels, H and W represent the height and width of the global feature respectively, and M and N are the width and height of the sliding window respectively;

[0086] Second, the global feature G is processed, and its size after dimension elevation is to make its dimension consistent with that of the neighborhood feature matrix D;

[0087] Third, the Hadamard product calculation is performed between each pixel point in the dimension-elevated global feature G and the neighborhood feature matrix D composed of the local feature patches around each pixel point to obtain a feature correlation matrix R. The calculation process of a certain pixel point is shown in Equation (1), and the calculation formula of the global feature G and the neighborhood feature matrix D is shown in Equation (2):

[0088] R(x) = G(x) ⊙ D(x) (1)

[0089] R = G ⊙ D (2)

[0090] where G(x) ∈ G, with a size of 1*1, x represents a certain pixel point in the global feature G, and x ∈ [1, H] × [1, W]; D(x) ∈ D, with a size of M×N, representing the local feature patch around the pixel point x; the dot product represents the Hadamard product, R(x) ∈ R, with a size of M×N, representing the correlation between the pixel point and its surrounding features; the Hadamard products of the pixel points G and the neighborhood feature matrix D in all dimensions together form the final correlation matrix R, and

[0091] Finally, in order to better integrate the features of the channel and spatial dimensions, a four-layer convolutional neural network is designed to extract features from the correlation matrix R to obtain the local features of the support set and the query set;

[0092] The fused features are obtained by combining the global and local information, as shown in Formulas (3) and (4); the design of the local feature extraction module is based on the global features, calculates the correlation matrix, and extracts the local structure, texture information, and color and intensity relationships of the input features; combining the global and local feature information is more helpful for distinguishing between different categories in the few-shot image classification task;

[0093]

[0094]

[0095] Among them, S represents the fused feature obtained by combining the global and local features of the support set, and Q represents the fused feature obtained by combining the global and local features of the query set; G s represents the global feature of the support set output by the feature embedding network; G q represents the global feature of the query set output by the feature embedding network; L s represents the local graph feature of the support set output by the feature embedding network; L q represents the local feature of the query set output by the feature embedding network;

[0096] After the support set data belonging to the same category is processed by the feature embedding network and the local feature extraction network, the class prototype center P of each support set category is obtained by using the average value calculation method, as shown in Formula (5), for subsequent similarity measurement:

[0097]

[0098] Among them, P Ι represents the class prototype center of the support set belonging to category Ι, K represents the number contained in each support set, and S Ι represents the global and local fused feature of the support set of category Ι;

[0099] S200. Design the similar features of the support set and the query set. Note that

[0100] In the few-shot image classification task based on metric learning, the features of the support set and the query set extracted are key factors affecting the classification effect. To mine rich feature information from a small number of image samples and improve the discriminability within and between classes, the present invention uses the correlation between the support set and the query set to design a similar feature attention module, focuses on the similar regions between the query set and the class prototype center, and makes effective distinctions by comparing the similarity scores between the two; Combine Figure 4The network structure of the similarity feature attention module shown, the class prototype center part: 1. Calculate the cosine similarity between the class prototype center feature and the query set feature; obtain the class prototype similarity matrix Z s and the query similarity matrix Z q ; 2. Similarity feature attention layer: Design a convolutional network, and use the class prototype similarity map and the query similarity map to calculate the class prototype attention map and the query set attention map respectively;

[0101] S210. Similarity calculation, by calculating the cosine distance between the support set class prototype center feature and the query set feature, obtain the correlation score between different feature vectors, as shown in formula (6):

[0102]

[0103] where Z ij represents the cosine similarity between two feature vectors, i and j respectively represent the local feature vectors of the class prototype center and the query set, and both have H×W local feature vectors; s i represents the i-th local feature vector of the class prototype center; q j represents the j-th local feature vector in the query set feature;

[0104] Combined Figure 4 shown, use cosine similarity calculation to find the cosine similarity between each local vector s i of the class prototype and each local vector q j in the query set feature, and obtain the similarity matrix Z of this class prototype with respect to this query set s ; Similarly, use the transposed matrix of the query set feature and the class prototype center feature representation for cosine calculation to obtain the similarity matrix Z of the query set with respect to each class prototype center q , where and M = H×W;

[0105] S220. Design a similarity feature attention layer. The similarity feature attention layer aims to find the similar regions between the class prototype similarity matrix Z s with respect to the query set and the query set similarity matrix Z q with respect to the class prototype; Z s can generate the class prototype attention map s a for a specific query set, while Z q can generate the query set attention map q a for a specific class prototype; The attention map highlights the features crucial for different category classification tasks;

[0106] The method to obtain the class prototype attention map s s from the class prototype similarity matrix Z a is asFigure 5 As shown, from the query set similarity matrix Z q obtain the query set attention map q a The method is also calculated according to this structure;

[0107] Specifically,

[0108] S221. For the class prototype similarity matrix Z s Adopt global average pooling operation in each channel, and calculate the average similarity R between each local area in the class prototype and a single local area in the query set s , and convert the class prototype similarity matrix into a region similarity of size M×1, as shown in formula (7):

[0109]

[0110] where, represents the average similarity between the overall class prototype and the i-th local feature in the query set;

[0111] S222. After obtaining the region similarity through step S221, obtain the similarity fusion coefficient through the similar feature attention layer, multiply the correlation scores in the class prototype similarity matrix by the average similarity of the class prototype with respect to a single region in the query set in the similarity fusion coefficient in turn, and sum to obtain the average similarity of each region in the class prototype with respect to the global feature of the query set;

[0112] S223. Use the SoftMax classifier to process the average similarity between each region in the class prototype and the global feature of the query set, assign a larger weight to the region with higher similarity, so as to highlight the region feature, and obtain s a ; On the basis of the class prototype fusion feature, enhance the similar region feature through similar feature attention, and use the final class prototype similar feature for classification; The method for determining the query set similar feature is the same, and finally obtain the similar features F s and F q , as shown in formula (8) and formula (9):

[0113] F s =S + S×s a (8)

[0114] F q =Q + Q×q a (9)

[0115] Through the above operations, the common similar feature regions between the class prototype center and the query set can be obtained. Focus on the similar features of the two for classification to avoid the interference of background information existing in the global features on the classification accuracy. At the same time, the similar feature regions and similarity scores between the support set and the query set of the same category should be higher than those between the support set and the query set of different categories.

[0116] S300, Similarity metric calculation,

[0117] Combined with Figure 6 As shown, referring to the idea of the relational network, the convolutional neural network predicts the features of the query set. Among them, the input end is the predicted category features, and the output is the probability value predicted for each category. The category type with the highest probability is the predicted category. Different from using a fixed metric function, the present invention designs a non-linear metric function, that is, an adaptive metric network. This function adaptively learns a metric network that is more conducive to classification according to the information provided by the input data, including,

[0118] Use the adaptive metric network to perform category prediction on the similar features F of the query set q The prediction loss calculation is as shown in formula (10):

[0119]

[0120] Among them: U is the number of test query sets, C is the number of categories; y ij is the true label that the i-th input data belongs to the j-th class. If the true label of the i-th input is the j-th class, then y ij = 1, otherwise y ij = 0; p ij is the probability that the model predicts the i-th input as the j-th class, and log is the natural logarithm;

[0121] To ensure the effectiveness of the similar feature attention module and make the attention map obtained between the class prototype and the query set more similar, category prediction is performed by calculating the similarity scores between the similar features of the query set and the similar features of each class prototype to ensure the effectiveness of the similar feature attention layer;

[0122] Calculate the similarity probability between the similar features of the query set and the similar features of the class prototype through cosine similarity. The category with the highest score is the predicted category. The cosine similarity calculation formula is as shown in formula (11), and the similarity probability calculation formula (12) is as follows:

[0123]

[0124] Among them, d() represents the cosine similarity calculation function, A and B are two vectors in the n-dimensional space, A = (x 11 , x 12 ,..., x1n ), B = (x 21 , x 22 ,..., x 2n );

[0125]

[0126] Among them, a represents the a-th query set, k represents the category of the class prototype, V represents the number of categories of the support set, represents the similar features of the k-th class prototype with respect to the a-th query set, represents the similar features of the a-th query set with respect to the k-th class prototype, and at the same time, the category label of this query set is consistent with the class prototype; through the above formula, the similar feature matrix of the class prototypes obtained for the query set a in multiple query sets a is calculated Among them, the probability of the similar features of the class prototype consistent with the category of the query set a ;

[0127] At the same time, calculate the loss predicted by the similar features of the class prototype and the similar features of the query set, as shown in Equation (13):

[0128]

[0129] Among them, p(y = k|Q a ∈ k) is the probability of the similar features of the class prototype consistent with the category of the query set a in the similar feature matrix of the class prototypes obtained for the query set a in multiple query sets Among them, the probability of the similar features of the class prototype consistent with the category of the query set a ; U is the total number of query sets; this loss function can ensure that the similar feature attention module strengthens the attention to the similar features shared between the class prototype and the query set; the final loss function is calculated according to formula (14):

[0130] loss = loss1 + loss2 (14).

[0131] Specific implementation plan two: A small-sample bearing surface defect classification system based on feature enhancement according to the present invention, which system has program modules corresponding to the above steps, and when running, executes the steps in the above-mentioned small-sample bearing surface defect classification method based on feature enhancement.

[0132] Other combinations and connection relationships in this implementation plan are the same as those in the specific implementation plan one.

[0133] Specific implementation plan three: A computer-readable storage medium according to the present invention, where the computer-readable storage medium stores a computer program, and the computer program is configured to implement the steps of the small-sample bearing surface defect classification method based on feature enhancement when called by a processor.

[0134] Other combinations and connection relationships in this implementation scheme are the same as those in the first specific implementation scheme.

[0135] Simulation experiment

[0136] Experimental data set:

[0137] The data sets used in the experiments of the present invention include the commonly used small-sample image classification public data set mini-ImageNet and the self-made bearing surface defect data set under actual factory production conditions. Mini-ImageNet is the most widely used benchmark data set in the field of small-sample learning. This data set contains a total of 100 categories, of which 64 categories are used for training, 16 categories are used for validation, and 20 categories are used for testing. And each category contains 600 images. During the training and testing phases, all input image data is adjusted to a size of 84*84. By comparing the classification effects of existing small-sample image classification methods on mini-ImageNet, the effectiveness of the experimental method of the present invention is evaluated.

[0138] Combined with Figure 7 As shown, to verify the performance of the method of the present invention under actual factory conditions, we constructed a bearing surface defect data set. All defect data was discovered manually on the actual production line and data collection was carried out. The surface defect data set mainly includes seven common defects in actual industrial scenarios, namely: groove, red rust, scratch, notch, qualified, mill scale and pitting. Each category contains 50 images. The collected bearing surface defect data does not participate in training and is used as the test set to verify the model effect. Also, the defect data is uniformly adjusted to a size of 84*84 for testing.

[0139] Experimental settings:

[0140] The experiments of the present invention are implemented based on the open-source deep learning framework Pytorch on the Ubuntu 18.04 system, where the Torch version is 1.8.0, the corresponding Torchvison version is 0.8.0, the processor is the GPU NVIDIA GeForce RTX3090, and the running memory is 48GB.

[0141] Both the training and testing phases of the model of the present invention are carried out based on the experimental method of few-shot learning with N-way K-shot. The model is trained using a cross-task mechanism, with settings of 5-way 1-shot or 5-way 5-shot, that is, each meta-task contains 5 categories, and each category contains 1 or 5 support set images and 15 query set images. Each epoch in the model training phase includes 2,000 tasks, and 100 epochs are carried out, including a total of 200,000 meta-tasks. In the model testing phase, the experimental settings of 5-way 1-shot and 5-way 5-shot are also used for verification. Among them, 15 query set images are selected for verification in each meta-task. Within a 95% confidence interval, the average classification accuracy obtained from 3,000 test tasks is used as the evaluation index.

[0142] Parameter settings: The training set, validation set, and test set all use the same parameter settings. All experiments use the SGD optimization algorithm to update the model parameters. The initial learning rate is set to 0.1. After the 50th epoch, the learning rate is halved every 10 epochs. The momentum is 0.9, and the weight decay is 0.0005. The remaining parameters use the default values.

[0143] Experimental results:

[0144] To verify the effectiveness of the algorithm proposed by the present invention, the present invention uses the small sample public dataset mini-ImageNet to compare with other small sample image classification methods. Among them, the feature embedding network all uses Conv4 (four-layer convolution) to ensure the fairness of the experimental results. The classification effects of different methods on the public dataset are shown in Table 1.

[0145] Table 1 Comparison of mini-ImageNet experimental results

[0146]

[0147]

[0148] First, from the analysis of the experimental results in Table 1, under the condition of the same feature embedding layer, compared with the classical meta-learning method (MAML) and the classical metric learning methods (MatchingNetwork, ProtoNetwork, Relation Network), the classification accuracy of the method proposed in the present invention is relatively high in both the 5-way 1-shot and 5-way 5-shot experimental settings. Second, compared with the metric methods (MATANet, MSLPN, TRNA, FMAF) that perform multi-scale feature fusion in the feature embedding stage or design attention mechanism networks, the method proposed in the present invention is also superior to them. Third, compared with the improved meta-learning methods MELR and BaseTransformers in the past two years, the method proposed in the present invention has also achieved a significant improvement in classification accuracy. Finally, compared with the RENet and CAN methods that classify by focusing on the common similar features between the support set and the query set, the classification accuracy of the method proposed in the present invention is increased by 1.04% and 2.23% respectively compared with RENet under the conditions of 5-way 1-shot or 5-way 5-shot. Compared with the CAN method, the classification accuracy of the method proposed in the present invention is increased by 4.01% and 6.23% respectively.

[0149] Meanwhile, during the training process of different few-shot image classification methods, after each round of task training is completed, in order to test the training effect of the network model in this round, the updated weights of the current network model are used to classify and test the categories that have not participated in the training. A total of 600 test tasks are carried out, and the average classification accuracy obtained from the 600 test tasks is averaged as the average test accuracy after this round of training is completed. Combining Figure 8 with the average accuracy curves of different few-shot image classification methods during the model training process, where classification tests are carried out after each round of training task is completed. Observing Figure 8 the curve trend in it, it can be clearly seen that as the number of training rounds increases, the classification accuracy of the model gradually improves and finally stabilizes. It is proved that the test accuracy of the few-shot image classification method proposed in the present invention on the benchmark mini-ImageNet dataset is significantly higher than that of other compared few-shot image classification methods.

[0150] The method proposed in the present invention achieves good classification results on the public dataset. Since there is no intersection between the data categories used in the training and testing phases, it demonstrates strong generalization ability. Due to the small amount of self-made bearing surface defect data, each category only contains 50 image data. Directly using a convolutional neural network for training and testing results in poor model generalization ability. Therefore, the present invention will transfer the network model trained on the mini-ImageNet public dataset to the self-made bearing surface defect dataset for classification. Applying the classification method learned on the public dataset to the classification of bearing surface defect data with a small amount of data, without going through training, directly using the differences between the features of different categories for classification, solves the problem of poor model generalization caused by the small number of bearing defect datasets.

[0151] As can be seen from the classification results of Table 1 on the public dataset mini-ImageNet, under the experimental method setting of 5way 5shot, the classification effect is better. Analyzing the reasons: Compared with using only one support set feature as the class prototype center feature of this category, using multiple support set features to jointly determine the class prototype center feature can better show the common features of this category and is helpful for distinguishing between different categories. Therefore, to improve the accuracy of bearing surface defect classification, the present invention selects the experimental method setting of 5way5shot for meta-testing.

[0152] Table 2 Experimental results of different few-shot image classification methods on bearing surface defect data

[0153]

[0154] From the analysis of Table 2, when the model trained using the mini-ImageNet public dataset for meta-training is transferred to the meta-testing phase for bearing surface defect data testing, compared with other classic few-shot image classification methods based on metric learning (Proto Network, Relation Network, DN4), the method proposed in the present invention achieves good classification results. Compared with the RENet and CAN methods that focus on classifying by the common similar features between the support set and the query set, the classification accuracy of the method proposed in the present invention is increased by 1.07% and 1.69% respectively compared with RENet and CAN under the 5-way 5-shot condition. It also proves the effectiveness of the method proposed in the present invention for bearing surface defect classification.

[0155] Combined with Figure 9The bar chart of the accuracy of classifying the self-made bearing surface defect data by different few-shot image classification methods shown below. The test phase was carried out under the experimental setting of 5-way 5-shot, and different few-shot image classification network models were all trained using the public dataset mini-ImageNet. 15 query sets were selected for testing in each category. Each epoch contained 400 test tasks, and a total of 5 epochs were carried out. The average classification accuracy of the 400 test tasks in each epoch was calculated. The purpose of using more test tasks for testing was to ensure the accuracy of the classification results. As can be seen from Figure 9 it, when comparing the classification test accuracies of different few-shot image classification methods in each epoch, the method proposed by the present invention is better than other schemes.

[0156] To further explore the classification effect of the model, the present invention uses a confusion matrix to show the prediction results of different defects, as shown in Figure 10 the figure below. At the same time, common image classification metrics, precision (P), recall (R), and F1 score (F1), were used to compare the classification effects of different few-shot models.

[0157] It was found through the confusion matrix that it is easy to classify the mill scale, pitting, and qualified data with widely distinguishable defect features, but it is easy to misjudge the defects of grooves, red rust, scratches, and indentations with smaller defect features and larger background information. The reason for the analysis is that the bearing background is relatively consistent and the defect features are small. When extracting similar features, it is very likely to pay attention to the consistent background information between the two, and classify using the background information as the main feature of this category, thus resulting in misjudgment.

[0158] As can be seen from Tables 3 and 4, the model proposed by the present invention has the best prediction effect on the defects of grooves and scratches compared with the other 5 few-shot classification methods for comparison.

[0159] Table 3 Evaluation results of the model of the present invention

[0160]

[0161]

[0162] Table 4 Comparison of evaluation results of few-shot image classification models

[0163]

[0164] Ablation experiment

[0165] To further verify the rationality and effectiveness of the method proposed in the present invention, ablation experiments are carried out on the local feature extraction module, the similar feature attention module and the adaptive metric network proposed in the network model structure. The method of the present invention is improved on the basis of the framework structure of the prototype network, and the prototype network mainly includes two parts: a feature embedding network and a metric function. Among them, the feature embedding network is a commonly used four-layer convolutional layer, and the Euclidean distance is selected as the metric function.

[0166] To ensure the rigor of the ablation experiment, during the ablation experiment on local feature extraction (LFE) and similar feature attention (SFA), the present invention changes the Euclidean distance of the metric function in the prototype network to an adaptive metric network, that is, the baseline network (Base) consists of four-layer convolutional layer and an adaptive metric network.

[0167] First, training and testing are carried out on the baseline network (Base); secondly, the local feature extraction module and the similar feature attention module are respectively introduced on the basis of the baseline network to form Base+LFE and Base+SFA; finally, the above two modules are added to the baseline network at the same time to form the final method of the present invention, Base+LFE+SFA.

[0168] Table 5 Ablation experiment of the method of the present invention on bearing surface defect data

[0169]

[0170] By observing the results of the ablation experiment in Table 5, it can be found that on the self-made bearing defect data set, under the experimental setting of 5N 5-shot, compared with the classification accuracy of the baseline network, after introducing the local feature extraction module, the classification accuracy is improved by 1.82%; after introducing the similar feature attention module, the classification accuracy is improved by 6.71%; after adding the two modules to the baseline network at the same time, the classification accuracy is improved by 8.60%. The ablation experiment proves that the method for classifying small-sample bearing surface defects designed by the present invention is effective.

[0171] Secondly, the present invention conducts ablation experiments on different metric functions for final class classification prediction. The experimental network includes four-layer convolutional layer, local feature extraction and similar feature attention module, and finally the Euclidean distance, cosine similarity and adaptive metric network are respectively selected to predict the class. The adaptive metric network proposed by the present invention has a 5.59% higher classification accuracy than the Euclidean distance; and a 12.27% higher classification accuracy than the cosine similarity. The ablation experiment proves that the adaptive metric network proposed by the present invention is effective.

[0172] Table 6 Ablation experiment of the method of the present invention on bearing surface defect data

[0173]

[0174] Experimental results:

[0175] To visually prove that the proposed similar feature attention module of the present invention focuses on the similar features between the support set and the query set, the present invention uses the Grad-CAM method to visualize the class prototype center and the features of the query set extracted from the base network (Base) respectively, as well as the class prototype center and the query set features after introducing local feature extraction and similar feature attention module (Base+LFE+SFA) on the base network by the method of the present invention.

[0176] The present invention selects 3 different categories of data from mini-ImageNet, and also selects 4 different categories of data from the self-made bearing surface defect dataset. Each category contains a support set and a query set. The features output by the network constructed by the base network and the method proposed by the present invention for each pair of data are visualized. The feature visualization heat map is as Figure 11 shown in Fig. 12.

[0177] By comparing the visualization results of the base network (Base) and the method of the present invention (Base+LFE+SFA), it can be seen that on the base network, whether on the publicly available mini-ImageNet dataset or on the self-made bearing surface defect dataset, the base network only completes the operation of converting the input image data into a feature matrix, without further requirements and restrictions on the features to be extracted, resulting in no regularity in the area of the feature information concerned, and not noticing the key information in the image that can distinguish this category from other categories, including too much background information that is not of the category, which affects the classification result. On the contrary, the proposed similar feature attention module of the present invention calculates the cosine similarity of the features between different class prototype centers and the query set. Compared with the base network, after the method of the present invention extracts the class prototype center features and the query set features, it pays more attention to the similar feature area between the two, assigns a higher weight to this area, and the query set is classified by comparing the similarity with the similar feature areas of different class prototype centers, which is more conducive to distinguishing between different categories and improving the classification accuracy.

[0178] Although the present invention is disclosed as above, the protection scope of the present invention is not limited thereto. Those skilled in the art of the present invention can make various changes and modifications without departing from the spirit and scope of the present invention, and these changes and modifications will all fall within the protection scope of the present invention.

Claims

1. A small sample bearing surface defect classification method based on feature enhancement, characterized in that: The specific steps are as follows: Including the fusion of global and local features, designing similar feature attention and adaptive measurement of support set and query set, S100, design a global and local feature fusion layer. The global features are obtained from the bearing surface defect image dataset through the feature embedding network layer. On this basis, a local feature extraction layer is designed. This layer calculates the correlation matrix between each region and its neighborhood on the global features, and fuses the correlation matrix in the channel and spatial directions to obtain local feature information. S200, using the correlation between the support set and the query set, design a similar feature attention module, obtain the common similar feature area between the class prototype center of the support set and the query set, and compare the similarity scores between the two to avoid the interference of background information existing in the global features on the classification accuracy; S300, designing an adaptive metric network, including predicting categories for similar features of the query set, and predicting categories for similarity scores between similar features of the query set and similar features of each class prototype, and then calculating the prediction loss of similar features of the class prototype and similar features of the query set and the category prediction loss to update the network, and completing classification according to the similarity metric.

2. The method for classifying surface defects of small sample bearings based on feature enhancement according to claim 1 is characterized in that: In step S100, it specifically includes mapping the pre-processed image data at the input end to the same feature space through a feature embedding network to obtain a global feature representation of the input support set and the query set in the feature space; S110, after data preprocessing, the support set and the query set are sent to the feature embedding module to obtain the global feature representation matrix of the support set and the query set of the input image data Among them, C represents the number of global feature channels, H and W are the height and width of the global feature respectively; S120, design a local feature extraction module to enhance the feature information of the input data, including the local feature extraction module processing the global features of the obtained support set and query set, extracting local features through a sliding window on each channel; setting the sliding window size to M×N, sliding on the global features in sequence, and each pixel point in the spatial dimension obtains a local block of M×N size around it, so there are C local blocks on the same channel, each of which is M×N in size, and the local feature block obtained by a pixel point is The local features extracted from each pixel on all channels together form the neighborhood feature matrix The global feature G of the support set and query set is upgraded to a size of Make its dimension consistent with the neighborhood feature matrix D; The Hadamard product is calculated for each pixel in the global feature G after dimensionality increase and the neighborhood feature matrix D composed of the local feature blocks around each pixel to obtain the feature correlation matrix R. The calculation process of a certain pixel is shown in formula (1). The calculation formula of the global feature G and the domain feature matrix D is shown in formula (2): R(x)=G(x)☉D(x) (1) R=G☉D (2) Among them, G(x)∈G, size is 1*1, x represents a pixel point in the global feature G, x∈[1,H]×[1,W]; D(x)∈D, size is M×N, represents the local feature block around the pixel point x; the dot product represents the Hadamard product, R(x)∈R, size is M×N, represents the correlation between the pixel point and its surrounding features; the Hadamard product of the pixel point G and the neighborhood feature matrix D in all dimensions together constitutes the final correlation matrix R, and A four-layer convolutional neural network is designed to extract features from the correlation matrix R and obtain local features of the support set and query set. S130, combining global and local information to obtain fusion features, as shown in formula (3) and formula (4); the design of the local feature extraction module calculates the correlation matrix based on the global features, and extracts the local structure, texture information, and color and intensity relationship of the input features; combining the global and local feature information is more helpful to distinguish between different categories in the small sample image classification task; Among them, S represents the fusion feature obtained by combining the global features and local features of the support set, Q represents the fusion feature obtained by combining the global features and local features of the query set; G s G represents the global features of the output support set of the feature embedding network; q represents the global features of the query set output by the feature embedding network; L s L represents the local graph features of the feature embedding network output support set; q Represents the local features of the query set output by the feature embedding network.

3. The method for classifying surface defects of small sample bearings based on feature enhancement according to claim 2 is characterized in that: After the support set data belonging to the same category is processed by the feature embedding network and the local feature extraction network, the average value calculation method is used to obtain the class prototype center P of each support set category, as shown in formula (5): Among them, P Ι represents the class prototype center belonging to the class support set of category Ι, K represents the number contained in each support set, S Ι Represents the global and local fusion features of the class support set.

4. The method for classifying surface defects of small sample bearings based on feature enhancement according to claim 3 is characterized in that: In step S200, specifically including: S210, similarity calculation, by calculating the cosine distance between the central feature of the support set class prototype and the query set feature, the correlation score between different feature vectors is obtained, as shown in formula (6): Among them, Z ij represents the cosine similarity between two feature vectors, i and j represent the local feature vector of the class prototype center and the local feature vector of the query set, and both have H×W local feature vectors; s i Represents the i-th local eigenvector of the class prototype center; q j Represents the jth local feature vector in the query set features; Use cosine similarity calculation to find each local vector s of the class prototype i With each local vector q in the query set feature j The cosine similarity between them is used to obtain the similarity matrix Z of the prototype of this class with respect to the query set. s Similarly, the transposed matrix of the query set features is used to perform cosine calculations with the class prototype center feature representation to obtain the similarity matrix Z of the query set with respect to each class prototype center. q ,in, And M = H × W; S220, design similar feature attention layer, used to obtain similarity matrix Z from class prototype to query set s and the query set about the class prototype similarity matrix Z q The similar areas between the two were found in Z s Can generate prototype-like attention maps s for a specific query set a , Z q Able to generate a query set attention graph q for a specific class prototype a ; By the class prototype similarity matrix Z s Get the class prototype attention map s a , by the query set similarity matrix Z q Get the query set attention graph q a .

5. The method for classifying surface defects of small sample bearings based on feature enhancement according to claim 4 is characterized in that: In step S220, specifically including: S221, similarity matrix Z for class prototypes s A global average pooling operation is performed in each channel to calculate the average similarity R between each local region in the class prototype and a single local region in the query set. s , convert the class prototype similarity matrix into the regional similarity of size M×1, as shown in formula (7): in, Represents the average similarity between the overall class prototype and the i-th local feature in the query set; S222, after obtaining the regional similarity through step S221, the similarity fusion coefficient is obtained after the similarity feature attention layer, and the correlation score in the class prototype similarity matrix is ​​multiplied with the average similarity of the class prototype with respect to a single region of the query set in the similarity fusion coefficient in turn, and the sum is obtained to obtain the average similarity of each region in the class prototype with respect to the global feature of the query set; S223, use the classifier to process the average similarity between each region in the class prototype and the global features of the query set, and give a larger weight to the region with higher similarity to highlight the characteristics of the region, and obtain s a ; Based on the fusion features of the class prototype, similar features are noticed to enhance the similar area features, and the final similar features of the class prototype are used for classification; the method for determining the similar features of the query set is the same, and finally the similar features F of the class prototype and the query set are obtained s and F q , as shown in formula (8) and formula (9): F s =S+S×s a (8) F q =Q+Q×q a (9) Through the above operations, the common similar feature areas between the class prototype center and the query set are obtained.

6. The method for classifying surface defects of small sample bearings based on feature enhancement according to claim 5 is characterized in that: In step S300, specifically including: Adaptive metric network is used to classify similar features F in the query set q For category prediction, the prediction loss is calculated as shown in formula (10): Among them, U is the number of test query sets, C is the number of categories; y ij The i-th input data is the true label of the j-th category. If the true label of the i-th input is the j-th category, then y ij =1, otherwise y ij =0; p ij is the probability that the model predicts that the i-th input is of the j-th class, and log is the natural logarithm; The similarity probability between similar features of the query set and similar features of the class prototype is calculated by cosine similarity, and the one with the highest score is the predicted category; the cosine similarity calculation formula is shown in formula (11), and the similarity probability calculation formula is shown in formula (12): Where d() represents the cosine similarity calculation function, A and B are two vectors in n-dimensional space, A=(x 11 ,x 12 ,...,x 1n ), B=(x 21 ,x 22 ,...,x 2n ); Among them, a represents the ath query set, k represents the category of the class prototype, and N represents the number of categories in the support set. represents the similar features of the k-th class prototype with respect to the a-th query set, Represents the similarity features of the a-th query set with respect to the k-th class prototype, and the query set category label is consistent with the class prototype; the above formula is used to calculate the similarity matrix of the query set a in multiple class prototypes obtained with respect to the query set a. The prototype-like features that are consistent with the query set a category probability; At the same time, the loss of prediction of similar features of the class prototype and similar features of the query set is calculated, as shown in formula (13): Among them, p(y=k|Q a ∈k) is the similarity matrix of the prototype obtained by query set a in multiple query set a The prototype-like features that are consistent with the query set a category The probability of; M is the total number of query sets; this loss function ensures that the similar feature attention module strengthens the attention to the similar features shared between the class prototype and the query set; the final loss function is calculated according to formula (14): loss=loss1+loss2 (14).

7. The method for classifying surface defects of small sample bearings based on feature enhancement according to claim 1, characterized in that: The feature embedding network is a Conv64 structure, including 4 convolution blocks; each convolution block includes an activation function layer, a batch normalization layer and a convolution layer; a maximum pooling layer is introduced after the first two convolution blocks.

8. The method for classifying surface defects of small sample bearings based on feature enhancement according to claim 1 is characterized in that: The input dataset is a small sample image classification public dataset mini-ImageNet and a bearing surface defect image dataset under actual factory production conditions; the bearing surface defect image dataset includes bearing images with grooves, red rust, abrasions, notches, conformity, oxide scale and pitting.

9. A small sample bearing surface defect classification system based on feature enhancement, characterized by: The system has a program module corresponding to the steps of any one of claims 1 to 8, and executes the steps in the above-mentioned small sample bearing surface defect classification method based on feature enhancement when running.

10. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores a computer program, and the computer program is configured to implement the steps of the small sample bearing surface defect classification method based on feature enhancement according to any one of claims 1 to 8 when called by a processor.

Images