Metal surface small target defect detection method based on adaptive spatial attention mechanism

Abstract

The application provides a metal surface small target defect detection method based on an adaptive spatial attention mechanism, which can better classify small defect targets on a metal surface without increasing computing resources, ensures that the training retains global information while paying more attention to important information of input features, and improves the recognition accuracy of the model without reducing the reasoning speed. The adaptive spatial attention mechanism is embedded into the existing network model, and in the adaptive spatial attention mechanism, adaptive pooling operation is performed on the output features; the feature map output by the adaptive pooling operation is spliced along the channel to obtain a feature map that is most representative and takes into account the context information, and then the empty spatial attention mechanism is used to increase the receptive field and improve the dependency relationship in the defect feature space; finally, the dimensionality of the feature map is reduced through channel splitting and merging operation, which can better retain the extracted features and reduce the computing resources.

Description

Technical Field

[0001] This invention relates to the field of deep learning target detection technology, specifically a method for detecting small target defects on metal surfaces based on an adaptive spatial attention mechanism. Background Technology

[0002] During the production of metal products, due to factors such as raw materials, hardware equipment, manufacturing processes, and the production environment, in addition to potential initial cracks inside, surface defects such as scratches, abrasions, pinholes, rust, oil stains, and dimensional deviations inevitably appear on the surface of metal products. Traditional defect detection methods based on image recognition suffer from serious false negative rates and poor accuracy, and no longer meet the requirements of production lines for high efficiency and high precision.

[0003] Technicians have developed a deep learning-based surface defect detection technology that continuously improves its detection performance as the scale of defect sample data increases. It can directly extract defect features from data samples, significantly reducing workload compared to manually designed features in traditional machine vision technology, and exhibits strong robustness against multiple types of defects. Many existing deep learning-based surface defect detection technologies utilize attention mechanisms to optimize convolutional neural networks and improve detection accuracy. Attention mechanisms primarily mimic the characteristics of human vision, suppressing unimportant background information and focusing on important information within a region. However, various problems exist in practical applications:

[0004] (1) Squeeze-and-Excitation Networks (SENet) was the first to use channel attention mechanisms and channel fusion strategies to suppress unimportant channel information. The subsequent convolutional block attention module (CBAM) built on SENet by considering the role of spatial information and combining channel attention mechanisms with spatial attention mechanisms. However, they all ignored the interaction between channel and spatial information and lost cross-dimensional information.

[0005] (2)Liu,Y.,Shao,Z.,&Hoffmann,N.(nd).Global Attention Mechanism:Retain Information to Enhance Channel-Spatial Interactions.

[0006] To enhance the cross-dimensional information interaction capabilities between features, Liu et al. proposed a Global Attention Mechanism (GAM) to capture important information about features in three dimensions. Coordinate attention (CA) embeds information from specific directions along the spatial dimension into channel attention and selects an appropriate dimensionality reduction rate for the channel dimension, achieving comparable performance. However, the cost of improving model accuracy is high computational cost and a large number of parameters, which is very detrimental to model deployment.

[0007] (3)Zhang,Q.-L.,&Yang,Y.-B.(2021).SA-Net:Shuffle Attention for DeepConvolutional Neural Networks.ICASSP 2021-2021IEEE International Conference on Acoustics,Speech and Signal Processing(ICASSP).Presented at the ICASSP2021-2021IEEE International Conference on Acoustics,Speech and SignalProcessing(ICASSP),Toronto,ON,Canada.https: / / doi.org / 10.1109 / icassp39728.2021.9414568.

[0008] Shuffle attention (SA) groups the channel dimension into multiple sub-features and addresses them in parallel, enabling efficient parallelization across multiple processors. However, in practical applications of small-target defect detection on metal surfaces, its actual detection performance is not very good. Summary of the Invention

[0009] To address the challenge of achieving a good balance between detection accuracy and inference speed when embedding attention mechanisms into existing metal surface defect detection methods, this invention provides a metal surface defect detection method based on an adaptive spatial attention mechanism. This method can better classify small metal surface defects without increasing computational resources, ensuring that training retains global information while focusing on important information from input features, thereby improving the model's recognition accuracy without reducing inference speed.

[0010] The technical solution of this invention is as follows: a method for detecting small target defects on metal surfaces based on an adaptive spatial attention mechanism, comprising the following steps:

[0011] S1: Construct a defect detection model;

[0012] S2: Based on the training dataset, train the defect detection model to obtain the trained defect detection model;

[0013] S3: Acquire an image of the workpiece to be inspected, input it into the trained defect detection model, and output the corresponding defect detection result using the trained model weights;

[0014] Its features are:

[0015] The defect detection model includes: an adaptive spatial attention mechanism;

[0016] The processing flow in the adaptive spatial attention mechanism includes:

[0017] a1: The first layer of the adaptive spatial attention mechanism is an adaptive feature pooling layer, which includes two adaptive average pooling layers and one adaptive max pooling layer;

[0018] The input graph of the adaptive spatial attention mechanism is denoted as: feature map F;

[0019] The feature map F is fed into an adaptive feature pooling layer;

[0020] a2: Upsample the context features at three scales output by the adaptive feature pooling layer to their original size, and then concatenate them along the channel direction to obtain the upsampled feature map;

[0021] a3: The upsampled feature map is fed into the hole space attention mechanism to obtain the dependencies in the feature space;

[0022] In the aforementioned hole spatial attention mechanism, the input feature map is fed into max pooling and mean pooling respectively, and then the results of the two pooling operations are concatenated into a feature map, which is then fed into the Conv Layer. The output of the Conv Layer is activated using the Sigmoid function, and the spatial weight coefficients are obtained at this time.

[0023] a4: Multiply the upsampled feature map with the spatial weight coefficient to obtain the product result; after standardizing the product result, a new feature map with spatial dependencies can be obtained;

[0024] a5: The new feature map is split into three feature tuples with the same shape as the feature map F along the channel direction. The three feature tuples and the input feature map of the adaptive spatial attention mechanism are added element by element along the channel direction. The resulting feature map is used as the output of the adaptive spatial attention mechanism.

[0025] Its further features are:

[0026] The defect detection model is built based on the YOLOv5 model;

[0027] The adaptive spatial attention mechanism is embedded between the backbone and neck of the YOLOv5 model, and finally outputs the result through the head of the output end.

[0028] After the input image is fed into the backbone network, the backbone network extracts features from the input image and outputs a feature map F. The feature map F is then fed into the adaptive spatial attention mechanism. After background suppression and key feature enhancement by the adaptive spatial attention mechanism, it is fed into the PANet neck for further feature extraction. Finally, it is fed into the output head for category prediction and bounding box localization, and the final prediction result is output.

[0029] The Conv Layer comprises: two 2D convolutions with a kernel size of 3 connected in sequence, and one 2D dilated convolution with a kernel size of 3.

[0030] This application provides a method for detecting small target defects on metal surfaces based on an adaptive spatial attention mechanism. This mechanism embeds the adaptive spatial attention mechanism into an existing network model. Within this mechanism, adaptive pooling is performed on the output features. While preserving the most prominent features in the image, it utilizes information beyond the main label to enhance the generalization ability of the surface defect recognition task, helping the model learn holistic features. The feature maps output by the adaptive pooling operation are then concatenated along channels to obtain the most representative feature map that also takes into account contextual information. A dilated spatial attention mechanism is then used to increase the receptive field while enhancing the spatial dependencies of defect features. Finally, dimensionality reduction of the feature map through channel splitting and merging better preserves the extracted features while reducing computational resources. The adaptive spatial attention mechanism proposed in this method considers the dependencies between input feature channels and space, achieving improved detection accuracy with fewer parameters and less computation. Therefore, the small target detection model for metal surfaces proposed in this application, which incorporates the adaptive spatial attention mechanism, can more accurately classify small defects on metal surfaces without increasing computational resources. This ensures that training retains global information while paying more attention to important information in the input features, thereby improving the model's recognition accuracy and achieving a good balance between model accuracy and inference speed. Attached Figure Description

[0031] Figure 1 Flowchart of a metal surface defect detection model;

[0032] Figure 2 Flowchart of the adaptive spatial attention mechanism;

[0033] Figure 3 Here is the flowchart for the Conv Layer;

[0034] Figure 4 A heatmap showing the detection results of a metal surface defect detection example. Detailed Implementation

[0035] This application includes a method for detecting small target defects on metal surfaces based on an adaptive spatial attention mechanism, which includes the following steps.

[0036] S1: Construct a defect detection model.

[0037] The adaptive spatial attention mechanism proposed in this method is a plug-and-play module that not only considers capturing key information and location information of target features, but also meets the lightweight requirements of mobile model deployment. It can be embedded into any existing deep learning model, such as the YOLO series, VGG, LeNet, and AlexNet, to improve the model's detection accuracy. The specific embedding location can be anywhere in the model, such as between the backbone and neck of the YOLOv5 model, or between the neck and the output head; similarly, the attention mechanism in this application can also be inserted into any layer of the backbone network in VGG and LeNet. The final usage location can be adjusted and determined based on the actual detection results.

[0038] The defect detection model in this embodiment is based on You Only Look Once v5 7.0 (YOLOv5). The adaptive spatial attention mechanism is embedded between the backbone network and the neck of the YOLOv5 model. The defect detection model specifically includes: the CSPDarknet53 backbone network, the adaptive spatial attention mechanism, the PANet neck and the output head, which are connected in sequence.

[0039] Its testing process is as follows: Figure 3 As shown, after the input image is fed into the Darknet53 backbone network, the Darknet53 backbone network extracts features from the input image and outputs a feature map F. The feature map F is then fed into the adaptive spatial attention mechanism. After background suppression and key feature enhancement by the adaptive spatial attention mechanism, it is fed into the neck (PANetNeck) for further feature extraction. Finally, it is fed into the output head for class prediction and bounding box localization, and the final prediction result is output.

[0040] like Figure 2 As shown, the processing flow in the adaptive spatial attention mechanism includes the following steps.

[0041] a1: The feature map F extracted by the backbone network is fed into the adaptive feature pooling layer as input;

[0042] The adaptive pooling layer includes two adaptive average pooling layers (AvgPool) and one adaptive max pooling layer (MaxPool).

[0043] The defects to be identified are extracted using the CSPDarknet53 backbone network to obtain feature maps F, with dimensions S = h × w. F is then processed by an adaptive pooling layer to obtain contextual feature maps at three scales (γ1 × S, γ2 × S, γ3 × S). The common pooling factor γ of the adaptive pooling layer is [0.1, 0.64], where the specific pooling factor is adaptively determined by the input dataset size.

[0044] Max pooling retains the most prominent features in the image while ignoring secondary features, extracting significant edge and texture details, thus accelerating the model's feature extraction from the input feature map. Average pooling, on the other hand, preserves good background information, utilizing information beyond the main label to improve the generalization ability of surface defect recognition tasks. This helps the model learn holistic features, enabling a more comprehensive understanding of the relationship between the defective target and its surrounding environment, and allowing for the extraction of the target's positional information in the input feature map, thereby improving detection accuracy. Ultimately, the adaptive feature pooling layer ensures that it can extract crucial and comprehensive features from the input image.

[0045] a2: The contextual features at three scales output by the adaptive feature pooling layer are upsampled to their original size, and then concatenated along the channel direction to tightly pack the extracted features while making full use of computational resources. The final concatenated features retain more layers of contextual information, resulting in the upsampled feature map.

[0046] In practice, the context features at three scales after the adaptive feature pooling layer are upsampled to the original size, and then the cat function in torch is used to perform the concatenation operation. The upsampling step uses bilinear interpolation to obtain more continuous and refined features.

[0047] a3: The upsampled feature map is fed into the hole space attention mechanism to obtain the dependencies in the feature space.

[0048] In the dilated spatial attention mechanism, the input feature map is subjected to max pooling (MaxPool(1×h×w)) and mean pooling (AvgPool(1×h×w)) to reduce dimensionality along the channel dimension. The results of the two pooling operations are then concatenated into a single feature map, which is then fed into the Conv Layer. The output of the Conv Layer is activated using the Sigmoid function, thus obtaining the spatial weight coefficients.

[0049] like Figure 3 As shown, the Conv Layer consists of two sequentially connected 2D convolutions with kernel size of 3 (3×3Conv2d) and one 2D dilated convolution with kernel size of 3 (3×3DilatedConv2d). Using 2D dilated convolution in the Conv Layer not only ensures that the receptive field size of the original convolution remains unchanged, but also reduces the number of model parameters in the original spatial attention mechanism, thus lowering computational complexity.

[0050] a4: The upsampled feature map is multiplied by the spatial weight coefficients (Product operation) to obtain the product result; the product result is then standardized (BatchNorm) to accelerate network convergence, thus obtaining a new feature map with spatial dependencies. In this application, a spatial attention mechanism is used to assign weights to each pixel in the feature map to reduce the interference of background noise.

[0051] a5: Due to the concatenation operation in a2, the number of channels in the feature map is now three times the number of input channels. To maintain consistency with the number of input channels in the neck region, the new feature map is split along the channel direction (SplitChannel) into three feature tuples with the same shape as the feature map F. The three feature tuples are then added element-wise along the channel direction to the input feature map of the adaptive spatial attention mechanism to ensure the invariance of the original input features. The resulting feature map is then used as the input to the network neck region.

[0052] Among them, the feature tuple has the same shape as the feature map F, that is, the width, height and number of channels of the feature tuple are the same as those of the feature map F.

[0053] In step a5, the new feature map is split and its dimensionality is reduced to retain all feature elements. Compared with convolutional dimensionality reduction using a kernel of 1, this dimensionality reduction method can retain more feature information with the least amount of computation.

[0054] The adaptive spatial attention mechanism proposed in this application improves the generalization ability of the model for surface defect recognition by combining adaptive max pooling and average pooling. It also enhances the model's understanding of the dependencies between spatial features of defects through an improved spatial attention mechanism. Furthermore, the dimensionality reduction method by splitting channels preserves as many original defect features as possible, which helps improve model training performance. This adaptive spatial attention mechanism is plug-and-play, considering not only capturing key information and location information of target features but also meeting the lightweight deployment requirements of mobile devices. It tightly connects the target feature channels through pooling operations and then uses the spatial attention mechanism to assign weights to each pixel in the feature map to reduce background noise interference.

[0055] S2: Based on the training dataset, train the defect detection model to obtain the trained defect detection model.

[0056] In practice, the training dataset can use datasets from existing technologies or be constructed based on historical data using existing graph processing techniques.

[0057] S3: Acquire the image of the workpiece to be inspected, input it into the trained defect detection model, and output the corresponding defect detection result using the trained model weights.

[0058] like Figure 4 The image shown is a heatmap illustrating the detection results of a metal surface defect detection example.

[0059] The image of the product to be detected is used as input to the trained defect detection model. The resulting defect detection output includes three defects, each represented by a heatmap visualization.

[0060] To verify the performance of the technical solution proposed in this application, experiments were conducted on the PyTorch platform. The AdamW optimizer was used for model performance optimization, with a batch size of 16 and an initial learning rate of 0.001. The LambdaLR learning rate strategy was used, and the learning rate was adjusted based on the warmup period. The number of training epochs was set to 300. To verify the difference between the attention mechanism proposed in this application and existing attention mechanisms, the dataset and network structure remained consistent, with the network adopting the YOLOv5s network structure. Based on the coaxial package metal base defect dataset (ILS-MB), the accuracy (P), recall (R), F1 score (F1), mean precision (mAP@IoU=0.5), model size (Model), computational cost (GFLOPs), and inference speed (Latency) of YOLOv5s models with different attention mechanisms for metal surface defect detection were compared. Inference speed is measured in latency, which refers to the delay time; a higher latency indicates poorer performance.

[0061] To verify the superiority of the Adaptive Spatial Attention (ASAM) mechanism over existing attention mechanisms, 12 different attention mechanisms were used to study metal surface defects on the ILS-MB dataset. These included feature recalibration attention mechanisms such as Squeeze-and-Excitation (SE), Efficient Channel Attention (ECA), Selective Kernel Attention (SKAttention), and Convolutional Block Attention Module (CBAM); coordinate attention mechanisms such as CoordAttention, Global Attention Map (GAM), and Exponential Moving Average (EMA); similarity distance metric attention mechanisms such as SimilarityAttention Module (SimAM) and TripletAttention; and some other popular attention mechanisms such as Spatial Group-wise Enhancement (SGE). 2 -MLPv2 (S2Attention) and Bi-Level RoutingAttention (BiFomer) attention, Figure 1 To validate the basic framework of the model.

[0062] During validation, each attention mechanism was embedded between the backbone and neck of the YOLOv5s model, respectively. Under the same experimental conditions, their performance on the ILS-MB test set was tested, and the experimental results are shown in Table 1.

[0063] Table 1. Performance comparison of different attention mechanisms on the ILS-MB test set.

[0064]

[0065] The F1 score is a comprehensive indicator that measures both precision and recall. As shown in Table 1 of the experimental results:

[0066] The adaptive attention mechanism (ASAM) proposed in this application has the best overall performance in terms of accuracy and recall, which alleviates the problems of low accuracy and high false negative rate of small target defects on metal surfaces to a certain extent.

[0067] In terms of reasoning speed, there is no significant difference in reasoning time between different types of attention.

[0068] In terms of lightweight model deployment, the adaptive attention mechanism proposed in this application, while maintaining the same model computation and size, achieves 1.0% higher accuracy, 4.0% higher recall, 2.7% higher F1 score, and 2.7% higher mean precision than the original network model. Furthermore, compared to other lightweight attention mechanisms, the adaptive attention mechanism proposed in this application achieves a mean precision that is 3.9% higher than that of SE, ECA, CBAM, EMA, SimAM, and TripletAttention, respectively.

[0069] Compared to computationally intensive attention mechanisms, the average accuracy of the adaptive attention mechanism proposed in this application is 0.4%, 1.4%, 1.3%, 4.5%, and 4.5% higher than that of SKAttention, CoordAttention, GAM, S2Attention, and BiFormer, respectively. Although the average accuracy of SKAttention is only 0.4% lower than that of the adaptive attention mechanism proposed in this application, its high computational cost and model size make it extremely difficult to deploy in lightweight models. Therefore, the overall performance of the adaptive spatial attention mechanism proposed in this application is significantly better than the above 12 attention mechanisms.

[0070] To further verify the effectiveness of the adaptive spatial attention mechanism proposed in this application, three other models—YOLOv3-tiny, YOLOv7-tiny, and YOLOv5—were selected and tested on the ILS-MB test set. The experimental results are shown in Table 2.

[0071] Table 2 compares the performance of different ASAM models on the ILS-MB test set.

[0072]

[0073] Based on the experimental results in Table 2, we can conclude that:

[0074] From the perspective of model deployment, after embedding the adaptive spatial attention mechanism, the computational cost of each model remains unchanged while the inference time remains basically unchanged.

[0075] In terms of model detection accuracy, the average precision of YOLOv3-tiny with the embedded adaptive spatial attention mechanism improved by 1.1%. Although the detection accuracy decreased after embedding the adaptive spatial attention mechanism, the detection recall increased by 3.1%, which is very beneficial for alleviating the problem of high false negative rate in metal surface defect detection.

[0076] After YOLOv7-tiny incorporates an adaptive attention mechanism, its average precision improved by 2.3%, recall improved by 7.5%, and overall performance was improved.

[0077] For YOLOv5n, YOLOv5m, and YOLOv5l, after embedding the adaptive spatial attention mechanism, although the F1 scores of YOLOv5n and YOLOv5l decreased except for YOLOv5m, their average accuracy improved by 0.3%, 1.2%, and 1.3%, respectively.

[0078] Overall, the adaptive spatial attention mechanism proposed in this application can achieve performance improvement on different models.

Claims

1. A method for detecting small target defects on metal surfaces based on an adaptive spatial attention mechanism, comprising the following steps: S1: Construct a defect detection model; The defect detection model is built based on the YOLOv5 model; The defect detection model specifically includes: the CSPDarknet53 backbone network, the adaptive spatial attention mechanism, the PANet neck, and the output head, which are connected in sequence. S2: Based on the training dataset, train the defect detection model to obtain the trained defect detection model; S3: Acquire an image of the workpiece to be inspected, input it into the trained defect detection model, and output the corresponding defect detection result using the trained model weights; Its features are: The defect detection model includes: an adaptive spatial attention mechanism; The processing flow in the adaptive spatial attention mechanism includes: a1: The first layer of the adaptive spatial attention mechanism is an adaptive feature pooling layer, which includes two adaptive average pooling layers and one adaptive max pooling layer; The input graph of the adaptive spatial attention mechanism is denoted as: feature map F; The feature map F is fed into an adaptive feature pooling layer; a2: Upsample the context features at three scales output by the adaptive feature pooling layer to their original size, and then concatenate them along the channel direction to obtain the upsampled feature map; a3: The upsampled feature map is fed into the hole space attention mechanism to obtain the dependencies in the feature space; In the aforementioned hole spatial attention mechanism, the input feature map is fed into max pooling and mean pooling respectively, and then the results of the two pooling operations are concatenated into a feature map, which is then fed into the Conv Layer. The output of the Conv Layer is activated using the Sigmoid function, and the spatial weight coefficients are obtained at this time. The Conv Layer includes: two two-dimensional convolutions with a kernel size of 3 connected in sequence and one two-dimensional dilated convolution with a kernel size of 3. a4: Multiply the upsampled feature map with the spatial weight coefficient to obtain the product result; after standardizing the product result, a new feature map with spatial dependencies can be obtained; a5: The new feature map is split into three feature tuples with the same shape as the feature map F along the channel direction. The three feature tuples and the input feature map of the adaptive spatial attention mechanism are added element by element along the channel direction. The resulting feature map is used as the output of the adaptive spatial attention mechanism.

2. The method for detecting small target defects on metal surfaces based on an adaptive spatial attention mechanism according to claim 1, characterized in that: The adaptive spatial attention mechanism is embedded between the backbone and neck of the YOLOv5 model, and finally outputs the result through the head of the output end. After the input image is fed into the backbone network, the backbone network extracts features from the input image and outputs a feature map F; the feature map F is then fed into the adaptive spatial attention mechanism. After background suppression and key feature enhancement through an adaptive spatial attention mechanism, the data is fed into the PANet neckline for further feature extraction, and finally fed into the output headline for category prediction and bounding box localization, outputting the final prediction result.

Images