Feature enhancement-based small sample bearing surface defect classification method and system

Abstract

The invention provides a small sample bearing surface defect classification method and system based on feature enhancement, and belongs to the field of bearing surface defect classification. The problems that a traditional image processing method is poor in robustness and adaptability; and a deep learning method needs a large amount of data for training. A global and local feature fusion layer is designed, global features are obtained by a feature embedding network layer, a local feature extraction layer is designed on the basis, the layer calculates a correlation matrix of each region and a neighborhood of the global features, and the correlation matrixes are fused in channel and space directions to obtain local feature information so as to enrich and enhance detail features; besides, a similar feature attention layer is designed, cosine similarity between different class prototype centers and different query sets is calculated to assist in obtaining similar areas between the two, high weight is given, classification is performed by comparing the similar areas between the query sets and the class prototype centers, and the influence of image background and noise on classification results is effectively overcome.

Description

Technical Field

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

[0002] Bearings are a basic component used in a large number of mechanical equipment. The quality of bearings will seriously affect the stability of equipment operation. In recent years, with the vigorous development of my country's manufacturing industry, the demand for bearing products in various industries has continued to increase, and the requirements for bearing quality have also increased. Although my country's mechanical processing technology has reached a high level, it is inevitable that certain damage will occur in the mass production of bearings. Therefore, bearings must undergo strict testing before leaving the factory to ensure that the bearings can be used normally after leaving the factory. At present, domestic bearing manufacturers have basically realized automated production and assembly line assembly, but the surface quality inspection after bearing assembly still relies on human visual inspection. This detection method has high labor intensity, low work efficiency, high labor cost, and it is difficult to achieve standardization in speed and accuracy. In addition, personal experience and subjective factors are strong, which will have a certain impact on the detection results and seriously affect the efficiency of the entire production line. In order to save labor costs and improve production efficiency, more and more companies use automated bearing surface defect methods instead of manual labor. At present, the use of machine vision technology to achieve bearing surface defect detection tasks mainly includes two aspects: traditional image processing methods and methods based on deep learning.

[0003] Traditional image processing methods analyze images based on a series of techniques based on traditional mathematics and engineering principles, and then use classifiers for detection. Zhang Chuangchuang uses grayscale transformation, bilateral filtering, least squares method, and ROI extraction and positioning to realize bearing end and side defect detection; Bin Liu proposed an innovative bearing surface defect detection method, which collects images under different lighting conditions, and combines polar coordinate transformation and improved OTSU threshold method to achieve accurate detection of pits and scratches on the surface of the bearing dust cover. However, traditional methods are not robust enough, lack adaptability, are only effective for specific categories, and cannot automatically adjust parameters to adapt to different detection tasks. At the same time, the calculation complexity is high and the real-time performance is poor.

[0004] With the rapid development of artificial intelligence, the research on intelligent detection algorithms for bearing surface defects has gradually attracted 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 FasterR-CNN algorithm, combined with a feature pyramid network and introduced deformable convolution, adaptively adjusted the receptive field to detect bearing surface defects; Yuan Tianle proposed automatic extraction of detection area preprocessing based on the YOLOv5 network structure, and improved the multi-head self-attention mechanism in Transformer, which improved the extraction ability of small and medium-sized surface defects, and had better defect positioning ability and high accuracy. Li Yadong et al. proposed a bearing surface small target defect detection algorithm with multi-attention feature weighted fusion, introduced the Res2Block module to increase the fine-grained features of the image in the backbone network, and designed embedded coordinate attention to enhance the cross-fusion ability of shallow detail features and deep advanced semantic features, and significantly improved the recognition rate of small targets. The surface defect detection method based on deep learning is based on a large amount of training experimental data, and the common features of the bearing image are extracted by training the convolutional neural network to realize 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 bearing surface defect detection tasks in actual industrial scenarios.

[0005] Under actual industrial conditions, with the continuous improvement of automated production technology and manufacturing processes, defects are not easy to occur during bearing production and assembly, which makes it difficult to collect defect data. At the same time, the location, size, and shape of defects are random, and the distribution of defect data generated is unbalanced. The scarcity of bearing defect samples has brought challenges to the training of intelligent detection networks. Summary of the invention

[0006] The technical problems to be solved by the present invention are:

[0007] In order to solve the problems in the existing bearing surface defect classification methods, the traditional image processing methods have poor robustness and weak adaptability; and the deep learning methods require a large amount of data for training.

[0008] The present invention adopts the following technical solutions to solve the above technical problems:

[0009] The present invention provides a small sample bearing surface defect classification method based on feature enhancement, which specifically comprises the following steps: including the fusion of global features and local features, designing similar feature attention and adaptive measurement of support set and query set,

[0010] 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.

[0011] 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;

[0012] 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.

[0013] Further, in step S100, it specifically includes mapping the preprocessed 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;

[0014] 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;

[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, 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

[0016] 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;

[0017] 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):

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

[0019] R=G☉D (2)

[0020] 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

[0021] 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.

[0022] 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;

[0023]

[0024]

[0025] 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.

[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 average value calculation method is used to obtain the class prototype center P of each support set category, as shown in formula (5):

[0027]

[0028] 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.

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

[0030] 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):

[0031]

[0032] 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;

[0033] 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;

[0034] 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 ;

[0035] By the class prototype similarity matrix Z s Get the class prototype attention map s a , by the query set similarity matrix Z qGet the query set attention graph q a .

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

[0037] 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):

[0038]

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

[0040] 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;

[0041] 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):

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

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

[0044] Through the above operations, the common similar feature areas between the class prototype center and the query set are obtained.

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

[0046] Adaptive metric network is used to classify similar features F in the query set qFor category prediction, the prediction loss is calculated as shown in formula (10):

[0047]

[0048] 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;

[0049] 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):

[0050]

[0051] 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 );

[0052]

[0053] 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;

[0054] 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):

[0055]

[0056] Among them, p(y=kQ 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):

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

[0058] Furthermore, 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; and a maximum pooling layer is introduced after the first two convolution blocks.

[0059] Furthermore, 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, scratches, notches, conformity, oxide scale and pitting.

[0060] The present invention provides a small sample bearing surface defect classification system based on feature enhancement. The system has a program module 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] The present invention provides a computer-readable storage medium, wherein the computer-readable storage medium stores a computer program, wherein the computer program is configured to implement the steps of a small sample bearing surface defect classification method based on feature enhancement when called by a processor.

[0062] Compared with the prior art, the present invention has the following beneficial effects:

[0063] The present invention discloses a method and system for classifying surface defects of small-sample bearings based on feature enhancement. The design of local feature extraction is to further extract local feature information based on global features, and to perform more detailed feature representation on the features. Similar features pay attention to the similarity calculation between the class prototype and the query set to determine the similar feature area between the two, so as to avoid the interference of background information and be 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, a certain degree of improvement has been achieved on both public data sets and self-made bearing defect data sets. At the same time, the present invention applies the small-sample learning method for the first time to solve the task of classifying bearing surface defects, solves the challenging task of classifying surface defects of small-sample bearings, overcomes the limited data availability, and provides an innovative method for classification tasks where data is scarce or difficult to obtain. BRIEF DESCRIPTION OF THE DRAWINGS

[0064] Figure 1 It is a 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 This is a schematic diagram of a feature embedding network structure in an embodiment of the present invention;

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

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

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

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

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

[0071] Figure 8 This is a comparison chart of the classification accuracy curves of the small sample image classification method in the simulation experiment of the present invention on the benchmark mini-ImageNet dataset after each round of training;

[0072] Fig. 9 This is a 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 dataset, tested after each round of training;

[0073] Fig.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] Fig.11 The thermal comparison diagram of the output features of the mini-ImageNet dataset after the baseline network and the final method in the ablation experiment of the present invention is disclosed;

[0075] Fig.12 Thermal comparison map of the output features after the baseline network and the final method for the self-made bearing surface defect dataset in the ablation experiment of the present invention. DETAILED DESCRIPTION

[0076] In order to make the above-mentioned objects, features and advantages of the present invention more obvious and easy to understand, specific embodiments of the present invention are described in detail below with reference to the accompanying drawings.

[0077] Specific implementation plan 1: Combine Figures 1 to 6 As shown, 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 measurement;

[0078] The present invention is improved on the basis of the prototype network of the metric learning framework, and 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, which calculates the correlation matrix of each region on the global feature and its neighborhood, and fuses the correlation matrix in the channel and spatial directions to obtain local feature information to enrich and enhance the detail features; a similar feature attention layer is designed to calculate the cosine similarity between the centers of different class prototypes and different query sets to assist in obtaining the similar areas between the two, and assign a higher weight, and classify by comparing the similar areas between the query set and the class prototype center, effectively overcoming the influence of image background and noise on the classification results;

[0079] Specifically include:

[0080] S100, mapping the input image data to the same feature space through a feature embedding network, including passing the preprocessed input image through the feature embedding network to obtain a global feature representation of the input support set and the query set in the feature space; wherein the data set is a public data set mini-Imagenet, and the data set preprocessing is to resize the input data size to a uniform size of 84*84;

[0081] Then, based on the input data features extracted by the common feature embedding module, a local feature extraction module is designed to further extract the local features of the sample, so as to extract as much feature information as possible under the limitation of few samples and shallow feature embedding network;

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

[0083] After data preprocessing, the support set and query set are sent to the feature embedding module to obtain the global feature representation matrix of the support set and 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;

[0084] In the meta-training and meta-testing process of small sample image classification, the number of support sets for each category is small, and generally 1 or 5 data are used to represent the characteristics of the category, which is not conducive to distinguishing between different categories. In order to improve the accuracy of small sample image classification, it is necessary to obtain more information that can represent the characteristics of the category under limited data. Local self-similar descriptors can represent the local characteristic information of the image, including the target skeleton, color, edge or other common local feature information. Inspired by the local self-similar descriptor, combined with Figure 3 As shown, the present invention designs a local feature extraction module to enhance the feature information of the input data and increase the discrimination of different categories, including:

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

[0086] Secondly, the global feature G is processed and its size after dimension upgrading is Make its dimension consistent with the neighborhood feature matrix D;

[0087] Next, perform Hadamard product calculation on each pixel in the dimension-enhanced global feature G and the neighborhood feature matrix D composed of 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):

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

[0089] R=G☉D (2)

[0090] 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

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

[0092] The fusion feature is obtained by combining global and local information, as shown in formula (3) and formula (4); the design of the local feature extraction module calculates the correlation matrix based on the global feature, and extracts the local structure, texture information, and color and intensity relationship of the input feature; combining the global and local feature information is more helpful to distinguish between different categories in the small sample image classification task;

[0093]

[0094]

[0095] 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;

[0096] After the support set data belonging to the same category are 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), which is used for subsequent similarity measurement:

[0097]

[0098] 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;

[0099] S200, design support set and query set similar features attention,

[0100] In the small sample image classification task based on metric learning, the extracted support set and query set features are the key factors affecting the classification effect. In order to mine rich feature information from a small number of image samples and improve the discrimination 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, focusing on the similar area between the query set and the class prototype center, and effectively distinguishing them by comparing the similarity scores between the two; combined with Figure 4The network structure of the similar feature attention module shown in the figure, the prototype center part: 1. The cosine similarity between the prototype center feature and the query set feature is calculated; the prototype similarity matrix Z is obtained s and query similarity matrix Z q ; 2. Similar feature attention layer: Design a convolutional network and use the class prototype similarity map and query similarity map to calculate the class prototype attention map and query set attention map respectively;

[0101] 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):

[0102]

[0103] 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;

[0104] Combination Figure 4 As shown, the cosine similarity is used to calculate 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;

[0105] S220, design similar feature attention layer, the similar feature attention layer aims to obtain the similarity matrix Z of the query set from the class prototype s and the query set about the class prototype similarity matrix Z q The similar areas between the two were found in Z s Able to generate prototype-like attention maps for a specific query set a , while Z q Able to generate a query set attention graph q for a specific class prototype a ;Attention maps highlight features that are critical to different classification tasks;

[0106] By the class prototype similarity matrix Z s Get the class prototype attention map s a Methods such as Figure 5 As shown, the query set similarity matrix Z q Get the query set attention graph q a The method is also calculated according to this structure;

[0107] Specifically include:

[0108] 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):

[0109]

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

[0111] 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;

[0112] S223, use the SoftMax 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):

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

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

[0115] Through the above operations, we can obtain the common similar feature areas between the prototype center and the query set, focus on the similar features of the two for classification, and avoid the interference of background information of global features on classification accuracy; at the same time, the similar feature areas and similarity scores between the support set and the query set of the same category are higher than the areas and similarity scores between the support set and the query set of different categories;

[0116] S300, similarity measurement calculation,

[0117] Combination Figure 6 As shown, referring to the relational network concept, the convolutional neural network predicts the query set features, wherein the input is the predicted category feature, the output is the probability value predicted for each category, and the maximum probability is the predicted category type; different from the fixed metric function, the present invention designs a nonlinear metric function, namely, an adaptive metric network; the function adaptively learns a metric network that is more conducive to classification based on the information provided by the input data, including,

[0118] 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):

[0119]

[0120] Where: 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;

[0121] To ensure the effectiveness of the similarity feature attention module, the attention maps obtained between the class prototype and the query set are made more similar. The category prediction is performed by calculating the similarity score between the similar features of the query set and the similar features of each class prototype to ensure the effectiveness of the similarity feature attention layer;

[0122] 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 (12) is as follows:

[0123]

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

[0125]

[0126] Among them, a represents the ath query set, k represents the category of the class prototype, and V 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;

[0127] 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):

[0128]

[0129] Where 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; U 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):

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

[0131] Specific implementation scheme 2: The present invention provides a small sample bearing surface defect classification system based on feature enhancement. The system has a program module corresponding to the above steps, and executes the steps in the above-mentioned small sample bearing surface defect classification method based on feature enhancement during operation.

[0132] The other combinations and connection relationships of this embodiment are the same as those of the first embodiment.

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

[0134] The other combinations and connection relationships of this embodiment are the same as those of the first embodiment.

[0135] Simulation experiment

[0136] Experimental dataset:

[0137] The data sets used in the experiment 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. The data set contains a total of 100 categories, of which 64 categories are used for training, 16 categories are used for verification, and 20 categories are used for testing. Each category contains 600 images. All input image data are adjusted to 84*84 size during the training and testing stages. The effectiveness of the experimental method of the present invention is evaluated by comparing the classification effect of the existing small sample image classification method on mini-ImageNet.

[0138] Combination Figure 7 As shown, in order to verify the performance of the method of the present invention under real factory conditions, we constructed a bearing surface defect data set. All defect data were manually discovered on the actual production line and data collection was performed. The surface defect data set mainly includes seven common defects in actual industrial scenarios, namely: grooves, red rust, abrasions, notches, qualified, oxide scale and pitting, and each category contains 50 images. The collected bearing surface defect data does not participate in training, but is used as a test set to verify the model effect. The defect data is also uniformly adjusted to 84*84 size for testing.

[0139] Experimental setup:

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

[0141] The training and testing phases of the model of the present invention are both carried out under the setting of the small sample learning experimental method based on N-wayK-shot. The model is trained using a cross-task mechanism, and a 5-way 1-shot or 5-way 5-shot setting is adopted, that is, each meta-task contains a total of 5 categories, and each category contains 1 or 5 support set images and 15 query set images. In the model training phase, each epoch includes a total of 2000 tasks, and 100 epochs are performed, including a total of 200,000 meta-tasks. In the model testing phase, the 5-way1-shot and 5-way 5-shot experimental settings are also used for verification. Among them, 15 query set images are selected for verification in each meta-task, and the average classification accuracy obtained from 3000 test tasks within a 95% confidence interval is used as an evaluation indicator.

[0142] Parameter settings: The same parameter settings are used for the training set, validation set, and test set. 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, the weight decay is 0.0005, and the rest of the parameters use the default values.

[0143] Experimental results:

[0144] To verify the effectiveness of the algorithm proposed in this paper, this paper uses a small sample public dataset mini-ImageNet to compare with other small sample image classification methods, where the feature embedding network 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 min i-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, the classification accuracy of the method proposed in the present invention is higher than that of the classic meta-learning method (MAML) and the classic metric learning method (MatchingNetwork, ProtoNetwork, Relation Network), whether in the 5-way 1-shot or 5-way 5-shot experimental setting. Secondly, compared with the metric methods (MATANet, MSLPN, TRNA, FMAF) that perform multi-scale feature fusion or design attention mechanism networks in the feature embedding stage, the method proposed in the present invention is also better than them. Thirdly, 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 significant improvement in classification accuracy. Finally, compared with the RENet and CAN methods that focus on the common similar features between the support set and the query set for classification, the classification accuracy of the method proposed in the present invention is 1.04% and 2.23% higher than that of RENet under 5-way 1-shot or 5-way 5-shot conditions, respectively. Compared with the CAN method, the classification accuracy of the method proposed in the present invention is improved by 4.01% and 6.23% respectively.

[0149] At the same time, during the training process of different small sample image classification methods, after each round of training, in order to test the effect of the network model training in that round, the updated weights of the current network model are used to perform classification tests on the categories that did not participate in the training. A total of 600 test tasks were performed, and the average classification accuracy obtained from the 600 test tasks was taken as the average test accuracy after the end of that round. Figure 8 The average accuracy curve of the classification test after each round of training task during the model training process for different small sample image classification methods shown in the figure. Figure 8 It can be clearly seen from the curve trend that with the increase of training rounds, the accuracy of model classification gradually improves and eventually stabilizes. It proves that the test accuracy of the small sample image classification method proposed in the present invention on the benchmark mini-ImageNet dataset is significantly higher than that of other small sample image classification methods compared.

[0150] The method proposed in the present invention achieves good classification results on public data sets. Since the data categories used in the training and testing stages have no intersection, it shows strong generalization ability. Since the amount of homemade bearing surface defect data is small, each category contains only 50 image data, and the direct use of convolutional neural networks for training and testing results in poor model generalization ability. Therefore, the present invention migrates the network model trained using the mini-ImageNet public data set to the homemade bearing surface defect data set for classification. The classification method learned in the public data set is applied to the classification of bearing surface defect data with a small amount of data. No training is required, and the differences between the features of different categories are directly used for classification, which solves the problem of poor model generalization caused by the small number of bearing defect data sets.

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

[0152] Table 2 Experimental results of different small sample image classification methods on bearing surface defect data

[0153]

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

[0155] Combination Fig. 9The accuracy bar graph of the classification test of self-made bearing surface defect data by different small sample image classification methods is shown in the figure. The test phase is carried out under the 5-way 5shot experimental setting, and the different small sample image classification network models are trained using the public dataset mini-ImageNet. 15 query sets are selected for each category for testing. Each epoch contains 400 test tasks. A total of 5 epochs are carried out. The average classification accuracy of the 400 test tasks in each epoch is calculated. The purpose of using more test tasks for testing is to ensure the accuracy of the classification results. Fig. 9 It can be seen that by comparing the classification test accuracy of different small sample image classification methods in each epoch, the method proposed in the present invention is superior to other solutions.

[0156] In order to further explore the classification effect of the model, the present invention uses a confusion matrix to display the prediction results of different defects, such as Fig.10 At the same time, common image classification indicators, precision (P), recall (R) and F1 score (F1), are used to compare the classification effects of different small sample models.

[0157] Through the confusion matrix, it is found that the defect features distinguish widely-used oxide scale, pitting and qualified data, which are easy to classify, but the defects such as grooves, red rust, scratches and notches with small defect features and large background information are easy to misjudge. The reason for the analysis is that the bearing background is relatively consistent and the defect features are small. When similar features are extracted, it is very likely to focus on the consistent background information between the two, and use the background information as the main feature of the category for classification, which leads to misjudgment.

[0158] It can be seen from Tables 3 and 4 that the model proposed in the present invention has the best prediction effect on groove and scratch defects compared with the other five small sample classification methods.

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

[0160]

[0161]

[0162] Table 4 Comparison of evaluation results of small sample image classification models

[0163]

[0164] Ablation experiment

[0165] In order to further verify the rationality and effectiveness of the method proposed in the present invention, the present invention conducts ablation experiments on the local feature extraction module, similar feature attention module and 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. The prototype network mainly includes two parts: feature embedding network and metric function. Among them, the feature embedding network is a commonly used four-layer convolution layer, and the Euclidean distance is selected as the metric function.

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

[0167] First, training and testing are performed on the baseline network (Base); secondly, the local feature extraction module and the similarity feature attention module are 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 ablation experiment results in Table 5, it can be found that on the self-made bearing defect dataset, under the experimental setting of 5N 5shot, compared with the classification accuracy of the benchmark network, after the local feature extraction module is introduced, the classification accuracy is improved by 1.82%; after the similar feature attention module is introduced, the classification accuracy is improved by 6.71%; after adding the two modules to the benchmark 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 category classification prediction. The experimental network includes four convolutional layers, local feature extraction and similar feature attention modules, and finally selects Euclidean distance, cosine similarity and adaptive metric network to predict the category. Compared with Euclidean distance, the adaptive metric network proposed in the present invention improves classification accuracy by 5.59%; compared with cosine similarity, the classification accuracy is improved by 12.27%. Ablation experiments prove that the adaptive metric network proposed in 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] In order to intuitively prove that the similar feature attention module proposed in 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 prototype center and query set features extracted from the baseline network (Base) and the prototype center and query set features after the present invention introduces local feature extraction and similar feature attention module (Base+LFE+SFA) on the baseline network.

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

[0177] By comparing the visualization results of the baseline network (Base) and the method of the present invention (Base+LFE+SFA), it can be seen that on the baseline network, whether on the public mini-ImageNet dataset or on the self-made bearing surface defect dataset, the baseline 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 feature information area of ​​interest, and no attention to the key information in the image that can distinguish this category from other categories, and contains too much non-category background information, which affects the classification results. On the other hand, the similar feature attention module proposed in the present invention calculates the cosine similarity between the features of the prototype centers of different classes and the query set. Compared with the baseline network, after extracting the prototype center features and the query set features, the method of the present invention pays more attention to the similar feature areas between the two, and assigns a higher weight to the area. The query set is classified by comparing the similarity with the similar feature areas of the prototype centers of different classes, which is more helpful to distinguish between different categories and improve the accuracy of classification.

[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 may make various changes and modifications without departing from the spirit and scope of the present invention, and these changes and modifications will 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