CN115953386A - A detection method for surface defects of lightweight gears based on MSTA-YOLOv5 - Google Patents

Abstract

The invention relates to the technical field of product defect detection, and discloses a light-weight gear surface defect detection method based on MSTA-YOLOv5, which comprises the following steps: firstly, acquiring a gear surface defect image, marking and dividing the image, and constructing a gear surface defect data set; then constructing an MSTA-YOLOv5 detection model, and training the MSTA-YOLOv5 detection model based on the gear surface defect data set; and finally, sending the gear defect image to be detected into a trained MSTA-YOLOv5 detection model to obtain the defect type of the detected gear. The invention solves the problems that the computing resource demand is too large, the memory consumption is serious, the cost is higher, an enterprise needs a low-delay model, and a mobile device terminal needs a fast and accurate small model, realizes the detection and automatic sorting of the surface defects of the gear, and can improve the detection efficiency of the surface defects of the gear.

Description

A detection method for surface defects of lightweight gears based on MSTA-YOLOv5

technical field

The invention relates to the technical field of product defect detection, in particular to a method for detecting surface defects of lightweight gears based on MSTA-YOLOv5.

Background technique

With the development of science and technology and changes in social needs, large multi-faceted and complex structural workpieces are more popular in industrial production. Gears are widely used transmission components in the machinery industry, and their quality is particularly important in production. However, in the actual production process, due to the influence of factors such as process flow, production equipment and on-site environment, various defects appear on the surface of the gear. Benefits drop. Therefore, it is necessary to inspect the surface of the gear, and the traditional manual inspection has a large workload, which may easily cause visual fatigue of the inspectors, resulting in missed inspections and false inspections.

In recent years, with the rapid development of machine vision technology, inspection technology based on machine vision has been applied in product surface quality inspection. Most of the current gear defect detection technology adopts digital image processing technology, but the processing method and algorithm of this technology are single, and it is difficult to effectively extract defect targets for gears with high surface complexity, resulting in poor detection results. ideal.

For example, the patent document with the publication number CN115187820A discloses lightweight target detection methods, devices, equipment, and storage media. ShuffleNetv is used as the feature extraction module in the YOLOv4 network structure, but the amount of parameters and calculations are very large, and the feature extraction module The SE attention mechanism is used in the method, which has the disadvantage of insufficient precision;

For example, the patent document with the publication number CN112990325A discloses a light-weight network construction method for embedded real-time visual target detection, using the CBAM attention mechanism, which has the advantage of light weight, but with the increase in light weight, the loss of accuracy is large. Moreover, the Focus slicing operation is adopted in this technical solution, which increases the amount of parameters and weakens the advantages of light weight;

For example, the patent document whose publication number is CN114898171A discloses a real-time target detection method suitable for embedded platforms. Although the effect of light weight is achieved, the accuracy still causes a relatively large loss.

With the development of artificial intelligence technology, deep learning methods are widely used in the fields of image processing and workpiece quality inspection because of their excellent performance advantages in processing industrial images with complex backgrounds and weak defects. Using the deep learning method, it can accurately identify and segment the gear surface defects semantically, reducing the interference of background and other factors, thus effectively improving the detection accuracy. Although there are a lot of researches on improving different target detection networks and detecting defects in industrial products, which have achieved considerable results, there is no research on the models that require small size and fewer calculation parameters in the enterprise. Such models are in the cost budget. Devices with low computing power and relatively insufficient computing power can also achieve good detection speed and accuracy.

Contents of the invention

The deep learning method has greatly improved the accuracy in the field of image classification, but the current target detection algorithm based on deep learning requires too much computing resources and serious memory consumption, which makes the cost high. Aiming at the above problems and the problem that enterprises need low-latency models and mobile device terminals need fast and accurate small models, the present invention provides a light-weight gear surface defect detection method based on MSTA-YOLOv5, which realizes the detection and automatic detection of gear surface defects. Sorting can improve the detection efficiency of gear surface defect detection.

The technical scheme that the present invention solves technical problem is:

A light-weight gear surface defect detection method based on MSTA-YOLOv5, comprising the following steps: firstly obtain a gear surface defect image, and mark and divide the image to construct a gear surface defect data set; then construct MSTA-YOLOv5 detection The model is based on the gear surface defect data set to train the MSTA-YOLOv5 detection model; finally, the gear defect image to be detected is sent to the trained MSTA-YOLOv5 detection model to obtain the defect type of the detected gear.

The MSTA-YOLOv5 detection model includes:

Input part: Input the gear surface defect image into the MSTA-YOLOv5 network for adaptive anchor frame calculation and Mosaic9 data enhancement;

Backbone part: the feature extraction backbone network adopts the ShuffleNetv2 architecture, including sequentially connected CBRM operations, the first downsampling layer, the second convolutional normalization layer, the second downsampling layer, the third convolutional normalization layer, and the third Downsampling layer, the fourth convolutional normalization layer; the gear surface defect image processed by the downsampling layer uses 1*1 convolution for feature extraction, and the three gear surface defect feature maps obtained are respectively denoted as S2, S3, S4;

Neck part: The neck Neck structure adopts FPN+PAN, and the FPN layer transmits strong semantic information from top to bottom. S4 undergoes 3*3 convolution, and the feature map obtained is marked as Q4. connected, and then through 3*3 convolution, the obtained feature map is marked as Q3; Q3 is connected with S2 after transposed convolution upsampling, and then after 3*3 convolution, the obtained feature map is marked as Q2;

PAN transmits strong positioning information from the bottom up. The feature map Q2 is used as the bottom feature R2. After R2 is down-sampled, it is connected with Q3, and the feature map is marked as R3; after R3 is down-sampled, it is connected with Q4, and the obtained feature map is marked as is R4; R2, R3, and R4 undergo 3*3 convolution respectively to obtain feature maps T2, T3, and T4;

After the last three C3 modules of the neck Neck structure, an AMECA attention module is integrated respectively, and the feature maps T2, T3, and T4 are used as the original input feature maps, respectively, through the global average pooling module and the global maximum pooling module, and Add the two obtained feature maps, compress the spatial information, and then use 1*1 convolution to learn channel attention information, combine the obtained channel attention information with the original input feature map, and finally obtain a specific channel attention feature map D1, D2, D3;

Output part: Input the feature maps D1, D2, and D3 into the YOLOv5-MSTA detection head network respectively, and finally get the detection results.

Further, the Mosaic9 data enhancement includes: firstly take out a batch of data from the total data set, randomly take out 9 pictures from it each time, perform cropping and zooming at random positions, and synthesize new pictures; the above process is repeated batch-size times , and finally get a batch of new data including batch-size pictures after Mosaic9 data enhancement, and then pass it to the neural network for training.

Further, the CBRM operation includes Conv, BN, ReLU and MaxPool.

Further, the first downsampling layer, the second downsampling layer, and the third downsampling layer all include a Shuffle_Block(d) module, and the Shuffle_Block(d) divides the input features into two branches, the left side The branch has 2 convolutional layers, which are 3×3 depth convolution with a step size of 2 and 1×1 ordinary convolution; there are three convolutional layers on the right branch, which are 1×1 ordinary convolution, 3×3 depth convolution with a step size of 2 and 1×1 ordinary convolution; the left and right branches are spliced through Concat to fuse the features on the left and right sides, and finally the channel shuffling operation is used to enable the two branches communication between information.

Further, the second convolutional normalization layer, the third convolutional normalization layer, and the fourth convolutional normalization layer all include a Shuffle_Block(c) module, and the Shuffle_Block(c) module pairs Each channel is split and divided into two branches. According to the criterion of reducing the degree of fragmentation of the model, no operation is performed on the left branch, and there are 3 convolutional layers on the right branch, which are 1×1 ordinary Convolution, 3×3 depth convolution, 1×1 ordinary convolution, 1×1 ordinary convolution, 3×3 depth convolution, 1×1 ordinary convolution three convolutional layers have the same input and output channels, Two of the 1×1Convs are no longer group convolutions but changed to ordinary convolutions. After the three convolutions, the two branches are spliced by Concat.

Further, the operation steps of the transposed convolution upsampling include:

(1) Fill s-1 rows and columns 0 between the input feature map elements, where s represents the step size of the transposed convolution;

(2) Fill k-p-1 rows and columns 0 around the input feature map, where k represents the kernel_size of the transposed convolution, and p is the padding of the transposed convolution;

(3) Flip the convolution kernel parameters up and down, left and right;

(4) Do normal convolution operation, fill with 0, and step 1.

Further, the process of the AMECA attention module includes:

(1) First input the feature map X, the dimension of the feature map X is H*W*C;

(2) Perform spatial feature compression on the input feature map X; in the spatial dimension, use the global average pooling GAP to obtain the feature map F1 of 1*1*C; use the global maximum pooling GMP to obtain the 1*1*C feature map Feature map F2;

(3) F1 and F2 are fused to obtain a 1*1*C feature map F3 to obtain higher-level semantic information;

(4) Carry out channel feature learning on the fused feature map F3; through 1*1*1 convolution, learn the importance between different channels, and the dimension of the output feature map F4 at this time is still 1*1*C;

(5) Pass the feature map F4 through the σ function to obtain F41;

(6) Finally, channel attention is combined, and the feature map F41 of channel attention is multiplied channel-by-channel with the original input feature map X, and finally the feature map X' with channel attention is output;

Among them, H, W and C represent the height, width and channel number of the input feature map, respectively, and σ represents the activation function;

The feature maps T2, T3, and T4 are used as input feature maps X, and the output feature maps X' obtained respectively are feature maps D1, D2, and D3.

A computer readable medium having stored thereon a computer program for performing the method as described above.

Beneficial effects of the present invention:

In the present invention, by using Mosaic9 data enhancement at the input end, the small sample target is added while enriching the data set, and the training speed and generalization ability of the network are improved; in order to facilitate model deployment, the present invention uses ShuffleNetv2 as the backbone network to extract features, and the channel weight To achieve cross-group information exchange, build a YOLOv5 lightweight neural network model, reduce the amount of network parameters and improve the speed of model detection; by using transposed convolution for upsampling, semantic level upsampling is realized, so that features contain more Strong semantic information can make the network more lightweight; finally, add the AMECA attention mechanism to the Neck structure, adjust the information extraction method of the model channel through the attention module, strengthen the channel features, make the defect detection more accurate, and further enhance the gear defect Feature extraction capability to improve the performance of gear defect model detection.

Description of drawings

Fig. 1 is the network structure diagram of the MSTA-YOLOv5 detection model of the present invention;

Fig. 2 is the flow chart of Mosaic9 data enhancement of the present invention;

Fig. 3 is the structural diagram of two kinds of modules of ShuffleNetv2 of the present invention;

Fig. 4 is the structural diagram of the AMECA attention module of the present invention;

Figure 5 is a network structure diagram of YOLOv5;

Detailed ways

In order to clearly illustrate the technical features of the present solution, the present invention will be described in detail below through specific implementation methods and in conjunction with the accompanying drawings. The following disclosure provides many different embodiments or examples for implementing different structures of the present invention. To simplify the disclosure of the present invention, components and arrangements of specific examples are described below. Furthermore, the present invention may repeat reference numerals and/or letters in different instances. This repetition is for the purpose of simplicity and clarity and does not in itself indicate a relationship between the various embodiments and/or arrangements discussed. It should be noted that components illustrated in the figures are not necessarily drawn to scale. Descriptions of well-known components and processing techniques and processes are omitted herein to avoid unnecessarily limiting the present invention.

A light-weight gear surface defect detection method based on MSTA-YOLOv5, comprising the following steps: firstly obtain a gear surface defect image, and mark and divide the image to construct a gear surface defect data set; then construct MSTA-YOLOv5 detection The model is based on the gear surface defect data set to train the MSTA-YOLOv5 detection model; finally, the gear defect image to be detected is sent to the trained MSTA-YOLOv5 detection model to obtain the defect type of the detected gear. The gear surface defect types in the present invention include three types of tooth bottom black skin, tooth surface black skin and bumps.

As shown in Figure 1, the MSTA-YOLOv5 detection model includes:

Input part: Input the gear surface defect image into the MSTA-YOLOv5 network for adaptive anchor frame calculation and Mosaic9 data enhancement;

Backbone part: the feature extraction backbone network adopts the ShuffleNetv2 architecture, including sequentially connected CBRM operations, the first downsampling layer, the second convolutional normalization layer, the second downsampling layer, the third convolutional normalization layer, and the third Downsampling layer, the fourth convolutional normalization layer; the gear surface defect image processed by the downsampling layer uses 1*1 convolution for feature extraction, and the three gear surface defect feature maps obtained are respectively denoted as S2, S3, S4;

As shown in Figure 5 and Figure 1, on the basis of the traditional YOLOv5 model, the present invention uses the ShuffleNetV2 architecture instead of CSPDarknet53 as a feature extraction network to construct a YOLOv5 lightweight neural network model. ShuffleNetV2 not only inherits the characteristics of ShuffleNet group convolution and channel rearrangement, but also follows the four guidelines for designing lightweight networks. Under the same conditions, ShuffleNetV2 is faster and more accurate than other models. The MSTA-YOLOv5 model inputs the target feature quantity extracted by ShuffleNetV2, adaptively adjusts the parameters of the network model according to the loss value returned by each iteration, and waits for the loss value to converge and stabilize, then the detection model with the best evaluation index can be obtained. It greatly reduces the parameters of the model Parameters and the amount of calculation FLOPs, reducing the size of the model. Table 1 shows the comparison of the number of layers, the amount of parameters, and the amount of calculation of the two architectures of CSPDarknet53 and ShuffleNetv2.

Table 1 Comparison of backbone network parameters

Model Network layers Parameter amount Calculations CSPDarknet53 270 7.03M 16.0 GFLOPs ShuffleNetv2 308 3.79M 8.0 GFLOPs

Among them, FLOPs refers to the amount of calculation. For the convolutional layer, the calculation formula of FLOPs is as follows:

FLOPs＝2HW(C _in K ² +1)C _out (1)

Among them, C _in refers to the number of channels of the input tensor of the convolutional layer, C _out refers to the number of channels of the output tensor of the convolutional layer, and K refers to the size of the convolution kernel, and then the constant item is removed, which is simplified as:

FLOPs＝HW(C _in K ² )C _out (2)

For the convolutional layer, the calculation formula of the parameter quantity Parameters is as follows:

parameters＝Co×(Ci×K×K+1) (3)

Among them, Co is the number of output channels, Ci is the number of input channels, and K refers to the size of the convolution kernel;

H and W represent the height and width of the input feature map, respectively;

As shown in Figure 3, the ShuffleNetv2 backbone has two modules: Shuffle_Block(d), Shuffle_Block(c), Shuffle_Block(c) and Shuffle_Block(d) can be distinguished by the step size in the yaml configuration file, Shuffle_Block(c) The stride of stride=2, the stride of Shuffle_Block(d)=1, these two are used alternately in the present invention. Replace the Focus slice at the input end of the original YOLOv5 network with CBRM, a 3*3 convolution, replace all Conv+C3 in the backbone network with Shuffle_Block, and remove the SPP and a subsequent C3 structure, because the parallel operation of SPP will affect the speed.

Neck part: The neck Neck structure adopts FPN+PAN, and the FPN layer transmits strong semantic information from top to bottom. S4 undergoes 3*3 convolution, and the feature map obtained is marked as Q4. connected, and then through 3*3 convolution, the obtained feature map is marked as Q3; Q3 is connected with S2 after transposed convolution upsampling, and then after 3*3 convolution, the obtained feature map is marked as Q2;

PAN transmits strong positioning information from the bottom up. The feature map Q2 is used as the bottom feature R2. After R2 is down-sampled, it is connected with Q3, and the feature map is marked as R3; after R3 is down-sampled, it is connected with Q4, and the obtained feature map is marked as is R4; R2, R3, and R4 undergo 3*3 convolution respectively to obtain feature maps T2, T3, and T4;

After the last three C3 modules of the neck Neck structure, an AMECA attention module is integrated respectively, and the feature maps T2, T3, and T4 are used as the original input feature maps, respectively, through the global average pooling module and the global maximum pooling module, and Add the two obtained feature maps, compress the spatial information, and then use 1*1 convolution to learn channel attention information, combine the obtained channel attention information with the original input feature map, and finally obtain a specific channel attention feature map D1, D2, D3;

Output part: Input the feature maps D1, D2, and D3 into the YOLOv5-MSTA detection head network respectively, and finally get the detection results.

As shown in Figure 2, the Mosaic9 data enhancement includes: firstly take out a batch of data from the total data set, randomly take out 9 pictures from it each time, perform cropping and zooming at random positions, and synthesize new pictures; the above process repeats batch -size times, and finally get a batch of new data including batch-size images after Mosaic9 data enhancement, and then pass it to the neural network for training. The present invention adopts Mosaic9 data enhancement, uses 9 images for random cutting and zooming, and then randomly arranges and stitches them to form a picture, which improves and enriches the data set while increasing the small sample target and improving the training speed and generalization of the network Ability, the data of 9 pictures will be calculated at one time during the normalization operation, so the memory requirement of the model is reduced.

Specifically, the CBRM operation includes Conv, BN, ReLU and MaxPool.

Specifically, the first down-sampling layer, the second down-sampling layer, and the third down-sampling layer all include a Shuffle_Block(d) module, and the Shuffle_Block(d) does not perform a shunting operation, and divides the input feature There are two branches, the left branch has 2 convolutional layers, which are 3×3 depth convolution with a step size of 2 and 1×1 ordinary convolution; there are three convolutional layers on the right branch, which are 1 ×1 ordinary convolution, 3×3 depth convolution with a step size of 2, and 1×1 ordinary convolution; the left and right branches are spliced by Concat to fuse the features on the left and right sides, and finally shuffled through channels Action to enable communication of information between two branches. The difference from Shuffle_Block(c) is that a 3×3 depth convolution is introduced on the left and right sides to achieve downsampling.

Specifically, the second convolutional normalization layer, the third convolutional normalization layer, and the fourth convolutional normalization layer all include a Shuffle_Block(c) module, and the Shuffle_Block(c) module is Each channel is split and divided into two branches. According to the criterion of reducing the degree of fragmentation of the model, no operation is performed on the left branch, and there are 3 convolutional layers on the right branch, which are 1×1 ordinary Convolution, 3×3 depth convolution, 1×1 ordinary convolution, 1×1 ordinary convolution, 3×3 depth convolution, 1×1 ordinary convolution three convolutional layers have the same input and output channels, Two of the 1×1Conv are no longer group convolutions but changed to ordinary convolutions. After the three convolutions, the two branches are spliced by Concat; this can satisfy the input and output channels. Similarly, the two branches are spliced through Concat to perform channel shuffling operations, and information communication between the two branches is enabled through channel shuffling operations.

Specifically, the operation steps of the transposed convolution upsampling include:

(1) Fill s-1 rows and columns 0 between the input feature map elements, where s represents the step size of the transposed convolution;

(2) Fill k-p-1 rows and columns 0 around the input feature map, where k represents the kernel_size of the transposed convolution, and p is the padding of the transposed convolution;

(3) Flip the convolution kernel parameters up and down, left and right;

(4) Do normal convolution operation, fill with 0, and step 1.

The upsampling method of transposed convolution is used to obtain the optimal upsampling method through self-learning by the network, and the upsampling at the semantic level is realized, so that the feature contains stronger semantic information. The transposed convolution calculation process is to use each element value of the input as the weight of the convolution kernel, multiply it as the upsampling output corresponding to the element, and add the overlapping output parts of different inputs directly as the output.

The proposed AMECA attention module is shown in Figure 4. The feature map is passed through the global average pooling module and the global maximum pooling module, and the two obtained feature maps are added to compress the spatial information, and then use 1*1 Convolutional learning channel attention information, the obtained channel attention information is combined with the original input feature map, and finally a specific channel attention feature map is obtained. AMECA avoids dimensionality reduction and effectively captures cross-channel interaction information, enabling the network to locate and identify target areas more precisely.

The process of the AMECA attention module includes:

(1) First input the feature map X, the dimension of the feature map X is H*W*C;

(2) Perform spatial feature compression on the input feature map X; in the spatial dimension, use the global average pooling GAP to obtain the feature map F1 of 1*1*C; use the global maximum pooling GMP to obtain the 1*1*C feature map Feature map F2;

(3) F1 and F2 are fused to obtain a 1*1*C feature map F3 to obtain higher-level semantic information;

(4) Carry out channel feature learning on the fused feature map F3; through 1*1*1 convolution, learn the importance between different channels, and the dimension of the output feature map F4 at this time is still 1*1*C;

(5) Pass the feature map F4 through the σ function to obtain F41;

(6) Finally, channel attention is combined, and the feature map F41 of channel attention is multiplied channel-by-channel with the original input feature map X, and finally the feature map X' with channel attention is output;

Among them, H, W and C represent the height, width and channel number of the input feature map, respectively, and σ represents the activation function;

The feature maps T2, T3, and T4 are used as input feature maps X, and the output feature maps X' obtained respectively are feature maps D1, D2, and D3.

In another embodiment, a computer readable medium has a computer program stored thereon for performing the method as described above.

After the lightweight processing, the amount of network calculations and parameters are greatly reduced. After the last three C3 modules of the Neck structure of the YOLOv5 network structure model, the AMECA attention mechanism is introduced to adjust the information extraction method of the model space and channels. This method guarantees no loss of accuracy and can meet the needs of real-time detection of gear surface defects.

The present invention has done comparative test, as shown in table 2. The comparison test was measured from the amount of parameters, calculation amount, and model size. The smaller the amount of parameters, calculation amount, and model size, the lower the network complexity. The parameters, calculation amount and model size of the detection network using the ShuffleNetv2 module are lower than those of the detection network without the ShuffleNetv2 module under the same parameter conditions, which shows that the ShuffleNetv2 module has a lightweight effect.

Table 2 Comparison of parameter quantity, calculation quantity and model size of different models

According to the test results and the results in the table, the MSTA-YOLOv5 model has great advantages over the YOLOv3, YOLOv4 and YOLOv5s models. Compared with the original YOLOv5s model, the model, the amount of parameters and the amount of calculation have been greatly reduced, the amount of parameters has been reduced by about 46%, the amount of calculation has been reduced by 50%, and the size of the model has been reduced by about 44%. The new model is more streamlined, the complexity is significantly reduced, and the requirements for deployment on the mobile end are met. It has better detection results for gear surface defects.

Although the specific implementation of the invention has been described above in conjunction with the accompanying drawings, it does not limit the protection scope of the present invention. On the basis of the technical solution of the present invention, those skilled in the art can make various Modifications or variations are still within the protection scope of the present invention.

Claims (7)

1. A light-weight gear surface defect detection method based on MSTA-YOLOv5 is characterized by comprising the following steps: firstly, acquiring a gear surface defect image, marking and dividing the image, and constructing a gear surface defect data set; then constructing an MSTA-YOLOv5 detection model, and training the MSTA-YOLOv5 detection model based on the gear surface defect data set; finally, the gear defect image to be detected is sent to a trained MSTA-YOLOv5 detection model, and the defect type of the detected gear is obtained;

the MSTA-YOLOv5 detection model comprises:

an input section: inputting the gear surface defect image into an MSTA-YOLOv5 network, and performing self-adaptive anchor frame calculation and Mosaic9 data enhancement;

a backbone part: the feature extraction backbone network adopts a ShuffleNetv2 architecture and comprises a CBRM operation, a first downsampling layer, a second convolution normalization layer, a second downsampling layer, a third convolution normalization layer, a third downsampling layer and a fourth convolution normalization layer which are sequentially connected; respectively recording 3 gear surface defect feature maps obtained by performing feature extraction on the gear surface defect image subjected to the down-sampling layer processing by utilizing 1-by-1 convolution as S2, S3 and S4;

a neck portion: the Neck portion Neck structure adopts FPN + PAN, the FPN layers transmit strong semantic information from top to bottom, S4 is convoluted by 3 x 3 to obtain a characteristic diagram Q4, Q4 is connected with S3 after being subjected to transposition convolution and upsampling, and the characteristic diagram Q3 is obtained through 3 x 3 convolution; q3 is connected with S2 after being subjected to transposed convolution and up sampling, and the characteristic graph obtained through 3-by-3 convolution is marked as Q2;

PAN transmits strong positioning information from bottom to top, the characteristic diagram Q2 is used as a bottom characteristic R2, and the R2 is connected with the Q3 after down sampling to obtain a characteristic diagram R3; r3 is connected with Q4 after being downsampled, and the obtained characteristic diagram is marked as R4; r2, R3 and R4 are respectively convolved by 3 x 3 to obtain characteristic diagrams T2, T3 and T4;

respectively integrating an AMECA attention module behind the last 3C 3 modules of the Neck Neck structure, respectively taking the feature maps T2, T3 and T4 as original input feature maps, respectively passing through a global average pooling module and a global maximum pooling module, adding the two obtained feature maps, compressing spatial information, then using 1 × 1 convolution to learn channel attention information, combining the obtained channel attention information with the original input feature maps, and finally obtaining specific channel attention feature maps D1, D2 and D3;

an output section: and respectively inputting the characteristic diagrams D1, D2 and D3 into a YOLOv5-MSTA detection head network to finally obtain a detection result.

2. The MSTA-YOLOv 5-based light-weight gear surface defect detection method as claimed in claim 1, wherein the Mosaic9 data enhancement comprises: firstly, taking out a batch of data from a total data set, randomly taking out 9 pictures from the data set each time, cutting and zooming at random positions, and synthesizing new pictures; the process is repeated for the batch-size times, and finally a batch of new data comprising the batch-size pictures subjected to the Mosaic9 data enhancement is obtained and transmitted to the neural network for training.

3. The MSTA-YOLOv 5-based light-weight gear surface defect detection method as claimed in claim 1, wherein the CBRM operation comprises Conv, BN, reLU and MaxPool.

4. The MSTA-YOLOv 5-based lightweight gear surface defect detection method of claim 1, wherein the first downsampling layer, the second downsampling layer and the third downsampling layer respectively comprise a Shuffle _ Block (d) module, the Shuffle _ Block (d) module divides an input feature into two branches, the left branch has 2 convolution layers, and the convolution layers are respectively a 3 x 3 depth convolution layer with a step size of 2 and a common convolution with a step size of 1 x 1; three convolution layers are arranged on the right branch, namely 1 × 1 common convolution, 3 × 3 depth convolution with the step length of 2 and 1 × 1 common convolution respectively; splicing the left and right branches through Concat to fuse the characteristics of the left and right branches, and finally starting information communication between the two branches through channel shuffling operation;

the second convolution normalization layer, the third convolution normalization layer and the fourth convolution normalization layer respectively comprise a Shuffle _ Block (c) module, the Shuffle _ Block (c) module splits each channel and divides each channel into two branches, no operation is performed on the left branch according to the fragmentation degree criterion of a model to be reduced, 3 convolutional layers are arranged on the right branch and respectively comprise 1 × 1 ordinary convolution, 3 × 3 deep convolution and 1 × 1 ordinary convolution, the three convolutional layers of 1 × 1 ordinary convolution, 3 × 3 deep convolution and 1 × 1 ordinary convolution have the same input and output channels, two 1 × 1Conv are not group convolutional products and are changed into ordinary convolution, and after the 3 convolutions, the two branches are spliced through Concat.

5. The MSTA-YOLOv 5-based lightweight gear surface defect detection method of claim 1, wherein the operation step of transpose convolution upsampling comprises:

(1) Filling s-1 rows and columns 0 among the elements of the input feature graph, wherein s represents the step pitch of the transposition convolution;

(2) Filling k-p-1 rows and columns 0 around the input feature map, wherein k represents the kernel _ size of the transposition convolution, and p is filling of the transposition convolution;

(3) Turning the convolution kernel parameters up and down, left and right;

(4) And (5) performing normal convolution operation, filling 0 and step 1.

6. The MSTA-YOLOv 5-based light-weight gear surface defect detection method as claimed in claim 1, wherein the AMECA attention module process comprises:

(1) Firstly, inputting a feature map X, wherein the dimension of the feature map X is H, W and C;

(2) Performing spatial feature compression on the input feature map X; using the global average pooled GAP in the spatial dimension to obtain a feature map F1 of 1 × c; using global maximum pooling GMP, a feature map F2 of 1 × c was obtained;

(3) Fusing the F1 and the F2 to obtain a feature map F3 of 1 × C, and obtaining higher-level semantic information;

(4) Channel feature learning is carried out on the fused feature map F3; learning the importance among different channels through 1 × 1 convolution, wherein the dimension of the output feature map F4 is still 1 × C;

(5) Obtaining F41 by the characteristic diagram F4 through a sigma function;

(6) Finally, combining the attention of the channels, and performing channel-by-channel product on the characteristic diagram F41 of the attention of the channels and the original input characteristic diagram X to finally output a characteristic diagram X' with the attention of the channels;

wherein, H, W and C respectively represent the height, width and channel number of the input characteristic diagram, and sigma represents an activation function;

the characteristic diagrams T2, T3 and T4 are used as input characteristic diagrams X, and the obtained output characteristic diagrams X' are respectively characteristic diagrams D1, D2 and D3.

7. A computer-readable medium, characterized in that a computer program is stored thereon for performing the method according to any of claims 1-6.

Images