Attention-enhanced YOLOv8 bearing defect detection method based on dynamic up-sampling and frequency domain cooperation

Abstract

The invention relates to the field of computer vision and deep learning, in particular to a bearing defect detection method based on YOLOV8 improvement. In existing bearing flaw detection, due to the fact that bearing flaws are various in variety, different in shape, small in size and difficult to perceive, a traditional target detection method is often difficult to achieve an ideal detection effect. In order to improve the accuracy and efficiency of bearing defect detection, important improvement is carried out on the basis of YOLOV8, WTDown is introduced to solve the problem of information loss existing in a model in the down-sampling process, a Shuffle Attention module is introduced to enable the model to pay attention to important feature information, a DyUpsample module is introduced to improve the learning ability of the model, and the accuracy and efficiency of bearing defect detection are improved. The improved algorithm can greatly improve the accuracy of non-batch flaw detection.

Description

Technical Field

[0001] The present invention relates to the field of bearing surface defect detection, and in particular to a bearing surface defect detection method improved based on YOLOV8s. Background Art

[0002] This invention relates to the fields of computer vision and deep learning, and more specifically, to an improved bearing surface defect detection method based on YOLOv8s. Conventional bearing defect detection methods often struggle to achieve optimal detection results due to the wide variety and diverse forms of bearing defects, their often small size, and their subtle nature. To improve the accuracy and efficiency of bearing defect detection, this invention makes significant improvements based on YOLOv8s.

[0003] YOLOv8s, a real-time object detection algorithm, has demonstrated outstanding performance in multiple fields. However, YOLOv8's use of strided convolution for downsampling leads to information loss, and its use of interpolation for feature upsampling renders the upsampler incapable of learning. This also limits the model's capabilities by making it impossible to distinguish between important and minor features during forward propagation. To address this, the present invention introduces three improved modules: the Haar frequency-domain converter (WTDown), the dynamic upsampler (DU), and the shuffle attention (SA) mechanism, to comprehensively enhance YOLOv8s' performance in bearing defect detection.

[0004] The Haar frequency-domain converter WTDown is a frequency-domain-based lossless downsampling module that downsamples the input by 1 / 2 to obtain a low-frequency component A, a horizontal high-frequency component H, a vertical high-frequency component V, and a diagonal high-frequency component D. Through frequency-domain conversion, the high-frequency and low-frequency features of the input are retained while reducing the input resolution, achieving lossless downsampling.

[0005] The Dynamic Upsampler is a learnable upsampling module. Interpolation-based upsampling methods, such as the nearest neighbor method and bilinear interpolation, use fixed calculations and cannot leverage features learned from the data during upsampling, limiting the model's learning capabilities. The Dynamic Upsampler, on the other hand, introduces learnable parameters that can be dynamically adjusted based on data characteristics, improving the model's learning capabilities and performance.

[0006] Shuffle Attention combines channel shuffling and attention weight learning to achieve a novel feature interaction enhancement structure. It employs channel segmentation in the feature dimension and uses cross-group channel swapping to break redundancy between feature channels. This constructs multi-scale feature interaction paths, enabling the full integration of feature information from different receptive fields. In bearing defect detection, Shuffle Attention helps the model focus more closely on defect areas and less on background areas, thereby improving detection accuracy.

[0007] In summary, this paper significantly improves the performance of YOLO V8s in bearing defect detection by integrating three improved modules: the Haar frequency domain converter WTDown, the dynamic upsampler DynamicUpsampler, and the attention mechanism Shuffle Attention. This technological innovation not only improves detection accuracy and efficiency but also provides strong support for quality control during bearing production. Summary of the Invention

[0008] The technical problem to be solved by the present invention is that defects in the bearing production process are of many types, various shapes, and small sizes, making them difficult to detect. To address this problem, a bearing defect detection method based on an improvement of YOLOV8s is proposed. This method introduces three improved modules: the Haar frequency domain converter WTDown, the dynamic upsampler Dynamic Upsampler, and the attention mechanism ShuffleAttention, which comprehensively improve the performance of YOLOV8s in bearing defect detection. The method was verified on a self-built bearing defect detection dataset. The mAP50 of the native YOLOV8s algorithm was 83.3%, and the mAP50 of the improved method was 86.4%. The proposed method improved the mAP50 index of the native YOLOV8s algorithm by 3.1%, greatly increasing the detection accuracy of bearing defects and having extremely high practical application value.

[0009] To solve the technical problem of bearing defect detection difficulties, the present invention proposes an attention-enhanced YOLOv8 bearing defect detection method based on dynamic upsampling and frequency domain collaboration, which includes the following steps: Step 1: The Haar frequency-domain converter (WTDown) module is introduced into the YOLO V8s model. In traditional convolutional neural networks, feature downsampling is typically achieved using strided convolution or pooling operations. However, this inevitably results in feature information loss, which affects model performance. Frequency-domain conversion downsampling separates high-frequency and low-frequency features, preserving feature information while downsampling, achieving lossless downsampling.

[0010] Step 2: Use the Dynamic Upsampler to replace interpolation-based upsampling methods. Traditional interpolation-based upsampling methods, such as the nearest neighbor method and bilinear interpolation, use fixed calculation methods and cannot utilize features learned from the data during upsampling, limiting the model's learning capabilities. The Dynamic Upsampler introduces adaptively learnable parameters. During model training, these parameters can learn the model's characteristics, thereby improving model performance.

[0011] Step 3: The Shuffle Attention module was introduced into the YOLO V8s model. Shuffle Attention is a special self-attention mechanism designed to address the computational complexity and memory consumption issues faced by traditional self-attention models when processing high-resolution feature maps. It combines channel grouping, shuffling, and the attention mechanism to form a novel feature interaction enhancement structure.

[0012] Step 3.1 Shuffle Attention uses periodic channel partitioning (such as g×g partitioning) in the feature dimension, breaks the redundancy between feature channels through cross-group channel exchange operations, constructs multi-scale feature interaction paths, and fully integrates feature information from different receptive fields.

[0013] Step 3.2 Shuffle Attention introduces a lightweight convolution kernel (such as 1×1 convolution) to generate a channel attention mask, and dynamically adjusts the importance weight of each channel feature through the Sigmoid activation function to effectively suppress background noise interference.

[0014] Step 3.3 Shuffle Attention adopts a dual-branch parallel structure. The main branch retains the original feature flow, and the attention branch generates a weight mask, ultimately achieving adaptive enhancement of feature dimensions through weighted fusion.

[0015] Step 4: We verified the model performance on a self-built bearing defect detection dataset. The mAP50 of the native YOLOV8s algorithm was 83.3%, while the mAP50 of the improved method was 86.4%. The proposed method improved the mAP50 index by 3.1% over the native YOLOV8s algorithm. The experimental results are shown in the following table: Table 1 Ablation test results of improved model

[0016] The beneficial effects of adopting the above technical solution are: To address the diverse, complex, and small-sized nature of bearing defects, we introduced the Haar frequency domain converter (WTDown), the dynamic upsampler (Dynamic Upsampler), and the attention mechanism (Shuffle Attention) based on YOLOV8s. WTDown achieves lossless downsampling through frequency conversion, the dynamic upsampler introduces learnable parameters to improve the model's learning capabilities, and Shuffle Attention enables the model to extract more important features. This proposed method achieves a 3.1% improvement in the mAP50 metric over the native YOLOV8s algorithm, significantly increasing the accuracy of bearing defect detection and possessing high practical application value. BRIEF DESCRIPTION OF THE DRAWINGS

[0017] Figure 1 The original YOLOV8s network structure diagram provided by the embodiment of the present invention; Figure 2 This is a structural diagram of the Conv, SPPF, BootleNeck, and C2f modules in the YOLOV8s algorithm provided in an embodiment of the present invention; Figure 3 A schematic diagram of the Shuffle Attention method provided by an embodiment of the present invention; Figure 4 The structure diagram of the SAC2f module combined with Shuffle Attention provided by the example of the present invention; Figure 5 The SAC2f network structure diagram is added to the YOLOV8s model provided in an embodiment of the present invention; Figure 6 Schematic diagram of the WTDown module based on Haar wavelet transform provided in an embodiment of the present invention; Figure 7 A comparison chart of the results of Haar wavelet transform and pooling operation is provided for the embodiment of the present invention; Figure 8 A network structure diagram of an embodiment of the present invention that adds Haar wavelet transform WTDown to the YOLOV8s model; Figure 9 A schematic diagram of a Dynamic Upsampler according to an embodiment of the present invention; Figure 10 A network structure diagram of an embodiment of the present invention that adds a dynamic upsampling module DyUpsample to the YOLOV8s model; Figure 11A network structure diagram showing the addition of Shuffle Attention, WTDown, and Dynamic Upsampler to the YOLOV8s model provided in an embodiment of the present invention; Figure 12 A flow chart of bearing defect detection based on improved YOLOV8s provided in an embodiment of the present invention; Figure 13 A sample diagram of a bearing defect provided by an embodiment of the present invention; Figure 14 The results of detecting bearing defects using the original YOLOV8s algorithm provided by the embodiment of the present invention; Figure 15 The results of detecting bearing defects using the improved YOLOV8s algorithm provided by the embodiment of the present invention; DETAILED DESCRIPTION

[0018] The following embodiments of the present invention are described in further detail with reference to the accompanying drawings and examples. The following examples are used to illustrate the present invention but are not intended to limit the scope of the present invention.

[0019] This application uses a specific example to simulate the bearing defect data collected during the bearing production process to verify the effectiveness of the proposed prediction method. The following is an implementation example: This paper takes the bearing defect detection algorithm as an example for research. The network structure diagram of the YOLOV8s algorithm is as follows: Figure 1 As shown in the figure, the input image passes through BackBone to obtain feature maps of different scales. After passing through Neck, the feature maps are integrated with feature information of other scales. Finally, the feature maps of different scales are sent to the detection head to obtain the prediction results of the model. The YOLOV8s network consists of Conv, C2f, SPPF and Detect modules. Backbone contains 5 Conv modules, 4 C2f modules and 1 SPPF module. Neck contains 2 Upsample modules, 4 C2f modules, 4 Concat modules and 2 Conv modules. Head contains 3 Detect modules. The structure of C2f, SPPF, Conv and Detect modules is shown in the figure. Figure 2 As shown in the figure, the Conv module consists of a convolutional layer, a BatchNorm2d layer, and a SiLU layer; the Bottleneck module consists of two Conv modules with a residual connection structure; the C2f module divides the input feature map into multiple non-overlapping feature maps, and then combines different features after passing them through different Bottleneck layers to extract network features of different depths.

[0020] The principle of the Shuffle Attention module introduced in this invention is as follows Figure 3 As shown in Figure 2, the input is a feature map X, whose dimensions include width (W), height (H), and number of channels (C). First, the input feature map is grouped by channel dimension, dividing the number of channels C into several groups to prepare for subsequent feature processing in different branches. The grouped features are further split and fed into the channel attention module and spatial attention module, respectively. After integrating the features of different groups, the channels are positionally transformed so that each group incorporates information from other groups.

[0021] The SAC2f introduced in the present invention is as follows Figure 4 As shown in Figure 1, the Shuffle Attention mechanism is added to the original Conv and BottleNeck modules of YOLO V8s to obtain SAConv and SABottleNeck. In SAConv, ShuffleAttention replaces the traditional convolutional layers, allowing SABottleNeck and SAC2f to selectively focus on important feature information, improving model performance. Figure 5 The network structure diagram after adding SAC2f to the YOLOV8s model is shown.

[0022] The principle of the WTDown module based on Haar wavelet transform introduced in the present invention is as follows: Figure 6 The Haar wavelet transform is a frequency-domain lossless downsampling module that downsamples the input by 1 / 2 to produce a low-frequency component A, a horizontal high-frequency component H, a vertical high-frequency component V, and a diagonal high-frequency component D. Through frequency-domain conversion, the high-frequency and low-frequency features of the input are retained while reducing the input resolution, achieving lossless downsampling. Figure 7 This paper presents a comparison of the downsampling results of Maxpool, Avgpool, and the Haar wavelet transform, as well as their inverse upsampling results. It can be seen that Maxpool only retains the maximum value within the pooling window, and after de-pooling via indexing, a significant amount of information is lost. Mean pooling only retains the mean of each pixel within the window during downsampling. After de-pooling, image detail information is lost, and image clarity is significantly reduced. The Haar wavelet transform preserves both high-frequency and low-frequency features of the image, and during the inverse transform, the high-frequency and low-frequency components are used to achieve lossless image restoration. This demonstrates the lossless nature of the Haar wavelet transform. Using the Haar wavelet transform-based WTDown module preserves all information during the downsampling process, avoiding feature information loss. Figure 8 This is the structural diagram after adding the WTDown module to YOLOV8s.

[0023] The Dynamic Upsampler module introduced in the present invention is as follows Figure 9 As shown in the figure, the input feature map X (of dimensions H × W × C) is first processed by the sampling point generator. After normalization and pixel shuffle operations, the sampling point offsets O (of dimensions sH × sW × 2g) and G (of dimensions sH × sW × 2g) are obtained as the initial sampling coordinates. These two are added together to obtain the sampling coordinates. Combined with the original input X, a grid sampling operation is performed to finally output the upsampled feature map X' (of dimensions sH × sW × C). This achieves dynamic upsampling and flexible adjustment of feature map resolution, providing feature representations at richer scales for subsequent networks. Learnable parameters are also introduced to enhance the model's learning capabilities. Figure 8 This is the structural diagram after adding the Dynamic Upsampler module to YOLOV8s. The DyUpsample in the figure represents the Dynamic Upsampler module.

[0024] The bearing defect detection algorithm proposed in this invention is based on the improved YOLOV8s. Figure 11 In the Backbone part, the deep C2f module is replaced with the SAC2f module based on the Shuffle Attention mechanism, and the convolutional downsampling with stride is replaced with the WTDown downsampling module based on the Haar wavelet transform. This enhances Backbone's feature extraction capabilities and prevents information loss caused by downsampling. In the Neck part, the SAC2f and DyUpsample modules are introduced to improve the model's learning ability while focusing on important features.

[0025] The flow chart of the bearing defect detection algorithm based on improved YOLOV8s proposed in this invention is as follows: Figure 12 As shown in the figure, it mainly includes data collection, data set division, model training, model verification and model deployment. The specific steps are as follows: Step 1: Construct a bearing defect dataset. Collect 3293 pieces of bearing defect data from the Internet. The image resolution is 640x640. There are eight defect categories: casting_burr, polished_casting, burr, crack, pit, scratch, strain, and unpolished_casting.

[0026] Step 2: Split the 3293 data points into a training set and a validation set for training the model and verifying model performance, respectively. The training set contains 2561 data points, and the validation set contains 732 data points.

[0027] In step 3, model training and model verification are performed on the divided dataset. During the training process, the initial learning rate is 0.01, the batch training size of the model is 32, the number of training rounds is 100, and the resolution of the image during training is 640x640.

[0028] Step 3.1: Build the experimental environment required for model training and verification. The configured experimental environment is shown in Table 2.

[0029] Table 2 Experimental configuration information

[0030] Step 3.2: Train the improved models on the constructed datasets and verify the effectiveness of each model on the validation set. The experimental results are shown in Table 3. It can be seen that the modules introduced in this invention significantly improve the detection performance of the YOLOV8s model in the field of bearing defect detection. After adding the WTDown module to the YOLOV8s model, the mAP50 index increased by 1.6%. After adding the DyUpsample module to the YOLOV8s model, the mAP50 index increased by 1.4%. After adding Shuffle Attention (SA) to the YOLOV8s model, the mAP50 index increased by 1.4%. After adding the DyUpsample and SA modules to the YOLOV8s model, the mAP50 index increased by 1.7%. After adding WTDown and DyUpsample to the YOLOV8s model, the mAP50 index increased by 2.2%. After adding WTDown and SA to the YOLOV8s model, the mAP50 index increased by 2.2%. Compared with the original YOLOV8s model, the improved YOLOV8s model has a 3.1% increase in mAP50. Overall, it can be seen that the improved YOLOV8s model has better bearing defect detection performance.

[0031] Table 3 Comparison of improved model performance

[0032] Step 3.3 Verify the effects of each model on the validation set. The experimental results are shown in Table 4, which compares the mAP50 indicators of the proposed method with those of YOLOV3-tiny, YOLOV5u-s, YOLOV6u-s and other models on the validation set. It can be seen that compared with the best-performing YOLOV9s, the mAP50 of the proposed method is improved by 0.6%, and compared with the baseline model YOLOV8s, the mAP50 is improved by 2.8%.

[0033] Table 4 Performance comparison of the improved model with other models

[0034] Step 4: After verifying the model performance on the validation set, deploy the model for inference. Figure 13 This is an example of bearing defects. The bearing in this sample image contains three defects: Casting_burr, scratch, and crack. Figure 14 This is the detection result of YOLOV8s. It can be seen that YOLOV8s did not detect the scratch defect and repeatedly detected the crack defect. Figure 15 is the test result of the present invention, and Figure 13 The comparison shows that the detection results of the present invention are consistent with the label information. The experimental results show that the algorithm proposed in the present invention can effectively improve the detection accuracy of bearing defects and meet the needs of industrial production.

Claims

1. A bearing defect detection method based on an improved YOLOV8, which is characterized by improving the bearing defect detection capability of YOLOV8. The specific steps include: Step 1: The Haar frequency domain converter WTDown module is introduced into the YOLOV8s model. In traditional convolutional neural networks, convolution or pooling operations with strides are usually used to achieve feature downsampling. However, this inevitably leads to feature information loss, which affects the performance of the model. Frequency domain conversion achieves downsampling by separating high-frequency and low-frequency features. It can retain feature information while downsampling, thus achieving lossless downsampling. Step 2: Use the Dynamic Upsampler to replace the interpolation-based upsampling method. Traditional interpolation-based upsampling methods such as the nearest neighbor method and bilinear interpolation have fixed calculation methods and cannot utilize features learned from the data during upsampling, which limits the learning ability of the model. The dynamic upsampler introduces adaptive learning parameters. During the model training process, the relevant parameters can learn the characteristics of the model, thereby improving the performance of the model. Step 3: The Shuffle Attention module is introduced into the YOLOV8s model. Shuffle Attention is a special self-attention mechanism that aims to solve the computational complexity and huge memory consumption problems faced by traditional self-attention models when processing high-resolution feature maps. It forms a new feature interaction enhancement structure through channel grouping shuffling and attention mechanism. Step 3.1 Shuffle Attention uses periodic channel partitioning (such as g×g partitioning) in the feature dimension, breaks the redundancy between feature channels through cross-group channel exchange operations, constructs multi-scale feature interaction paths, and fully integrates feature information from different receptive fields. Step 3.2 Shuffle Attention introduces a lightweight convolution kernel (such as 1×1 convolution) to generate a channel attention mask, and dynamically adjusts the importance weight of each channel feature through the Sigmoid activation function to effectively suppress background noise interference. Step 3.3 Shuffle Attention adopts a dual-branch parallel structure. The main branch retains the original feature flow, and the attention branch generates a weight mask. Finally, the adaptive enhancement of the feature dimension is achieved through weighted fusion. Step 4: The model performance was verified on a self-built bearing defect detection dataset. The mAP50 of the native YOLOV8s algorithm was 83.3%, and the mAP50 of the improved method was 86.4%. The proposed method improved the mAP50 index by 3.1% compared with the native YOLOV8s algorithm.

Images