Strip steel surface defect detection method based on attention mechanism and improved UNet network

Abstract

The invention discloses a strip steel surface defect detection method based on an attention mechanism and an improved UNet network, relates to the field of industrial image processing, and aims to solve the defects of a traditional detection method in defect positioning precision, multi-scale feature fusion and class imbalance. According to the technical scheme, an improved UNet network architecture is adopted, and the method specifically comprises the following steps: in a coding stage, designing a UELAN feature extraction module to improve parameter efficiency; a coordinate attention mechanism (CA) is introduced to enhance the focusing capability on a defect area; and capturing multi-scale high-level semantic features in combination with a cavity space pyramid pooling module (ASPP). In a decoding stage, a convolutional block attention mechanism (CBAM) is integrated, and a channel and spatial feature fusion are dynamically optimized. Besides, a combined loss function is put forward, focus loss and Dice loss are fused, the problem of boundary uncertainty is relieved in combination with a label fuzzification strategy, and category imbalance is coped with through weighted cross entropy. According to the invention, high-precision pixel-level segmentation of strip steel surface defects is realized.

Description

Technical Field

[0001] The present invention relates to the field of industrial image processing, and in particular to a strip steel surface defect detection method based on an attention mechanism and an improved UNet network. Background Art

[0002] As an important industrial raw material, strip steel is widely used in automobiles, home appliances, construction and other fields. In the production process of strip steel, surface defects (such as scratches, plaques, and internal debris) will seriously affect product quality. The detection of surface defects is a key link in quality control. Traditional defect detection methods rely on manual visual inspection or simple image processing technology. Although manual visual inspection can identify defects to a certain extent, it has problems such as low efficiency, high false detection rate, and difficulty in adapting to complex defects. In recent years, deep learning technology, especially semantic segmentation algorithms, has provided new solutions for strip surface defect detection. However, when dealing with strip surface defects, the existing semantic segmentation model still has problems such as large defect scale differences, blurred boundaries, and unbalanced categories, resulting in insufficient detection accuracy and high false detection and missed detection rates. In addition, simple image processing technologies (such as edge detection, threshold segmentation, etc.) often have difficulty in accurately identifying the boundaries and types of defects when dealing with complex defects, especially when the defect scales vary greatly and the background is complex, the detection effect is poor. Therefore, there is an urgent need for a method that can solve the shortcomings of traditional detection methods in terms of defect positioning accuracy, multi-scale feature fusion and category imbalance, and achieve high-precision pixel-level segmentation of strip surface defects, making it suitable for industrial quality inspection. In recent years, with the development of deep learning technology, especially the application of convolutional neural networks (CNN) and semantic segmentation algorithms, the accuracy and efficiency of strip surface defect detection have been significantly improved. However, the existing semantic segmentation model still faces the following challenges when dealing with strip surface defects: 1. Large differences in defect scales: The size of defects on the surface of strip steel varies greatly, ranging from a few millimeters to tens of centimeters. The existing single-scale feature extraction method is difficult to process defects of different scales at the same time; 2. Blurred boundaries: The boundaries of surface defects on the strip are often unclear, especially the pixel classification in the transition area is difficult to accurately define, resulting in inaccurate segmentation results; 3. Category imbalance: The number of pixels of surface defects in the strip is much less than that of normal areas, which causes the model to tend to predict normal areas and ignore defective areas during training.

[0003] In view of the above problems, the present invention proposes a strip surface defect detection method based on attention mechanism and improved UNet network. By optimizing the UNet network structure and combining multiple attention mechanisms and multi-scale feature extraction modules, the detection accuracy of strip surface defects is improved. Summary of the invention

[0004] The present invention provides a strip surface defect detection method based on an attention mechanism and an improved UNet network, which improves the detection accuracy of strip surface defects by optimizing the UNet network structure and combining multiple attention mechanisms and multi-scale feature extraction modules. The specific technical solution is as follows: 1. Data collection and preprocessing: The image dataset of strip surface defects is obtained. The image size is 200×200, including three types of defects: inclusions, plaques, and scratches. Data enhancement operations are performed on the images, including random rotation, flipping, scaling, etc., to increase the diversity of training data. The defect labels are fuzzified by Gaussian filtering to generate smooth labels to solve the problem of blurred defect boundaries. The formula is as follows:

[0005] in, is the standard deviation of the Gaussian kernel, which controls the smoothness of the filter.

[0006] 2. Model design: Build a semantic segmentation model based on UNet. The specific optimization strategies include: Optimization of the encoding stage: redesign the feature extraction module, introduce the UELAN module, improve the feature extraction capability; introduce the coordinate attention mechanism (CA) to enhance the model's attention to defect areas; introduce the atrous spatial pyramid pooling module (ASPP) at high levels in the encoding stage to extract multi-scale features. The calculation formula of the UELAN module is as follows:

[0007] in, is the input feature map, is the output feature map, and Respectively represent 3 3 and 1 1 Convolution operation, Concat represents feature connection operation.

[0008] Decoding stage optimization: In the decoding stage, the convolutional block attention mechanism (CBAM) is introduced to dynamically adjust channel and spatial features to improve the model's ability to capture defect details. The calculation formula for channel attention is as follows:

[0009] Among them, F is the input feature map, and They represent global average pooling and global maximum pooling respectively, MLP represents multi-layer perceptron, is the sigmoid function.

[0010] The calculation formula of spatial attention is as follows:

[0011] in, Indicates 7 7 convolution operations, is the sigmoid function.

[0012] Loss function optimization: A combined loss function is used, combining focal loss and Dice loss, to solve the problems of category imbalance and difficult and easy samples.

[0013] The formula of the focal loss function is as follows:

[0014] in, is the class weight, is the modulated sparsity coefficient, which is used to increase the attention to difficult samples. and denote the class probability and the predicted class probability respectively. and Represent the sample label value and sample prediction value respectively.

[0015] The Dice coefficient is a set similarity measurement function, and its calculation formula is:

[0016] The calculation formula for Dice loss is:

[0017] Among them, A and B are the pixel sets of the predicted results and the true labels, respectively.

[0018] Combining the two as the overall loss function of the network, the calculation formula is:

[0019] Among them, β represents the adjustment weight coefficient of the Dice loss function, and the value in this paper is 0.5.

[0020] 3. Model training: The Adam optimizer is used for model training, the initial learning rate is set to 0.001, and the learning rate decay strategy is adopted. During the training process, the combined loss function is used for back propagation to ensure that the model can converge stably.

[0021] 4. Defect detection and output: The trained model is applied to the task of strip surface defect detection, the pixel-level segmentation results of the defects are output, and the model performance is evaluated by indicators such as intersection over union (IoU). The calculation formula of IoU is as follows:

[0022] Preferred embodiment: (1) Data preprocessing: Data enhancement and label fuzzification are performed on the surface defect images of strip steel to improve the generalization ability of the model.

[0023] (2) Optimization of the encoding stage: In the encoding stage of the UNet network, the UELAN module is used to replace the original convolution module, and the coordinate attention mechanism (CA) and the atrous spatial pyramid pooling module (ASPP) are introduced to enhance the feature extraction capability and multi-scale processing capability.

[0024] (3) Decoding stage optimization: In the decoding stage of the UNet network, the convolutional block attention mechanism (CBAM) is introduced to dynamically adjust the channel and spatial attention in the feature fusion process.

[0025] (4) Loss function design: A combined loss function is used, combining focal loss and Dice loss, to solve the problems of category imbalance and difficult and easy samples.

[0026] (5) Model training and testing: The SD-900 strip steel surface defect dataset is used for model training and testing to evaluate the performance of the model.

[0027] Compared with the prior art, the present invention has the following technical advantages: (1) High-precision detection: By introducing the coordinate attention mechanism and the convolutional block attention mechanism, the model can more accurately locate and identify surface defects of strip steel, especially defects of small size and complex shape; (2) Multi-scale processing capability: Through the Atrous Spatial Pyramid Pooling (ASPP) module, the model can handle defects of different scales and enhance the detection capability of multi-scale objects; (3) Efficient parameter utilization: The design of the UELAN module improves parameter efficiency and reduces computational effort while maintaining high feature extraction capabilities; (4) Robust loss function: The design of the combined loss function effectively solves the problems of class imbalance and difficult and easy samples, and improves the model's attention to defect areas; (5) Strong generalization ability: Through data enhancement and label fuzzification processing, the model is not prone to overfitting during training and has good generalization ability. BRIEF DESCRIPTION OF THE DRAWINGS

[0028] Figure 1 The schematic diagram of the UNet network structure after the convolutional block attention mechanism CBAM is introduced into the present invention; Figure 2 This is a schematic diagram of the UNet network structure of the semantic segmentation model of the present invention; Figure 3(a) is a schematic diagram of the basic modules of the UNet encoding stage; Figure 3(b) is a schematic diagram of the optimized basic module UELAN in the encoding stage; Figure 4(a) is a schematic diagram of embedding coordinate information; Figure 4(b) is a schematic diagram of output feature map calculation; Figure 4(c) is a schematic diagram of coordinate attention generation; Figure 5(a) is a schematic diagram of a one-dimensional ordinary convolution; Figure 5(b) is a schematic diagram of a one-dimensional dilated convolution with a dilation factor of r=2; Figure 6 Schematic diagram of the structure of the atrous spatial pyramid pooling module (ASPP) of the present invention; Figure 7(a) is a schematic diagram of channel attention; Figure 7(b) is a schematic diagram of spatial attention; Figure 7(c) is a schematic diagram of convolutional block attention. DETAILED DESCRIPTION

[0029] The present invention will be further described in detail below in conjunction with the accompanying drawings, which are intended to explain the present invention rather than to limit it.

[0030] A strip steel surface defect detection method based on an attention mechanism and an improved UNet network comprises the following steps: 1. Data Acquisition and Preprocessing (Step S1) A dataset of surface defect images of strip steel is obtained. The image size is 200×200, and it contains three types of defects: inclusions, plaques, and scratches. Data enhancement operations are performed on the images, including random rotation, flipping, scaling, etc., to increase the diversity of training data. Gaussian filtering is used to fuzzify the defect labels to generate smooth labels to solve the problem of blurred defect boundaries.

[0031] 2. Model design (step S2) Build a semantic segmentation model based on UNet. The specific optimization strategies include: S2.1 Optimization of the encoding stage: The improved UNet network structure of the present invention is as follows Figure 1 As shown, the network structure is derived from the classic UNet network structure as shown Figure 2As shown in Figure 3 (a), the two-dimensional 3*3 convolution kernel structure used for the UNet network encoder is shown in Figure 3 (a). The structure is composed of two 3*3 convolution kernels stacked together. In the figure, C represents the number of channels. The number of output channels of the first convolution kernel is doubled, and the number of input and output channels of the second convolution kernel remains the same. The feature extraction module is redesigned. First, the UELAN module is introduced. The feature extraction core module ELAN of this module has higher parameter efficiency and learning ability. Its basic structure is shown in Figure 3 (b). The upper branch contains three 3*3 convolutions and the lower branch contains one 1*1 convolution. The number of input and output channels of each convolution is C. The outputs of the four convolutions are spliced, and then 1*1 convolution is used for fusion and channel dimensionality reduction, which improves the feature extraction ability; secondly, the coordinate attention mechanism (CA) is introduced, which can encode both spatial and channel dimensions at the same time to better locate and identify key areas. The module first embeds the coordinate information into the network, as shown in Figure 4 (a), which can generate a pair of direction-aware feature maps. The feature captures long-distance dependencies along one spatial direction and retains accurate location information along another direction. Coordinate attention generation refers to the process of processing the aggregated features to obtain coordinate attention weights. As shown in Figure 4 (b), it concatenates the aggregated features in the horizontal and vertical directions, and then uses two layers of 1*1 convolution to obtain channel information. The number of channels in the middle layer is, which can control the amount of convolution operation. Finally, the features are split into horizontal and vertical directions, that is, attention weights are obtained, which enhances the model's attention to the defect area. The calculation of the output feature map refers to the process of re-weighting the feature value using attention weights, as shown in Figure 4 (c). The weight coefficients come from the horizontal and vertical directions respectively. The value of each feature point is equal to the product of its own value and the attention weight corresponding to the spatial position; finally, the high-level atrous spatial pyramid pooling module (ASPP) is introduced in the encoding stage. It uses multiple atrous convolutions with different dilation factors to perform parallel operations on high-level feature maps to obtain multi-scale features, such as Figure 6 The ASPP module realizes multi-scale feature information extraction by replacing ordinary convolution with dilated convolution. The one-dimensional ordinary convolution is shown in Figure 5 (a), and the dilated convolution is shown in Figure 5 (b).

[0032] S2.2 Optimization of decoding stage: The convolutional block attention mechanism (CBAM) is introduced in the decoding stage to dynamically adjust the channel and spatial features, thereby improving the model's ability to capture defect details. Among them, the channel attention module uses global average and global maximum pooling operations to extract global information, and uses a shared fully connected layer to generate channel attention weights. Its structure is shown in Figure 7 (a). The spatial attention module uses channel aggregation to extract the global information and significant spatial features of the channel, and uses feature concatenation and convolution operations to learn the attention weights in the spatial dimension to generate a spatial attention map. Its structure is shown in Figure 7 (b). The calculation process of the CBAM attention mechanism is shown in Figure 7 (c). It first uses the channel attention mechanism to generate channel weights, multiplies the input feature map and the channel weights to obtain intermediate features with channel attention, and then uses the intermediate features to generate spatial attention weights. The intermediate features and the spatial attention weights are multiplied to obtain output features with spatial attention.

[0033] S2.3 Loss function optimization: To address the problem of category imbalance, the focus loss function is used to increase attention to difficult samples. The Dice loss function is designed from the perspective of set similarity to improve the overall segmentation effect of the model on defect areas. The focus loss and Dice loss are combined to form the final combined loss function for model training. The formula of the focus loss function is as follows:

[0034] The Dice coefficient is a set similarity measurement function, and its calculation formula is:

[0035] The calculation formula for Dice loss is:

[0036] Combining the two as the overall loss function of the network, the calculation formula is:

[0037] 3. Model training (step S3) The SD-900 strip surface defect dataset was used for model training, with an input image resolution of 224×224, and a combined loss function was used for back propagation optimization. During the training process, the Adam optimizer was used, the initial learning rate was set to 0.001, and the learning rate decay strategy was adopted to ensure stable convergence of the model.

[0038] 4. Defect detection and output (step S4) S4.1 Training process: Apply the trained model to the task of strip surface defect detection, output the pixel-level segmentation results of the defects, and evaluate the model performance through indicators such as intersection over union (IoU). The calculation formula is as follows:

[0039] S4.2 Testing process: Evaluate the performance of the model on the test set, using mIoU (mean intersection over union) as the main evaluation indicator to verify the segmentation accuracy of the model. The calculation formula is as follows:

[0040] 5. Defect detection and output (step S5) S5.1 Performance evaluation: Through comparative experiments, the superiority of the optimized UNet model in the task of strip surface defect detection is verified. The experimental results show that the optimized model is significantly better than other common semantic segmentation models in terms of mIoU indicators.

[0041] S5.2 Visualization: The detection effect of the model before and after optimization is displayed through the segmentation effect diagram to verify the model's ability to accurately segment the defect edge.

[0042] 6. Overview of the complete process (step S6) The operation process of the entire strip surface defect detection method is as follows: S6.1 Startup phase: Initialize the equipment and model, and start data collection.

[0043] S6.2 Data acquisition and preprocessing: Real-time acquisition of strip surface images and data enhancement and label fuzzification processing.

[0044] S6.3 Model prediction: The preprocessed image passes through the encoding stage, decoding stage and classification module in sequence, and finally the segmentation result of the surface defects of the strip is obtained.

[0045] S6.4 Real-time monitoring and adjustment: During the strip production process, surface defects are detected in real time, the quality of the strip is monitored, and the accuracy of the detection is ensured.

[0046] S6.5 Shutdown and Evaluation: After production is completed, evaluate model performance and make optimization adjustments.

[0047] The present invention relates to a strip surface defect detection method based on an attention mechanism and an improved UNet network, in particular to a semantic segmentation technology for strip surface defects. The present invention combines the UNet network structure, the coordinate attention mechanism (CA), the atrous spatial pyramid pooling module (ASPP) and the convolutional block attention mechanism (CBAM), and proposes an efficient strip surface defect detection method. Through the optimized UNet network, the present invention can more efficiently extract the multi-scale features of strip surface defects, and effectively capture the local and global information of defects through the attention mechanism, thereby achieving high-precision segmentation of different types of defects. It aims to accurately identify and locate different types of defects through deep learning analysis of strip surface images, provide support for automated detection of strip production quality control, and is particularly suitable for the field of strip surface defect detection based on visual data.

[0048] In addition, the present invention adopts a combined loss function in the model training process, combining focal loss and Dice loss to solve the problems of category imbalance and difficult and easy samples, and optimizes the model parameters through the Adam optimizer, so that the model can stably and accurately predict the surface defects of strip steel. By combining the advantages of the UNet network and multiple attention mechanisms, the present invention effectively copes with the challenges of large defect scale differences, blurred boundaries and category imbalance in strip steel surface defect detection, and significantly improves the accuracy and stability of defect detection.

[0049] Through this multi-module combined deep learning model, the present invention has high detection accuracy and robustness when processing large-scale data and coping with complex industrial production environments. It can realize fast and accurate strip surface defect detection and effectively ensure the quality stability of strip production.

[0050] The above contents are only for explaining the technical idea of ​​the present invention and cannot be used to limit the protection scope of the present invention. Any changes made on the basis of the technical solution in accordance with the technical idea proposed by the present invention shall fall within the protection scope of the claims of the present invention.

Claims

1. A strip surface defect detection method based on attention mechanism and improved UNet network, characterized in that: The following steps are involved: Data collection and preprocessing: The public dataset of SD-900 strip surface defect images is obtained, which contains 900 images of three different types of defects, including scratches, plaques, and inclusions. Each type of defect has 300 images, and data enhancement and label fuzzification are performed. Optimization of the encoding and decoding stages: Build a semantic segmentation model based on UNet, optimize the structure of the encoding and decoding stages, replace the convolution module of the UNet encoding stage with the UELAN module, combine the coordinate attention mechanism (CA) and the atrous spatial pyramid pooling module (ASPP) to enhance the feature extraction and multi-scale processing capabilities, and introduce the convolutional block attention mechanism (CBAM) in the UNet decoding stage to dynamically adjust the channel and spatial attention in the feature fusion process; Loss function design: A combined loss function is used, combining focal loss and Dice loss to solve the problems of category imbalance and difficult and easy samples, and the model parameters are optimized through the Adam optimizer; Model training and testing: Use the steel strip surface defect dataset for model training and testing to evaluate the model performance.

2. The strip surface defect detection method based on attention mechanism and improved UNet network according to claim 1 is characterized in that: The data preprocessing step includes: The strip surface image is randomly rotated, scaled, and flipped to achieve data enhancement; the defect labels are Gaussian filtered to generate fuzzy labels to smooth the pixel classification in the transition area.

3. The strip surface defect detection method based on attention mechanism and improved UNet network according to claim 1 is characterized in that: The encoding stage optimization includes: The convolution module in the UNet encoding phase is replaced by the UELAN module, which improves the feature extraction efficiency through a multi-branch convolution structure. The coordinate attention mechanism (CA) is introduced after downsampling to enhance the model's ability to focus on defect areas; The Atrous Spatial Pyramid Pooling (ASPP) module is introduced at the high level of the encoding stage to obtain multi-scale feature information through atrous convolutions with different dilation factors.

4. The strip surface defect detection method based on attention mechanism and improved UNet network according to claim 1 is characterized in that: The decoding stage optimization step includes: The convolutional block attention mechanism (CBAM) is introduced after upsampling and feature concatenation in the decoding stage to dynamically adjust the channel and spatial attention in the feature fusion process; The CBAM module is used to enhance the model's ability to identify defective areas and improve segmentation accuracy.

5. The strip surface defect detection method based on attention mechanism and improved UNet network according to claim 1 is characterized in that: The combined loss function includes: Use the focal loss function to increase attention to difficult samples and solve the problem of class imbalance; The Dice loss function is used to improve the overall segmentation effect of the model on the defect area from the perspective of set similarity; Combine focal loss and Dice loss to form a combined loss function for model training; The formula of the focal loss function is as follows: ； The Dice coefficient is a set similarity measurement function, and its calculation formula is: ； The calculation formula for Dice loss is: ； Combining the two as the overall loss function of the network, the calculation formula is: 。 6. The strip surface defect detection method based on attention mechanism and improved UNet network according to claim 1 is characterized in that: The model training and testing steps include: The model was trained using a dataset of strip surface defects, with an input image resolution of 224×224, and a combined loss function was used for back propagation optimization; The performance of the model is evaluated on the test set, and mIoU (mean intersection over union) is used as the main evaluation indicator to verify the segmentation accuracy of the model. The formula is as follows: 。 7. The strip surface defect detection method based on attention mechanism and improved UNet network according to claim 1 is characterized in that: The model uses the Adam optimizer during training, with the learning rate set to 0.001 and the batch size to 32. The learning rate is dynamically adjusted using the learning rate decay strategy to ensure stable convergence of the model.

8. The strip surface defect detection method based on attention mechanism and improved UNet network according to claim 1 is characterized in that: The evaluation indicators of the model during the training process include average loss (AvgLoss) and mean intersection over union (mIoU), which are used to evaluate the prediction effect of the model.

9. The strip surface defect detection method based on attention mechanism and improved UNet network according to claim 1 is characterized in that: After the model training is completed, the surface image of the strip is predicted, and the prediction result is compared and evaluated with the actual defect label to verify the detection accuracy of the model.

10. The strip surface defect detection method based on attention mechanism and improved UNet network according to claim 1 is characterized in that: During real-time detection, the model can issue timely alarms and adjust production parameters based on the type, location and size of defects to ensure that the quality of the strip meets the standards.

Images