A lightweight gear surface defect detection method based on MSTA-YOLOv5

Abstract

The present application relates to product defect detection technical field, disclose a kind of light weight gear surface defect detection method based on MSTA-YOLOv5, comprising the following steps: first, gear surface defect image is obtained, and the image is labeled and divided, and gear surface defect dataset is constructed;Then, MSTA-YOLOv5 detection model is constructed, and MSTA-YOLOv5 detection model is trained based on gear surface defect dataset;Finally, the gear defect image to be detected is sent into the trained MSTA-YOLOv5 detection model, and the defect type of detection gear is obtained.The present application solves the problem that the demand for computing resources is too large, the cost is high due to serious memory consumption, enterprises need low-latency models and mobile device terminals need small models that are both fast and accurate, realizes the detection and automatic sorting of gear surface defects, and improves the detection efficiency of gear surface defect detection.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of product defect detection, and in particular to a lightweight gear surface defect detection method based on MSTA-YOLOv5. BACKGROUND

[0002] With the development of science and technology and the change of social needs, large-scale multi-faceted complex structure workpieces are more popular in industrial production. Gears are widely used transmission components in the mechanical industry, and their quality is particularly important in production. However, in the actual production process, due to factors such as process flow, production equipment and on-site environment, various defects may occur on the surface of the gear. If not handled in a timely manner, it will affect the apparent quality, performance and service life of the gear, resulting in a decline in enterprise production efficiency. Therefore, the surface of the gear needs to be detected, and the traditional manual detection is labor-intensive and prone to visual fatigue of the detection personnel, resulting in missed detection and false detection.

[0003] In recent years, with the rapid development of machine vision technology, detection technology based on machine vision has been applied in product surface quality detection. However, most of the current gear defect detection technologies use digital image processing technology, but this technology has a single processing method and algorithm, and it is difficult to effectively extract defect targets when facing gears with high surface complexity, resulting in unsatisfactory detection results.

[0004] The patent document with publication number CN115187820A discloses a lightweight target detection method, device, equipment and storage medium, which uses ShuffleNetv as a feature extraction module in the YOLOv4 network structure. However, the parameter quantity and calculation quantity are very large, and the SE attention mechanism is used in the feature extraction module, which has the disadvantage of insufficient precision.

[0005] The patent document with publication number CN112990325A discloses a lightweight network construction method for embedded real-time visual target detection, which uses the CBAM attention mechanism and has the advantage of lightweight. However, the precision loss is large with the improvement of lightweight, and the Focus slicing operation is used in this technical solution, which increases the parameter quantity and weakens the advantage of lightweight.

[0006] The patent document with publication number CN114898171A discloses a real-time target detection method suitable for embedded platforms, which achieves the effect of lightweight, but the precision loss is still relatively large.

[0007] With the development of artificial intelligence technology, deep learning methods have been widely used in image processing and workpiece quality detection fields due to their superior performance in handling complex background and weak defects. Using deep learning methods can accurately identify and segment gear surface defects, reducing the interference of background and other factors, thereby effectively improving the detection accuracy. Although there are a large number of studies on improving different target detection networks and detecting defects in industrial products, the results are satisfactory, but there is no research on small volume and fewer calculation parameter models for enterprises. Such models can also achieve good detection speed and accuracy in devices with low cost budget and relatively insufficient computing power. SUMMARY

[0008] Deep learning methods have greatly improved accuracy in image classification, but current deep learning-based target detection algorithms require too much computing resources and consume too much memory, making the cost high. To solve the above problems and the problem of low-latency models required by enterprises and small models required by mobile device terminals that are both fast and accurate, the present application provides a lightweight gear surface defect detection method based on MSTA-YOLOv5, which realizes the detection and automatic sorting of gear surface defects and can improve the detection efficiency of gear surface defect detection.

[0009] The technical solution of the present application to solve the technical problem is:

[0010] A lightweight gear surface defect detection method based on MSTA-YOLOv5, comprising the following steps: first, acquiring a gear surface defect image, and labeling and dividing the image to construct a gear surface defect dataset; then constructing an MSTA-YOLOv5 detection model, training the MSTA-YOLOv5 detection model based on the gear surface defect dataset; finally, inputting the gear defect image to be detected into the trained MSTA-YOLOv5 detection model to obtain the defect type of the detected gear.

[0011] The MSTA-YOLOv5 detection model comprises:

[0012] Input part: input the gear surface defect image into the MSTA-YOLOv5 network, perform adaptive anchor frame calculation and Mosaic9 data enhancement;

[0013] The backbone part: the feature extraction backbone network adopts a ShuffleNetv2 architecture, and includes sequentially connected CBRM operations, a first down-sampling layer, a second convolutional normalization layer, a second down-sampling layer, a third convolutional normalization layer, a third down-sampling layer and a fourth convolutional normalization layer; three gear surface defect feature maps obtained after 1*1 convolution feature extraction on the gear surface defect image processed by the down-sampling layer are respectively denoted as S2, S3 and S4;

[0014] The neck part: the neck Neck structure adopts FPN+PAN, the FPN layer transmits strong semantic information from top to bottom, S4 is subjected to 3*3 convolution to obtain a feature map denoted as Q4, Q4 is subjected to transposed convolution up-sampling and connected with S3, and then subjected to 3*3 convolution to obtain a feature map denoted as Q3; Q3 is subjected to transposed convolution up-sampling and connected with S2, and then subjected to 3*3 convolution to obtain a feature map denoted as Q2;

[0015] The PAN transmits strong positioning information from bottom to top, the feature map Q2 is taken as a bottom layer feature R2, R2 is subjected to down-sampling and connected with Q3 to obtain a feature map denoted as R3; R3 is subjected to down-sampling and connected with Q4 to obtain a feature map denoted as R4; R2, R3 and R4 are respectively subjected to 3*3 convolution to obtain feature maps T2, T3 and T4;

[0016] An AMECA attention module is integrated after the last three C3 modules of the neck Neck structure, the feature maps T2, T3 and T4 are respectively taken as original input feature maps, subjected to a global average pooling module and a global maximum pooling module, and the two obtained feature maps are added to compress spatial information, and then 1*1 convolution is used to learn channel attention information, and the obtained channel attention information is combined with the original input feature map to finally obtain specific channel attention feature maps D1, D2 and D3;

[0017] The output part: the feature maps D1, D2 and D3 are respectively input into a YOLOv5-MSTA detection head network to finally obtain a detection result.

[0018] Further, the Mosaic9 data enhancement includes: first, a batch of data is taken out from the total data set, and 9 pictures are randomly taken out each time, and new pictures are synthesized by random position cutting and scaling; the above process is repeated batch-size times, and finally a batch of new data including batch-size pictures after Mosaic9 data enhancement is obtained, and then the neural network is transmitted for training.

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

[0020] Further, the first down-sampling layer, the second down-sampling layer and the third down-sampling layer each comprise a Shuffle_Block(d) module, the Shuffle_Block(d) module divides the input features into two branches, the left branch has two convolution layers, which are a 3*3 depth convolution with a step of 2 and a 1*1 normal convolution respectively; the right branch has three convolution layers, which are a 1*1 normal convolution, a 3*3 depth convolution with a step of 2 and a 1*1 normal convolution respectively; the features on the left and right sides are fused by concatenating the left and right branches, and finally the channel shuffling operation is performed to enable information communication between the two branches.

[0021] Further, the second convolution normalization layer, the third convolution normalization layer and the fourth convolution normalization layer each comprise a Shuffle_Block(c) module, the Shuffle_Block(c) module divides each channel into two branches, according to the rule of reducing the fragmentation degree of the model, no operation is performed on the left branch, and three convolution layers are arranged on the right branch, which are a 1*1 normal convolution, a 3*3 depth convolution and a 1*1 normal convolution, the three convolution layers have the same input and output channels, and two 1*1Conv are changed from group convolution to normal convolution, and the two branches are spliced by Concat after the three convolutions.

[0022] Further, the operation steps of the transposed convolution up-sampling comprise:

[0023] (1) fill s-1 rows and 0 columns between input feature elements, where s represents the step distance of the transposed convolution;

[0024] (2) fill k-p-1 rows and 0 columns around the input feature Figure Four

[0025] (3) flip the convolution kernel parameters up and down and left and right;

[0026] (4) perform normal convolution operation, fill 0 and step 1.

[0027] Further, the process of the AMECA attention module comprises:

[0028] (1) first input the feature map X, the dimension of the feature map X is H*W*C;

[0029] ​(2) Spatial feature compression is performed on the input feature map X; in the spatial dimension, a global average pooling GAP is used to obtain a feature map F1 of 1*1*C; a global maximum value pooling GMP is used to obtain a feature map F2 of 1*1*C;

[0030] (3) F1 and F2 are fused to obtain a feature map F3 of 1*1*C, and higher-level semantic information is obtained;

[0031] (4) Channel feature learning is performed on the fused feature map F3; through 1*1*1 convolution, the importance between different channels is learned, and the dimension of the output feature map F4 is still 1*1*C;

[0032] (5) The feature map F4 is subjected to a sigma function to obtain F41;

[0033] (6) Finally, channel attention is combined, the channel attention feature map F41 is multiplied with the original input feature map X, and a feature map X' with channel attention is finally output;

[0034] Wherein, H, W and C represent the height, width and channel number of the input feature map respectively, and sigma represents an activation function;

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

[0036] A computer readable medium has a computer program stored thereon for executing the method as described above.

[0037] The beneficial effects of the present application are as follows:

[0038] The present application uses Mosaic9 data enhancement at the input end, enriches the data set while increasing the small sample target, and improves the training speed and generalization ability of the network; in order to facilitate model deployment, the present application uses ShuffleNetv2 as a backbone network to extract features, rearranges the channels to realize cross-group information exchange, constructs a YOLOv5 lightweight neural network model, reduces the network parameter quantity while improving the model detection speed; by using the transposed convolution method for upsampling, semantic-level upsampling is realized, which enables the feature to contain stronger semantic information and makes the network more lightweight; finally, the AMECA attention mechanism is added in the Neck structure, the information extraction mode of the model channel is adjusted through the attention module, the channel features are strengthened, the defect detection is more accurate, the gear defect feature extraction capability is further enhanced, and the performance of the gear defect model detection is improved. BRIEF DESCRIPTION OF DRAWINGS

[0039] Figure 1 It is a network structure diagram of the MSTA-YOLOv5 detection model of the present application;

[0040] Figure 2 Flow structure diagram for Mosaic9 data enhancement of the application;

[0041] Figure 3 Structure diagram of two modules of ShuffleNetv2 of the application;

[0042] Figure 4 Structure diagram of AMECA attention module of the application;

[0043] Figure 5 Network structure diagram of YOLOv5; DETAILED DESCRIPTION

[0044] In order to clearly illustrate the technical features of the scheme, the application will be described in detail below through specific implementation manners, and in conjunction with the accompanying drawings. The following disclosure provides many different embodiments or examples for implementing different structures of the application. In order to simplify the disclosure of the application, the components and settings of specific examples are described below. In addition, the application can repeatedly refer to numerals and / or letters in different examples. Such repetition is for the purpose of simplification and clarity, and does not indicate the relationship between the various embodiments and / or settings discussed. It should be noted that the components illustrated in the drawings are not necessarily drawn to scale. The application omits the description of well-known components and processing techniques and processes to avoid unnecessary limitation of the application.

[0045] A lightened gear surface defect detection method based on MSTA-YOLOv5, comprising the following steps: first, acquiring a gear surface defect image, and labeling and dividing the image to construct a gear surface defect dataset; then constructing a MSTA-YOLOv5 detection model, training the MSTA-YOLOv5 detection model based on the gear surface defect dataset; finally, inputting a gear defect image to be detected into the trained MSTA-YOLOv5 detection model to obtain the defect type of the detected gear. The gear surface defect types of the application include three types of tooth bottom black skin, tooth surface black skin and knock.

[0046] As shown in Figure 1 , the MSTA-YOLOv5 detection model comprises:

[0047] Input part: input the gear surface defect image into the MSTA-YOLOv5 network to perform adaptive anchor frame calculation and Mosaic9 data enhancement;

[0048] The backbone part: the feature extraction backbone network adopts a ShuffleNetv2 architecture, and includes sequentially connected CBRM operations, a first down-sampling layer, a second convolution normalization layer, a second down-sampling layer, a third convolution normalization layer, a third down-sampling layer and a fourth convolution normalization layer; three gear surface defect feature maps obtained after the gear surface defect image processed by the down-sampling layer is subjected to feature extraction by 1*1 convolution are respectively denoted as S2, S3 and S4;

[0049] As shown in Figure 5 , Figure 1 YOLOv5 lightweight neural network model is constructed by replacing CSPDarknet53 with ShuffleNetV2 architecture as the feature extraction network on the basis of the traditional YOLOv5 model. ShuffleNetV2 inherits the characteristics of ShuffleNet grouping convolution and channel rearrangement, and also follows the four guidelines for designing lightweight networks. Under the same conditions, ShuffleNetV2 is faster than other models, and has better accuracy. The MSTA-YOLOv5 model inputs the target feature quantity extracted by ShuffleNetV2, and according to the loss value returned by each iteration, the parameters of the network model are adaptively adjusted, and when the loss value converges to be stable, the detection model with the best evaluation index can be obtained. The parameter quantity Parameters and the calculation quantity FLOPs of the model are greatly reduced, and the size of the model is reduced. The number of layers, parameter quantity and calculation quantity of the two architectures of CSPDarknet53 and ShuffleNetv2 are shown in Table 1.

[0050] Table 1 Comparison of parameter quantities of backbone networks

[0051] Model Number of network layers Parameter amount Computational amount CSPDarknet53 270 7.03M 16.0 GFLOPs ShuffleNetv2 308 3.79M 8.0 GFLOPs

[0052] Where FLOPs refers to the size of the calculation quantity, and for the convolution layer, the calculation formula of FLOPs is as follows:

[0053] FLOPs=2HW(C in K 2 +1)C out (1)

[0054] Where C in is the channel number of the input tensor of the convolution layer, C out is the channel number of the output tensor of the convolution layer, and K is the size of the convolution kernel, and then the constant term is removed, and simplified as:

[0055] FLOPs=HW(C in K 2 )C out (2)

[0056] For the convolution layer, the calculation formula of the parameter quantity Parameters is as follows:

[0057] parameters=Co×(Ci×K×K+1) (3)

[0058] Wherein, Co is the output channel number, Ci is the input channel number, and K refers to the size of the convolution kernel.

[0059] H and W represent the height and width of the input feature map, respectively.

[0060] As shown in the formula (1), the calculation formula of the parameter quantity Parameters of the convolution layer is as follows: Figure 3 As shown in the formula (1), the calculation formula of the parameter quantity Parameters of the convolution layer is as follows:

[0061] Neck part: the neck Neck structure adopts FPN+PAN, FPN layer transmits strong semantic information from top to bottom, S4 is subjected to 3*3 convolution, and a feature map is obtained and recorded as Q4; Q4 is subjected to transposed convolution upsampling and connected with S3, and then subjected to 3*3 convolution, and a feature map is obtained and recorded as Q3; Q3 is subjected to transposed convolution upsampling and connected with S2, and then subjected to 3*3 convolution, and a feature map is obtained and recorded as Q2;

[0062] PAN transmits strong positioning information from bottom to top, the feature map Q2 is taken as a bottom layer feature R2, R2 is subjected to downsampling and connected with Q3, and a feature map is obtained and recorded as R3; R3 is subjected to downsampling and connected with Q4, and a feature map is obtained and recorded as R4; R2, R3 and R4 are respectively subjected to 3*3 convolution, and feature maps T2, T3 and T4 are obtained.

[0063] An AMECA attention module is integrated after the last three C3 modules of the neck structure, respectively taking the feature maps T2, T3 and T4 as the original input feature maps, respectively passing through a global average pooling module and a global maximum pooling module, adding the two feature maps obtained to compress the spatial information, then using a 1*1 convolution to learn the channel attention information, combining the channel attention information obtained with the original input feature map to finally obtain specific channel attention feature maps D1, D2 and D3.

[0064] The output part: input the feature maps D1, D2 and D3 into the YOLOv5-MSTA detection head network to finally obtain the detection result.

[0065] As shown in Figure 2 , the Mosaic9 data enhancement includes: first taking a batch of data from the total data set, randomly taking 9 pictures each time, performing random position cutting and scaling to synthesize a new picture; the above process is repeated batch-size times, and finally a batch of new data including batch-size pictures after Mosaic9 data enhancement is obtained, which is then passed to the neural network for training. The present application adopts Mosaic9 data enhancement, which uses 9 images to perform random cutting and scaling, and then randomly arranges and splices to form a picture. This improvement realizes the enrichment of the data set while increasing the small sample target, improves the training speed and generalization ability of the network, and calculates the data of 9 pictures at a time during the normalization operation, so that the demand of the model for memory is reduced.

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

[0067] Specifically, the first down-sampling layer, the second down-sampling layer and the third down-sampling layer each include a Shuffle_Block(d) module. The Shuffle_Block(d) no longer performs a shunting operation to divide the input feature into two branches. The left branch has two convolution layers, which are a 3*3 depth convolution with a step of 2 and a 1*1 normal convolution, respectively. The right branch has three convolution layers, which are a 1*1 normal convolution, a 3*3 depth convolution with a step of 2 and a 1*1 normal convolution, respectively. The features on the left and right sides are fused by concatenating the left and right branches, and finally the information communication between the two branches is enabled by a channel shuffling operation. Unlike the Shuffle_Block(c), a 3*3 depth convolution is introduced on both the left and right sides to realize down-sampling.

[0068] Specifically, the second convolution normalization layer, the third convolution normalization layer and the fourth convolution normalization layer each include a Shuffle_Block(c) module, the Shuffle_Block(c) module performs shunting on each channel, divides into two branches, according to the criterion of reducing the fragmentation degree of the model, does not perform any operation on the left branch, and has three convolution layers on the right branch, which are 1*1 ordinary convolution, 3*3 depth convolution and 1*1 ordinary convolution, 1*1 ordinary convolution, 3*3 depth convolution and 1*1 ordinary convolution, the three convolution layers have the same input and output channels, wherein the two 1*1Conv are no longer group convolution but ordinary convolution, after the three convolutions, the two branches are spliced through Concat; in this way, the input and output channels are the same, the two branches are spliced through Concat, and the channel mixing operation is performed on the result of the two branches, and the information communication between the two branches is enabled through the channel mixing operation.

[0069] Specifically, the operation steps of the transposed convolution upsampling include:

[0070] (1) fill s-1 rows and 0 columns between input feature map elements, wherein s represents the stride of the transposed convolution;

[0071] (2) fill k-p-1 rows and 0 columns around the input feature Figure Four map, wherein k represents the kernel_size size of the transposed convolution, and p represents the padding of the transposed convolution;

[0072] (3) flip the convolution kernel parameters up and down and left and right;

[0073] (4) perform normal convolution operation, fill 0 and stride 1.

[0074] The transposed convolution upsampling mode is adopted to learn the optimal upsampling mode through the network, realize semantic level upsampling, and make the features contain stronger semantic information. The transposed convolution calculation process is to take each element value of the input as the weight of the convolution kernel, multiply it as the corresponding upsampling output of the element, and directly add the overlapping output parts of different inputs as the output.

[0075] The proposed AMECA attention module is as shown in Figure 4 The feature map is subjected to a global average pooling module and a global maximum pooling module, and the two obtained feature maps are added to compress the spatial information, then a 1*1 convolution is used to learn the channel attention information, the obtained channel attention information is combined with the original input feature map, and finally the specific channel attention feature map is obtained. AMECA avoids dimensionality reduction and effectively captures cross-channel interaction information, so that the network can more accurately locate and identify the target region.

[0076] The flow of the AMECA attention module comprises:

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

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

[0079] (3) Fuse F1 and F2 to obtain a feature map F3 of 1*1*C, and obtain higher-level semantic information;

[0080] (4) Perform channel feature learning on the fused feature map F3; learn the importance between different channels through 1*1*1 convolution, and the dimension of the output feature map F4 is still 1*1*C;

[0081] (5) Obtain F41 by passing the feature map F4 through the sigma function;

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

[0083] Wherein, H, W and C represent the height, width and channel number of the input feature map respectively, and sigma represents an activation function;

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

[0085] In another embodiment, a computer readable medium has a computer program stored thereon for executing the method as described above.

[0086] The network calculation amount and parameter amount after light processing are greatly reduced, the AMECA attention mechanism is introduced after the last 3 C3 modules of the Neck structure of the YOLOv5 network structure model, and the information extraction mode of the model space and channel is adjusted. This method well guarantees that the precision is not lost, and can meet the demand of real-time detection of gear surface defects.

[0087] The present application has made a comparative test, as shown in Table 2. The comparative test measures the parameter amount, calculation amount and model size, and the smaller the parameter amount, calculation amount and model size, the lower the network complexity. The parameter amount, calculation amount and model size of the detection network using the ShuffleNetv2 module are all lower than those of the detection network without using the ShuffleNetv2 module under the same parameter condition, so that the ShuffleNetv2 module has the effect of lightening.

[0088] Table 2 Parameter quantity, calculation quantity, model size comparison of different models

[0089]

[0090]

[0091] According to the detection results and the results in the table, the MSTA-YOLOv5 model has great advantages compared with the YOLOv3, YOLOv4 and YOLOv5s models. Compared with the original YOLOv5s model, the model, parameter quantity and calculation quantity are greatly reduced, the parameter quantity is reduced by about 46%, the calculation quantity is reduced by 50%, and the model size is reduced by about 44%. The new model is more compact, the complexity is obviously reduced, the requirements of mobile deployment are realized, and the gear surface defects have good detection results.

[0092] Although the specific embodiments of the application are described above with reference to the drawings, the description is not a limitation on the scope of protection of the application. Various modifications or variations made by those skilled in the art on the basis of the technical solutions of the application without creative labor are still within the scope of protection of the application.

Claims

1. A method for detecting surface defects in lightweight gears based on MSTA-YOLOv5, characterized in that, Includes the following steps: First, images of surface defects on gears are acquired, and then these images are labeled and segmented to construct a gear surface defect dataset. Then, an MSTA-YOLOv5 detection model is constructed and trained based on a gear surface defect dataset. Finally, the gear defect image to be detected is fed into the trained MSTA-YOLOv5 detection model to obtain the defect type of the detected gear. The MSTA-YOLOv5 detection model includes: Input section: Input the image of the gear surface defect into the MSTA-YOLOv5 network for adaptive anchor frame calculation and Mosaic9 data augmentation; Backbone: The feature extraction backbone network adopts the ShuffleNetv2 architecture, which includes a CBRM operation, a first downsampling layer, a second convolutional normalization layer, a second downsampling layer, a third convolutional normalization layer, a third downsampling layer, and a fourth convolutional normalization layer connected in sequence. The gear surface defect image after downsampling is processed by 1*1 convolution to extract features, resulting in three gear surface defect feature maps, denoted as S2, S3, and S4. Neck section: The neck structure adopts FPN+PAN. The FPN layer transmits strong semantic information from top to bottom. S4 is convolved by 3*3 to obtain feature map Q4. Q4 is upsampled by transposed convolution and then connected to S3. After another 3*3 convolution, the feature map Q3 is obtained. Q3 is upsampled by transposed convolution and then connected to S2. After another 3*3 convolution, the feature map Q2 is obtained. PAN transmits strong localization information from bottom to top. Feature map Q2 is used as the bottom feature R2. After R2 is downsampled, it is connected with Q3 to obtain feature map R3. After R3 is downsampled, it is connected with Q4 to obtain feature map R4. R2, R3, and R4 are convolved with 3*3 to obtain feature maps T2, T3, and T4 respectively. An AMECA attention module is integrated after each of the last three C3 modules in the neck structure. Feature maps T2, T3, and T4 are used as the original input feature maps, respectively. They are then processed by a global average pooling module and a global max pooling module. The two resulting feature maps are added together to compress spatial information. Subsequently, a 1*1 convolution is used to learn channel attention information. The obtained channel attention information is combined with the original input feature maps to finally obtain specific channel attention feature maps D1, D2, and D3. Output section: Input feature maps D1, D2, and D3 into the YOLOv5-MSTA detection head network to obtain the final detection results; The CBRM operations include Conv, BN, ReLU, and MaxPool.

2. The method for detecting surface defects in lightweight gears based on MSTA-YOLOv5 as described in claim 1, characterized in that, The Mosaic9 data augmentation includes: first, taking a batch of data from the total dataset, randomly selecting 9 images each time, cropping and scaling them at random positions, and synthesizing new images; repeating the above process batch-size times, finally obtaining a batch of new data including batch-size images that have undergone Mosaic9 data augmentation, and then passing it to the neural network for training.

3. The method for detecting surface defects in lightweight gears based on MSTA-YOLOv5 as described in claim 1, characterized in that, The first, second, and third downsampling layers all include a Shuffle_Block(d) module. The Shuffle_Block(d) divides the input features into two branches. The left branch has two convolutional layers: a 3×3 depthwise convolution with a stride of 2 and a 1×1 ordinary convolution. The right branch has three convolutional layers: a 1×1 ordinary convolution, a 3×3 depthwise convolution with a stride of 2, and a 1×1 ordinary convolution. The left and right branches are concatenated using Concat to fuse the features from both sides. Finally, channel shuffling is used to enable communication between the two branches. The second, third, and fourth convolutional normalization layers all include a Shuffle_Block(c) module. This Shuffle_Block(c) module splits each channel into two branches. Based on the principle of reducing model fragmentation, no operations are performed on the left branch, while the right branch contains three convolutional layers: 1×1 ordinary convolutions. 3×3 depthwise convolution, 1×1 ordinary convolution, 1×1 ordinary convolution The three convolutional layers, namely 3×3 depthwise convolution and 1×1 ordinary convolution, have the same input and output channels. The two 1×1 Convolution layers are no longer group convolutions but ordinary convolutions. After the three convolutions, the two branches are concatenated by Concat.

4. The method for detecting surface defects in lightweight gears based on MSTA-YOLOv5 as described in claim 1, characterized in that, The transposed convolution upsampling operation steps include: (1) Fill the input feature map elements with s-1 rows and 0 columns, where s represents the stride of the transposed convolution; (2) Fill the input feature map with kp-1 rows and 0 columns, 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 vertically and horizontally; (4) Perform normal convolution operation, fill with 0, and step 1.

5. The method for detecting surface defects in lightweight gears based on MSTA-YOLOv5 as described in claim 1, characterized in that, The process of the AMECA attention module includes: (1) First, input the feature map. X Feature map X The dimension is ; (2) Input feature map X Spatial feature compression is performed; in the spatial dimension, global average pooling (GAP) is used to obtain 1. 1 Feature map F1 of C; using global max pooling (GMP) yields 1. 1 Feature map F2 of C; (3) F1 and F2 are fused to obtain 1 1 The feature map F3 of C is used to obtain higher-level semantic information; (4) Channel feature learning is performed on the fused feature map F3; through 1*1*1 convolution, the importance between different channels is learned, and the output feature map F4 still has a dimension of 1. 1 C; (5) Obtain F41 by passing the feature map F4 through the σ function; (6) Finally, channel attention is combined by multiplying the channel attention feature map F41 with the original input feature map X channel by channel to finally output the feature map X' with channel attention. Where H, W, and C represent the height, width, and number of channels 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 resulting output feature maps X' are feature maps D1, D2, and D3, respectively.

6. A computer-readable medium, characterized in that, It contains a computer program for performing the method as described in any one of claims 1-5.

Images