Metal surface defect detection method based on YOLOX improved model

Abstract

The invention discloses a metal surface defect detection method based on a YOLOX improved model, and belongs to the technical field of target detection. Comprising the following steps: acquiring a metal surface defect image data set; a YOLOX network model is deployed; modifying the model, wherein the model modification comprises introducing a channel attention and space attention mechanism before entering a feature pyramid, adding an ASFF feature fusion algorithm at the tail of the feature pyramid, modifying a loss function, introducing wavelet convolution, modifying down-sampling at the feature pyramid into wavelet convolution down-sampling, and retaining part of feature details; and training and optimizing the YOLOX improved model, and carrying out defect detection on the metal surface by using the optimized model. According to the method, four modifications are adopted in a YOLOX model, a channel attention and space attention mechanism is introduced, an ASFF feature fusion algorithm is introduced, a loss function is modified, and wavelet convolution is introduced; the defect detection precision of the whole network is improved.

Description

Technical Field

[0001] The invention relates to a metal surface defect detection method based on an improved YOLOX model, and belongs to the technical field of target detection. Background Art

[0002] Metal surface defect detection is crucial in industrial manufacturing and is directly related to product quality and safety. Traditional manual detection methods are not only time-consuming and laborious, but also easily affected by subjective factors and difficult to ensure consistency. With the advancement of computer vision technology, automated surface defect detection has gradually become mainstream. Through image processing and machine learning algorithms, the system can identify tiny defects such as scratches, cracks, and pits, with the advantages of fast detection speed and high accuracy. At present, the application of deep learning, especially convolutional neural networks (CNN), has made the detection of complex surface defects more efficient and accurate. However, there are still some challenges in practical applications, such as how to deal with the impact of different lighting conditions and complex texture backgrounds on defect detection accuracy, and how to improve the real-time and robustness of the algorithm in a large-scale production environment. These issues are the hot spots of current research and the driving force for the development of technology in this field.

[0003] Computer vision object detection technology plays a key role in modern intelligent systems and is widely used in autonomous driving, security monitoring, industrial automation and other fields. The main task of object detection is to identify specific objects in images or videos and accurately locate their positions. Early object detection relied on traditional image processing techniques and hand-designed features such as SIFT and HOG. However, these methods perform poorly when faced with complex scenes and diverse targets.

[0004] With the development of deep learning, object detection methods based on convolutional neural networks (CNNs) have gradually become mainstream. Models such as R-CNN, Fast R-CNN, and Faster R-CNN have significantly improved detection accuracy by introducing region proposal networks. Single-stage detectors such as YOLO and SSD have further improved detection speed and are suitable for real-time applications. Since its launch, the YOLO (You Only Look Once) series of models have been widely used in various computer vision tasks with their end-to-end real-time detection capabilities and high accuracy. YOLOX inherits the fast detection advantages of the YOLO family and introduces some key technical improvements, such as Decoupled Head design, advanced Anchor-free mechanism, and more reasonable label allocation strategy. Admittedly, YOLOX has relatively good detection results on some public data sets, but there is still a lot of room for improvement, and the detection speed of the model still needs to be controlled within a reasonable range. Summary of the invention

[0005] In view of the above problems, the present invention provides a metal surface defect detection method based on the improved YOLOX model, which improves the accuracy of defect detection without significantly reducing the detection speed.

[0006] The present invention provides a metal surface defect detection method based on an improved YOLOX model.

[0007] The technical solution of the present invention is: the present invention provides a metal surface defect detection method based on the YOLOX improved model, comprising the following steps:

[0008] S1: Obtain a dataset of metal surface defect images: Establish a dataset to collect different types of metal surface defects, with a comparable number of defects.

[0009] S2: Deploy the YOLOX network model and select a model with a small number of weights and relatively high accuracy to meet real-time detection requirements.

[0010] S3: Modify the YOLOX model, introduce channel attention and spatial attention (SA), add the ASFF feature fusion algorithm, modify the loss function, and introduce wavelet convolution.

[0011] S4: Train and optimize the improved model of YOLOX, and use the trained improved model of YOLOX to detect metal surface defects.

[0012] In S1, the sizes of the collected metal surface defect images are uniformly cropped to make the image resolutions of the same size, and the collected images are defect labeled according to different types of defects, and finally the data set is divided into a training set and a verification set.

[0013] In S2, YOLOX has multiple official models, each with different weights. Select the s-type weight model to balance detection speed and detection accuracy.

[0014] Furthermore, step S3 specifically includes:

[0015] S31: Introduce channel attention and spatial attention mechanism (SA) to enhance the connection between channels, increase the weight of certain important details and reduce the weight of unimportant details.

[0016] Furthermore, step S31 specifically includes:

[0017] S311: Channel attention and spatial attention take half of the original feature map channels respectively, and calculate the channel attention and spatial attention mechanisms respectively.

[0018] S312: Channel attention adopts two branches, combining the global maximum pooling with the global average pooling operation, and then using a 1*1 convolution kernel to perform additional independent adjustments on the two feature maps. Secondly, the two feature maps are added together to obtain a new feature map, which is then processed using the sigmoid function. Finally, this feature map is multiplied with each of the original feature maps to generate a new feature map. At this time, the number of channels is half of the number of channels of the original feature map.

[0019] S313: Spatial attention uses normalization to normalize the feature map on each channel so that its mean is 0 and its variance is 1. The normalized feature map is weighted. The weighted feature map is processed by the sigmoid function and finally multiplied with the original feature map to generate a new feature map with half the number of channels of the original feature map.

[0020] S314: Concatenate the new feature maps generated by channel attention and spatial attention to obtain a feature map with the same number of channels as the initial feature map. Then shuffle all the channels to disrupt the order of each channel, enhance the relationship expression between each channel, and finally output it to other modules of the model. This attention mechanism has a small number of parameters and does not have much impact on the detection speed of the model.

[0021] Furthermore, step S3 also includes:

[0022] S32: Add the ASFF feature fusion algorithm to enhance the fusion of feature maps at the feature pyramid and better integrate the information of multi-scale feature maps.

[0023] Further, step S32 specifically includes:

[0024] S321: The ASFF feature fusion algorithm is added at the connection between the feature pyramid and the YoloHead. According to the number of yolo heads being 3 and the different detection sizes of the three Yolo Heads, three ASFF modules of different sizes are added.

[0025] S322: Among them, the three ASFF modules use different sizes as input and output respectively. The output at the end of the feature pyramid is also the input of the ASFF module, which is a feature map of three different sizes. The ASFF module performs size transformation on the feature maps of three different sizes.

[0026] S323: After processing the size transformation, each ASFF generates a convolution layer for each scale feature map, and then performs weighted fusion of the three feature maps with the convolution layer to generate a weight map of the corresponding size, and finally multiplies the feature map after the size transformation with the weight map and outputs it. This fusion algorithm has a small number of parameters and will not significantly delay the detection speed of the model.

[0027] Furthermore, step S3 also includes:

[0028] S33: Modify the loss function and replace the original IOU loss with EIOU loss. Compared with IOU loss, EIOU loss increases the center distance error between the predicted box and the actual box, and increases the difference in the aspect ratio between the predicted box and the actual box, making the loss calculation more accurate.

[0029] Furthermore, step S3 also includes:

[0030] S34: Introduce wavelet convolution, replace the downsampling convolution in the feature pyramid part of the model with wavelet convolution to achieve the downsampling function. Wavelet convolution has the function of retaining detail features and reduces the problem of information loss when the feature map is resized. The number of parameters introduced by wavelet convolution is small, and it has little effect on the detection speed of the model.

[0031] Furthermore, step S4 specifically includes:

[0032] S41: Use YOLOX to improve the model for training and adjust the model hyperparameters based on the training results.

[0033] S42: Save the results and weight values ​​after each training, evaluate the detection accuracy of different weights, and find the optimal weight.

[0034] The beneficial effects of the present invention are as follows: the present method adopts four modifications in the YOLOX model, introduces channel attention and spatial attention mechanisms, introduces the ASFF feature fusion algorithm, modifies the loss function, and introduces wavelet convolution. The extraction and fusion of different scale features of the entire network model are enhanced, and the accuracy of defect detection is improved without significantly reducing the detection speed. BRIEF DESCRIPTION OF THE DRAWINGS

[0035] Figure 1 This is a schematic diagram of the detection method process in an embodiment of the present invention;

[0036] Figure 2 This is a schematic diagram of the structure of the YOLOX improved model;

[0037] Figure 3 It is a schematic diagram of the structure of channel attention and spatial attention mechanism (SA).

[0038] Figure 4 It is a schematic diagram of the structure of the ASFF fusion algorithm;

[0039] Figure 5 is the true label map of a certain defect in the dataset;

[0040] Figure 6This is a comparison chart of the detection effects of the original YOLOX model and the improved YOLOX model;

[0041] Figure 7 It is a heat map comparison between the original YOLOX model and the improved YOLOX model;

[0042] Figure 5-Figure 7 The background is the same picture. DETAILED DESCRIPTION

[0043] In order to make the purpose, technical solution and advantages of this application more clear, the present application is further described in detail with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain this application and are not used to limit this application.

[0044] It should be noted that the YOLOX models mentioned in the present invention are all based on the official structure. When the model structure is the same, the model depth must also be the s-type weight mentioned in the text. YOLOX models of other depths can also use the improved method mentioned in the present invention, but it is not related to the present invention. Similarly, the wavelet convolution module mentioned in the present invention includes a wavelet transform. And this module also has depth, and the selection of a wavelet convolution module of other depths is not related to the present invention.

[0045] The present invention does not mention the hardware facilities for training, because during the training process of the model, the results produced by each complete training will have a certain error value. Therefore, any hardware facilities can be used to implement this method, but the error of the results needs to be understood in advance.

[0046] The number of channels mentioned in the present invention is the number of layers of the feature map in each link. The size mentioned in the present invention is the resolution of the image. For example, the image format used in the present invention is 640*640, but this method does not limit the image size.

[0047] The hyperparameters such as the number of training epochs and the optimizer mentioned in the present invention can be selected by yourself and are not limited to the method in the present invention. However, it needs to be ensured that the hyperparameters of the YOLOX model before and after the improvement are the same.

[0048] The present application provides a metal surface defect detection method based on the YOLOX improved model, the method comprising:

[0049] S1: Establish a data set to collect different types of metal surface defects, with a similar number of defects.

[0050] The S1 step includes two sub-steps:

[0051] S11: Determine the type of metal surface defects. This example uses the metal surface defect dataset published by Northeastern University, which has 6 types of defects, namely crazing, inclusion, patches, pitted surface, rolled-in_scale, and scratches. The number of each defect is similar, with a total of 1,800 images. This dataset is divided into a training set, a validation set, and a test set, with a ratio of approximately 8:1:1.

[0052] S12: Change the format of the collected data set to yolo. This data set already contains labels. If not, labeling is required. The image size of this data set is 640*640, ensuring that the image size is consistent.

[0053] S2: Deploy the YOLOX network model and select a model with a small number of weights and relatively high accuracy to meet real-time detection requirements.

[0054] The S2 step consists of three sub-steps:

[0055] S21: You can download the official YOLOX model or use one written by others. You need to ensure that the model structure is consistent with the official one.

[0056] S22: YOLOX officially launched a variety of models with different weight sizes. These different models are roughly the same in structure. The biggest difference is the depth. The size of the weight ranges from a few MB to hundreds of MB. In this example, the weight code s is selected, which is 34MB in size. It takes into account both detection speed and detection performance, and is more suitable for horizontal and vertical research.

[0057] S23: Put the processed data set into the trusted path so that the YOLOX model can run training directly.

[0058] S3: Modify the YOLOX model, introduce channel attention and spatial attention, add the ASFF feature fusion algorithm, modify the loss function, introduce wavelet convolution, and retain some feature details.

[0059] The S3 step specifically includes:

[0060] S31: Introduce channel attention and spatial attention mechanisms to enhance the association between channels, increase the weight of certain important details and reduce the weight of unimportant details.

[0061] The sub-step S31 includes four sub-steps:

[0062] S311: Add channel attention and spatial attention (SA) at the connection between the backbone network and the feature pyramid, such as Figure 4The symbols SA in the figure are shown as follows. The channel attention and spatial attention take half of the original feature map channel respectively, and calculate the channel attention and spatial attention mechanisms respectively, as shown in Figure 4 As shown in , the number of channels of the input (X1) is divided into G groups, and the number of channels in each group is C / G. Figure 4 C, h, and w are the number of channels of the feature map, and the height and width of the feature map size. Channel attention and spatial attention use C / 2 channels respectively, and each group uses C / (2*G) channels.

[0063] S312: Channel attention uses two branches, combining global maximum pooling with global average pooling operations. After the two pooling operations, the specification of the feature map becomes C / (2*G), 1, 1. Then use a 1*1 convolution kernel to adjust the two different feature maps independently. Then add the two feature maps to generate (X9). Each channel has a corresponding feature map. Then, the feature map is weighted and processed using the sigmoid function to generate (X10). The feature map value is limited to a certain range. Finally, this feature map (X10) is multiplied with the original feature map (X3) to generate a new feature map (X11). Its specification is C / (2*G). h, w The feature map (X10) can be regarded as a weight map of the same specification as (X3). At this time, the number of channels is half of the number of channels in the original feature map.

[0064] S313: Spatial attention uses normalization operations to normalize the feature map on each channel so that the mean of all pixels is 0 and the variance is 1. The normalized feature map is weighted and processed with the sigmoid function to generate the feature map (X14), and then the generated (X14) is multiplied one by one with the feature map (X12) in the form of pixels to generate the feature map (X15), whose specifications are C / (2*G), h, w. The feature map (X14) can be viewed as a weight map of the same specifications as (X12), and the number of channels is half of the number of channels of the original feature map.

[0065] S314: Concatenate the new feature maps (X11) and (X15) generated by channel attention and spatial attention to obtain a feature map (X16) with the same number of channels as the original feature map, with the specification of C / G, h, w. Then shuffle all the channels in order, disrupt the order of each channel, enhance the relationship expression between each channel, and finally output it to other modules of the model. This module enhances the relationship flow between channels in the feature map and the spatial relationship on a single channel, enhances the association between channels, increases the weight of some important details and reduces the weight of unimportant details. The number of parameters of this attention mechanism is small, and there is not much loss on the detection speed of the model.

[0066] Furthermore, step S3 also specifically includes:

[0067] S32: Add the ASFF feature fusion algorithm to enhance the fusion of feature maps at the feature pyramid and better integrate the information of multi-scale feature maps.

[0068] Furthermore, step S32 specifically includes three sub-steps:

[0069] S321: According to the number of yolo heads being 3 and the different detection sizes of the three Yolo heads, three ASFF feature fusion algorithm modules are added at the connection between the feature pyramid and the YoloHead, such as Figure 4 shown.

[0070] S322: The three ASFF modules use different sizes as input and output respectively. The output at the end of the feature pyramid is also the input of the ASFF module, which is a feature map of three different sizes. The three feature maps are input into the three ASFF modules respectively, such as Figure 3 As shown in Figure 2, X1, X2, and X3 represent feature maps of different sizes. X1-1 means changing the size of X1 to X 1 X1-2 means changing the size of X1 to X 2 X1-3 means changing the size of X1 to X 3 The size of each ASFF module is the same as the others. After receiving the input, each ASFF module adds the three input feature maps together. The three different ASFFs correspond to three different size processing processes.

[0071] S323: Figure 3 As shown, the feature map size processed by the first ASFF module is X 1 The feature map size processed by the second ASFF module is X 2 The third ASFF module processes feature maps of size X 3 . After processing the size transformation, each ASFF generates a convolution layer for each scale feature map, and then performs weighted fusion of the three feature maps with the convolution layer to generate a feature map of the corresponding size. At this time, this feature map can be regarded as a weight map. Finally, the feature map after the size transformation is multiplied by the weight map and output. This module enhances the fusion of feature maps at the feature pyramid and better integrates the information of multi-scale feature maps. The number of parameters of this fusion algorithm is small and will not significantly delay the detection speed of the model.

[0072] Furthermore, step S3 specifically includes:

[0073] S33: Modify the loss function and replace the original IOU loss with EIOU loss. Compared with IOU loss, EIOU loss increases the center distance error between the predicted box and the actual box, and increases the difference in the aspect ratio between the predicted box and the actual box, making the loss calculation more accurate.

[0074] Furthermore, step S3 specifically includes:

[0075] S34: Introduce wavelet convolution. Wavelet convolution contains wavelet transform. Different wavelet transform times can be selected in this module, that is, the depth of the module. This method selects a module depth of 3, which means that this module has a total of three wavelet transforms. Replace the downsampling convolution in the feature pyramid part of the model with wavelet convolution to achieve the downsampling function. Wavelet convolution has the function of retaining detail features and reduces the problem of information loss when the feature map is resized. The number of parameters introduced by wavelet convolution is small, and it has little effect on the detection speed of the model.

[0076] S4: Train and optimize the improved model of YOLOX, and use the trained improved model of YOLOX to detect metal surface defects.

[0077] In step S4, the following sub-steps are included:

[0078] S41: Use YOLOX to improve the model for training, set the number of training epochs to 300, use mosaic dataset enhancement, select the optimizer as sgd, select the learning rate descent method as cos, set the model's learning rate range to (0.01-0.0001), and set the batchsize to 8. The system used in this example is ubuntu22.04, cuda11.8, cudnn8.9.7, the cpu used is 13700 (server 8 cores), the graphics card used is 4090D, and the memory is 16G. Use the pytorch deep learning framework.

[0079] S42: After completing all epochs of training, save the results and weight values ​​after each training, and evaluate the detection accuracy of different weights. Find the set of weight values ​​with the lowest loss value of the validation set, which is the optimal weight value, and use the test set to evaluate the map accuracy of this weight. Calculate the detection speed. The main feature of this method is that it enhances feature fusion, especially feature fusion at different scales, while also retaining more details of the feature map, which has a better detection effect on the defects of this dataset.

[0080] YOLOX improved model effect analysis: Figure 5 , as shown in Figures 6 and 7. Figure 5 is a picture with real labels, Figure 6This is a comparison chart of the detection effects of the original YOLOX model and the YOLOX improved model. The left side is the detection effect of the original YOLOX model, and the right side is the detection effect of the YOLOX improved model. It can be seen that in this picture, Figure 5 Compared with the original YOLOX model, the original YOLOX model missed two targets, while the improved YOLOX model correctly identified one more target than the original YOLOX model, indicating that the prediction ability of the improved YOLOX model has been improved. Figure 7 The heat map of Figure 7 As shown, the left side is the heat map generated by the original YOLOX model on this picture, and the right side is the heat map generated by the YOLOX improved model on this picture. The characteristic of the heat map is that the more highlights there are, the more attention the model pays to this area. It can be seen that the highlights are from Figure 6 The change from the left side to the right side corresponds to the position of the defect in this picture, which shows that the YOLOX improved method has improved the feature extraction capability here and increased the attention.

[0081] The specific implementation modes of the present invention are described in detail above in conjunction with the accompanying drawings, but the present invention is not limited to the above implementation modes, and various changes can be made within the knowledge scope of ordinary technicians in this field without departing from the purpose of the present invention.

Claims

1. A metal surface defect detection method based on the YOLOX improved model, characterized in that: The following steps are involved: S1: Acquire metal surface defect image dataset; S2: Deploy the YOLOX network model; S3: Modify the YOLOX model, introduce channel attention and spatial attention, add the ASFF feature fusion algorithm, modify the loss function, and introduce wavelet convolution; S4: Train and optimize the improved model of YOLOX, and use the trained improved model of YOLOX to detect metal surface defects.

2. The method for detecting metal surface defects based on the improved YOLOX model according to claim 1, characterized in that: In step S1, the size of the collected metal surface defect images is uniformly cropped to make the image resolution of the same size, and the collected images are defect labeled according to different types of defects. Finally, the data set is divided into a training set and a verification set.

3. The method for detecting metal surface defects based on the improved YOLOX model according to claim 1, characterized in that: Step S3 specifically includes: S31: Channel attention and spatial attention mechanisms are introduced to enhance the association between channels, increase the weight of certain important details and reduce the weight of unimportant details.

4. The metal surface defect detection method based on the YOLOX improved model according to claim 3 is characterized in that: Step S31 specifically includes: S311: channel attention and spatial attention take half of the original feature map channel respectively, and calculate the channel attention and spatial attention mechanisms respectively; S312: Channel attention uses two branches, combining global maximum pooling with global average pooling operations, and then using a 1*1 convolution kernel to perform additional independent adjustments on the two feature maps. Next, the two feature maps are added together to obtain a new feature map, which is then processed using a sigmoid function. Finally, this feature map is multiplied with each of the original feature maps to generate a new feature map, where the number of channels is half of the original number of channels. S313: The spatial attention adopts the normalized form, normalizing the feature map on each channel to make its mean 0 and variance 1; weighting the normalized feature map; the weighted feature map is processed by the sigmoid function, and finally multiplied with the original feature map to generate a new feature map, in which the number of channels is half of the number of channels of the original feature map; S314: Concatenate the new feature maps generated by channel attention and spatial attention to obtain a feature map with the same number of channels as the initial feature map; then shuffle all the channels in order to disrupt the order of each channel, enhance the relationship expression between each channel, and finally output it to other modules of the model.

5. The method for detecting metal surface defects based on the improved YOLOX model according to claim 1, characterized in that: Step S3 also includes: S32: Add the ASFF feature fusion algorithm to enhance the fusion of feature maps at the feature pyramid and better integrate the information of multi-scale feature maps.

6. The method for detecting metal surface defects based on the improved YOLOX model according to claim 5, characterized in that: Step S32 specifically includes: S321: adding the ASFF feature fusion algorithm to the connection between the feature pyramid and the yolo head, adding the ASFF module, and adding three ASFF modules of different sizes according to the number of yolo heads being 3 and the different detection sizes of the three Yolo heads; S322: The three ASFF modules use different sizes as input and output respectively. The output at the end of the feature pyramid is also the input of the ASFF module, which is a feature map of three different sizes. The ASFF module performs size transformation on the feature maps of three different sizes. S323: After processing the size transformation, each ASFF generates a convolution layer for the feature map of each scale, and then performs weighted fusion on the three feature maps and the convolution layer to generate a weight map of the corresponding size. Finally, the feature map after the size transformation is multiplied by the weight map and output.

7. The method for detecting metal surface defects based on the improved YOLOX model according to claim 1, characterized in that: Step S3 also includes: S33: Modify the loss function and replace the original IOUloss with EIOUloss.

8. The method for detecting metal surface defects based on the improved YOLOX model according to claim 1, characterized in that: Step S3 also includes: S34: Introduce wavelet convolution and replace the downsampling convolution in the feature pyramid part of the model with wavelet convolution to achieve the downsampling function.

9. The method for detecting metal surface defects based on the improved YOLOX model according to claim 1, characterized in that: Step S4 specifically includes: S41: Use YOLOX to improve the model for training and adjust the model hyperparameters based on the training results; S42: Save the results and weight values ​​after each training, evaluate the detection accuracy of different weights, and find the optimal weight.

Images