LAH-YOLO-based steel rail surface defect detection method

Abstract

The invention discloses a steel rail surface defect detection method based on LAH-YOLO. The detection precision of the YOLOv5-s algorithm on a steel rail surface defect data set still needs to be improved. The method comprises the following steps: acquiring a steel rail surface image, and inputting the steel rail surface image into a constructed LAH-YOLO defect detection network; the LAH-YOLO defect detection network comprises a backbone network, a feature fusion network and a detection head; the backbone network comprises a GRF-C3 module, and the feature fusion network is an AFPN feature fusion network. The method is suitable for small target and small sample data sets, the detection precision is remarkably improved while the parameter quantity is reduced, and the method has better detection performance on the steel rail surface defect data set.

Description

Technical Field

[0001] The invention belongs to the technical field of rail defect detection, and in particular relates to a rail surface defect detection method based on LAH-YOLO. Background Art

[0002] Rails are the most important components of railway tracks. They provide a continuous, smooth and least-resistance rolling surface for wheels, guide trains, and directly bear the huge pressure of wheels. Surface defects such as rail peeling and scratches directly affect the comfort and safety of train operation. For the massive amount of on-site rail photos taken by railway inspection vehicles, the surface defect information can be efficiently detected, and data support can be provided for maintenance by engineering departments in a timely manner, ensuring the smooth operation of trains and extending the service life of tracks, which has significant social and economic benefits. At present, rail surface defect detection based on image recognition is a hot research direction for scholars.

[0003] The rail surface peels off block defects, which usually occupy very few pixels and are classified as small targets. The anchor-base detection mechanism allows multiple candidate boxes of different sizes to be generated at one location, so as to better adapt to the shape and position changes of small targets. The YOLOv5 algorithm adopts the anchor-base detection mechanism and adopts a variety of data enhancement strategies to achieve good detection performance on small targets and small sample data sets.

[0004] The backbone network of the YOLOv5 algorithm uses the CSPDarknet53 fully convolutional neural network; the feature fusion network uses the PAN network, which contains two top-down feature pyramid networks to effectively fuse feature maps of different scales. The CSPDarknet53 network is composed of a combination of multiple 3×3 CBR convolutional layers and C3 modules. Each CBR convolutional layer is composed of a standard convolution followed by a BN layer and a rectified linear unit (ReLU) activation function. The design of the CSPDarknet53 network reduces the amount of calculation while enhancing the reusability of features and ensuring good feature extraction capabilities. The fast spatial pyramid pooling (SPPF) module is used at the end of the backbone network. This module can extract feature information of different scales without changing the size of the feature map, further enhancing the performance of feature extraction. The backbone network outputs three feature maps of different scales to the PAN network for feature fusion. In the YOLOv5 algorithm, the convolutional layer in the PAN network is replaced by the C3 module to further enhance the fusion capability of multi-scale features. The fused features are passed through a 1×1 standard convolution of the detection head to generate the final prediction result. To supervise the regression of the prediction box, CIoU loss is used, and Binary Cross Entropy Loss (BCE Loss) is used to supervise the category prediction.

[0005] The YOLOv5 algorithm is divided into multiple versions, n, s, m, l, and x, according to the design depth and width of its backbone network. The algorithm version can be switched by adjusting the parameters depth_multiple and width_multiple. Although the detection accuracy of the YOLOv5 algorithm can be improved by adjusting the depth and width parameters, it is also accompanied by an increase in the number of algorithm parameters. The YOLOv5-s algorithm has a small number of parameters and is suitable for actual deployment applications, but its detection accuracy on the rail surface defect dataset still needs to be improved. Summary of the invention

[0006] In order to make up for the shortcomings of the prior art, the present invention provides a rail surface defect detection method based on LAH-YOLO, which is suitable for small targets and small sample data sets. While reducing the number of parameters, the detection accuracy is significantly improved, and the detection performance on the rail surface defect data set is better.

[0007] In order to achieve the above object, the technical solution adopted by the present invention is:

[0008] A rail surface defect detection method based on LAH-YOLO, characterized by:

[0009] Acquire a rail surface image, and input the rail surface image into a constructed LAH-YOLO defect detection network; the LAH-YOLO defect detection network includes a backbone network, a feature fusion network, and a detection head;

[0010] The backbone network includes a GRF-C3 module, and the feature fusion network is an AFPN feature fusion network.

[0011] Furthermore, the GRF-C3 module first divides the input feature map into two parts according to half the number of channels, one part passes through the 1×1 CBR convolution layer and GRF-BottleNeck, and the other part directly passes through the 1×1 CBR convolution layer. The two parts are then concatenated according to the channel dimension and finally output after passing through the 1×1 CBR convolution layer.

[0012] Furthermore, the GRF-BottleNeck uses 3×3 and 5×5 convolution kernels, uses a CBR convolution layer before a depth-wise separable convolution, and applies an ECA attention mechanism to the residual connection part in GRF-BottleNeck. The features at the residual connection need to be processed by the ECA attention mechanism to assign different weights to each channel.

[0013] Furthermore, the AFPN fusion network first fuses the two adjacent layers of feature maps in the bottom layer, and then cross-fuse them after upsampling and downsampling; after feature fusion, 4 residual units are used to continue to extract features, and each residual unit includes 2 3×3 standard convolutions; after the fusion of the two layers of features in the bottom layer, the fused features are cross-fused with the top layer features for three layers of features, and after processing by the C3 module, the features are finally input into the detection head.

[0014] Furthermore, the fusion principle of the three-layer AFPN feature fusion network is as follows:

[0015] y 3 =α 3 ·x 1→3 +β 3 ·x 2→3 +γ 3 x 3→3

[0016] Among them, y 3 Represents the third layer features after fusion output; α 3 , β 3 , γ 3 They represent the weight factors of the input first-layer features, second-layer features, and third-layer features on the output third-layer features. Numerically, the sum of the weight factors is 1, i.e., α 3 +β 3 +γ 3 =1;x 1→3 、x2→3 、x 3→3 They represent the first-layer features, second-layer features, and third-layer features of the input respectively.

[0017] Furthermore, the detection head adopts an anchor-base detection mechanism, in which three detection heads perform detection on feature maps of three different scales respectively, and output five physical quantities: predicted category, x and y coordinates of the center point of the predicted box, and width and height of the predicted box.

[0018] Furthermore, the loss function of the object detection task includes three parts: classification loss function, bounding box loss function and confidence loss function;

[0019] The bounding box loss function adopts the Wise-IoU loss function, and the expression of the Wise-IoU loss function is as follows:

[0020] L WIOU = r·R WIoU ·L IoU

[0021] Among them, L WIoU represents the Wise-IoU loss; r represents the gradient gain coefficient; R WIoU represents the distance attention coefficient, whose value range is between [1, e), which can significantly amplify the IoU loss of ordinary anchor boxes; L IoU represents the IoU loss.

[0022] Furthermore, the gradient gain coefficient r is expressed as follows:

[0023]

[0024] Wherein β represents the clustering degree to describe the quality of the anchor frame, and its value range is between [0,). It is a dynamically adjusted value to adjust the gradient gain coefficient r to the best. α and δ represent two artificially set hyperparameters for controlling the clustering degree β. The present invention sets α and δ to 1.9 and 3 respectively.

[0025] Furthermore, the distance attention coefficient R WIoU The expression is as follows:

[0026]

[0027] Where x and y represent the horizontal and vertical coordinates of the center point of the prediction box, respectively. gt ,y gt Respectively represent the horizontal and vertical coordinates of the center point of the truth box, W g , H g Respectively represent the width and height of the smallest bounding box of the prediction box and the true value box. The superscript * indicates that W g , H gSeparating operations from the computation graph prevents R WIoU Produces gradients that hinder convergence.

[0028] Furthermore, the IoU loss L IoU The expression is as follows:

[0029]

[0030] Among them, S represents the area of ​​the prediction box, S gt Represents the area of ​​the true value box. The IoU loss is obtained by calculating the ratio of the intersection area of ​​the predicted box and the true value box to the union area.

[0031] Beneficial effects of the present invention:

[0032] 1) In the algorithm of the present invention, a BottleNeck with a larger receptive field (Greater Receptive Field BottleNeck, GRF-BottleNeck) combining the attention mechanism with BottleNeck is designed to enhance the feature extraction capability of the backbone network; at the same time, an Adaptive Feature Pyramid Network (AFPN) is used to replace the original PAN network, and some convolutions in the AFPN network are replaced by C3 modules, which further improves the fusion effect of multi-scale features and achieves good detection effect on the rail surface defect dataset, which is lightweight and high-precision;

[0033] 2) The GRF-BottleNeck of the present invention increases the receptive field by increasing the convolution kernel, and introduces the ECA attention mechanism in the residual connection part, so that the backbone network can focus on important local feature information, reduce feature information loss, and improve detection accuracy;

[0034] 3) The AFPN fusion network of the present invention adopts a progressive fusion method, realizes the fusion interaction of non-adjacent feature layers, and can efficiently fuse multi-scale features. BRIEF DESCRIPTION OF THE DRAWINGS

[0035] Figure 1 This is a diagram of the algorithm architecture of the present invention;

[0036] Figure 2 This is the GRF-C3 module structure diagram;

[0037] Figure 3 This is the GRF-BottleNeck structure diagram;

[0038] Figure 4 This is the AFPN feature fusion network structure diagram;

[0039] Figure 5The detection results of different detection algorithms for block drop caused by nuclear damage;

[0040] Figure 6 Detection results of different detection algorithms for peeled blocks caused by contact fatigue. DETAILED DESCRIPTION

[0041] The present invention is described in detail below in conjunction with specific implementation modes.

[0042] The rail surface defect detection method based on LAH-YOLO of the present invention is specifically as follows:

[0043] Obtain the rail surface image and input the rail surface image into the constructed LAH-YOLO defect detection network; Figure 1 As shown, the LAH-YOLO defect detection network includes a backbone network, a feature fusion network and a detection head; the backbone network includes a GRF-C3 module, and the feature fusion network is an AFPN feature fusion network.

[0044] The present invention designs a GRF-C3 module with stronger feature extraction capability in the backbone network, while reducing the number of parameters of the backbone network. In the feature fusion network part, a progressive adaptive fusion network AFPN is adopted. First, the feature maps of size 128×128 and 256×256 are adaptively fused; then the output feature maps of 128×128 and 256×256 are fused with the feature maps of size 512×512 output by the backbone network. This progressive fusion network helps to improve the feature fusion effect. The feature fusion method uses the adaptive spatial feature fusion (ASFF) method to distinguish the importance of different features. After each round of fusion, the C3 module is used to further increase the depth and receptive field of the network to improve the representation ability of the feature. In terms of the detection head, the LAH-YOLO algorithm detection mechanism and the coupled detection head, one branch predicts the regression box and category at the same time. The LAH-YOLO algorithm designed the GRF-C3 module and introduced the AFPN feature fusion network, which improved the effects of feature extraction and feature fusion while greatly reducing the number of parameters.

[0045] 1) GRF-C3 module

[0046] like Figure 2 As shown in the figure, the GRF-C3 module first divides the input feature map into two parts according to half the number of channels. One part passes through the 1×1 CBR convolution layer and GRF-BottleNeck, and the other part directly passes through the 1×1 CBR convolution layer. The two parts are then concatenated according to the channel dimension and finally output after passing through the 1×1 CBR convolution layer.

[0047] like Figure 3 As shown in the figure, GRF-BottleNeck uses 3×3 and 5×5 convolution kernels, uses CBR convolution layer before depth-wise separable convolution, and applies ECA attention mechanism to the residual connection part in GRF-BottleNeck. The features at the residual connection need to be processed by ECA attention mechanism to give different weights to each channel.

[0048] GRF-BottleNeck uses 3×3 and 5×5 convolution kernels. This combination of convolution kernels can maximize the acquisition of feature information without a large increase in the number of parameters. GRF-BottleNeck uses a CBR convolution layer before the depthwise separable convolution to reduce the loss of feature information and thus improve detection accuracy.

[0049] GRF-BottleNeck improves the backbone network's feature extraction capabilities by increasing the convolution kernel, replacing the Ghost convolution module with a CBR convolution layer, and introducing an attention mechanism. GRF-BottleNeck uses the attention mechanism to enable the backbone network to focus on extracting key features and improve the feature extraction effect of small targets. The attention mechanism enables the backbone network to shift its focus from the whole to important areas by allocating weights.

[0050] Based on SENet, the Efficient Channel Attention (ECA) mechanism introduces a channel attention mechanism in the convolution operation, and adopts a method of adaptively selecting the convolution kernel size. When the ECA attention mechanism is applied to the case where the dimension of the feature map is reduced, it has better feature focusing performance than the first two attention mechanisms. The present invention applies the ECA attention mechanism to the residual connection part in GRF-BottleNeck. The features at the residual connection need to be processed by the ECA attention mechanism, and different weights are assigned to each channel. This operation enhances the feature extraction capability without affecting the complexity of the algorithm, and plays a positive role in improving the accuracy.

[0051] 2) AFPN feature fusion network

[0052] The backbone network outputs three layers of feature maps of different scales. The feature maps at the bottom layer contain the richest spatial information, while the feature maps at the top layer have the richest semantic information. There are significant differences in feature information between non-adjacent feature maps. Figure 4 As shown in the figure, the AFPN fusion network is a progressive FPN fusion method, which first fuses the feature maps of 256 and 512 sizes, and then fuses them with the feature map of 1024 size;

[0053] First, the two adjacent layers of feature maps in the bottom layer are fused, and then cross-fused after upsampling and downsampling; after feature fusion, 4 residual units are used to continue to extract features, each of which includes 2 3×3 standard convolutions, similar to the ResNet structure; the LAH-YOLO algorithm uses the C3 module to replace the residual unit in the AFPN fusion network to process the fused features. Compared with using 2 3×3 standard convolutions, the C3 module performs a dimensionality reduction operation on the features before performing a convolution operation, which reduces the number of parameters without losing accuracy. After the fusion of the two layers of features in the bottom layer, the fused features are cross-fused with the top layer features for three layers of features. The implementation method and idea are consistent with the fusion of two layers of features. After being processed by the C3 module, the features are finally input into the detection head.

[0054] Taking the three-layer AFPN feature fusion network as an example, the fusion principle is as follows:

[0055] y 3 =α 3 ·x 1→3 +β 3 ·x 2→3 +γ 3 x 3→3

[0056] Among them, y 3 Represents the third layer features after fusion output; α 3 , β 3 , γ 3 They represent the weight factors of the input first-layer features, second-layer features, and third-layer features on the output third-layer features. Numerically, the sum of the weight factors is 1, i.e., α 3 +β 3 +γ 3 =1;x 1→3 、x 2→3 、x 3→3 They represent the first-layer features, second-layer features, and third-layer features of the input respectively.

[0057] The weight factor is calculated based on the scale of the input feature through convolution and softmax functions, and the weight factor can be dynamically adjusted based on the back-propagation operation to achieve the best feature fusion effect. Due to the differences in the number of fused feature layers, multiple fusion modules need to be designed according to actual needs.

[0058] 3) Detection head

[0059] The detection head of the present invention adopts an anchor-base detection mechanism. Three detection heads perform detection on feature maps of three different scales respectively, and output five physical quantities: predicted category, x and y coordinates of the center point of the prediction box, and the width and height of the prediction box. Each detection head uses three anchor boxes with the same center point but different width and height to preliminarily select the target. At the same time, the intersection-and-union ratio threshold is set to 0.5 to determine whether the grid corresponds to the foreground or background to be detected. Specifically, when the intersection-and-union ratio of the predicted box to the true value box is greater than or equal to 0.5, it is judged as the foreground; when the intersection-and-union ratio is less than 0.5, it is judged as the background. If it is judged to be the background, it is directly ignored. If it is judged to be the foreground, the algorithm regresses and locates the anchor box to generate a prediction box, and predicts and classifies the target in the box.

[0060] 4) Loss Function

[0061] The loss function of the target detection task includes three parts: classification loss function, bounding box loss function and confidence loss function; the expression is as follows:

[0062] L=λ1·L cls +λ2·L obj +λ3·L loc

[0063] Among them, L represents the overall loss, L cls represents the classification loss, L obj represents the confidence loss, L loc represents the bounding box loss, λ1, λ2, and λ3 represent the weights corresponding to the loss functions of each part respectively; the rail surface defect dataset has only one category, and the classification loss has little effect on the total loss; referring to the parameter setting experience of the YOLOv5 algorithm, the present invention sets λ1, λ2, and λ3 to 0.2, 0.4, and 0.4 respectively;

[0064] The bounding box loss function adopts the Wise-IoU loss function, and the expression of the Wise-IoU loss function is as follows:

[0065] L WIoU = r·R WIoU ·L IoU

[0066] Among them, L WIoU represents the Wise-IoU loss; r represents the gradient gain coefficient; R WIoU represents the distance attention coefficient, whose value range is between [1, e), which can significantly amplify the IoU loss of ordinary anchor boxes; L IoU represents the IoU loss.

[0067] The defects in the rail surface defect dataset are small and have irregular contours, and low-quality true value boxes are very likely to appear during the annotation process. If the regression positioning of the prediction box is strengthened according to these low-quality true value boxes, it will have an adverse effect on the positioning accuracy of the detection algorithm. Wise-IoU loss adopts a dynamic focusing mechanism, uses "outlier degree" instead of IoU to evaluate the quality of the anchor box, and provides a clear gradient gain allocation strategy. This strategy reduces the competitiveness of high-quality anchor boxes while reducing the harmful gradients generated by low-quality instances. Wise-IoU can focus on anchor boxes of ordinary quality.

[0068] The gradient gain coefficient r is expressed as follows:

[0069]

[0070] Wherein β represents the clustering degree to describe the quality of the anchor frame, and its value range is between [0,). It is a dynamically adjusted value to adjust the gradient gain coefficient r to the best. α and δ represent two artificially set hyperparameters for controlling the clustering degree β. The present invention sets α and δ to 1.9 and 3 respectively.

[0071] Distance attention coefficient R WIoU The expression is as follows:

[0072]

[0073] Where x and y represent the horizontal and vertical coordinates of the center point of the prediction box, respectively. gt ,y gt Respectively represent the horizontal and vertical coordinates of the center point of the truth box, W g , H g Respectively represent the width and height of the smallest bounding box of the prediction box and the true value box. The superscript * indicates that W g , H g Separating operations from the computation graph prevents R WIoU Produces gradients that hinder convergence.

[0074] IoU loss L IoU The expression is as follows:

[0075]

[0076] Among them, S represents the area of ​​the prediction box, S gt Represents the area of ​​the true value box. The IoU loss is obtained by calculating the ratio of the intersection area of ​​the predicted box and the true value box to the union area.

[0077] 5) Experimental verification

[0078] The following comparative experiments with other algorithms in the YOLO series are used to verify the detection performance of the LAH-YOLO algorithm in the field of rail surface defect detection, and ablation experiments are performed to verify the performance improvement effect of each module.

[0079] a) Experimental setup

[0080] The experimental environment was built on the Ubuntu 18.04.5 operating system, and was programmed in Python 3.9. Pytorch 1.8.0 and Cuda 11.2 were used to build the deep learning framework. As for the hardware equipment, an AMD 5900X with 32GB memory and a single NVIDIA GeForce RTX 3090 graphics card were used. In this experiment, the input image size was automatically adjusted to 640×640, and the algorithm model was trained using the Stochastic Gradient Descent (SGD) optimizer. The initial learning rate was set to 0.01, the weight decay was set to 0.001, the final learning rate was set to 0.001, the batch-size was 16, the confidence threshold was set to 0.25, and the IoU threshold was set to 0.45. During training, due to the poor convergence performance of the YOLOv6 algorithm, its total number of iterations was set to 400, and the total number of iterations of the other algorithms was 300. During the training process, data enhancement methods such as Mixcup and Mosaic are used, and the optimal training weights are taken to calculate the final detection accuracy.

[0081] b) Evaluation indicators

[0082] The commonly used accuracy evaluation indicator in the field of target detection is mAP 50 and mAP 50-95 , the subscript number represents the Intersection over Union (IoU) threshold. In the object detection task, the IoU threshold is usually set to filter out some inaccurate prediction boxes. mAP 50 It represents the mean average precision (mAP) when the IoU threshold is set to 0.5. 50-95 It represents the average mAP at different thresholds (0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95). The larger the mAP value, the higher the detection accuracy and the better the detection performance of the algorithm.

[0083] In addition to measuring the accuracy evaluation index, the LAH-YOLO algorithm proposed in the present invention also compares the parameter quantity and computational complexity with other detection algorithms. The parameter quantity of the algorithm represents the spatial complexity of the algorithm, that is, the size of the computer memory required by the algorithm, mainly including the parameters of the convolution operation and the full connection operation. The smaller the parameter quantity, the lower the requirement of the detection algorithm for hardware resources, and the easier it is to implement deployment. The number of floating point operations (Floating Point Operations, FLOPs) represents the number of floating point operations required by the algorithm when performing a forward propagation. Since the number of floating point operations is generally large, it is usually measured in G, and GFLOPs is used as an evaluation indicator. The fewer GFLOPs, the faster the calculation speed under the same hardware conditions, and the better the real-time performance of the detection algorithm.

[0084] c) Experimental results analysis

[0085] The experimental results of the comparison between the LAH-YOLO algorithm and other algorithms in the YOLO series are shown in Table 1. As can be seen from Table 1, the number of parameters of the LAH-YOLO algorithm is 5.9M, the GFLOPs is 13.7, and the mAP is 50-95 The mAP is 64.1%. 50 97.5%; the YOLOv5-s algorithm has 7.0M parameters, 15.9 GFLOPs, and mAP 50-95 The mAP is 61.9%. 50 The LAH-YOLO algorithm reduces the number of parameters by 15.7%, GFLOPs by 13.8%, and mAP by 96.3%. 50 Improved by 1.2%, mAP 50-95 It has increased by 2.2%. The number of parameters and computational complexity of YOLOX-s, YOLOv6-s, and YOLOv8-s algorithms are much higher than those of the above two algorithms, and their detection accuracy is low. From the experimental results, it can be seen that algorithms such as YOLOX, YOLOv6, and YOLOv8 that use anchor-free detection mechanisms have poor detection effects on rail surface defect datasets and cannot identify small target defects well; while the YOLOv5 algorithm performs better than advanced algorithms such as YOLOv8 on rail surface defect datasets, proving that the anchor-base detection mechanism is more suitable for small target detection scenarios. The LAH-YOLO algorithm is improved based on the YOLOv5-s algorithm, and its detection performance is the best. It reduces the number of parameters while maintaining the highest detection accuracy. Experimental data prove that the algorithm research conducted by the present invention has achieved certain results in both parameter quantity and detection accuracy.

[0086] Table 1 Experimental results of performance comparison between LAH-YOLO algorithm and other algorithms in the YOLO series

[0087]

[0088]

[0089] The detection results of different detection algorithms for the block drop caused by nuclear damage are as follows Figure 5 As shown, from left to right are the original image, YOLOv5-s, YOLOX-s, YOLOv6-s, YOLOv8-s, and LAH-YOLO. Figure 5 It can be seen that all the detection algorithms in this experiment did not miss the falling blocks on the rail surface. However, when detecting defects located in the upper part of the image, the YOLOv8 algorithm generated two overlapping prediction boxes, and the redundant prediction boxes generated were not filtered out, resulting in poor regression positioning performance. The prediction boxes generated by the YOLOX and YOLOv6 algorithms deviate from the defect contour, and the deviation of the YOLOX algorithm is greater. Both algorithms do not perform regression positioning of the defect boundary well. In terms of confidence score, it can be seen from the detection results that the LAH-YOLO algorithm has the highest confidence score, with a confidence score of 0.93 for the upper defect and a prediction score of 0.88 for the lower defect. The confidence scores of the YOLOv5 algorithm for the two defects are 0.87 and 0.82 respectively. The confidence scores of the YOLOv6 algorithm for upper defects are 0.69, and for lower defects are 0.80; the confidence scores of the YOLOX algorithm for upper defects are 0.71, and for lower defects are 0.86; the confidence score of the YOLOv8 algorithm for lower defects is 0.7. The confidence scores predicted by the latter three algorithms are low, and the prediction and classification performance is poor. From the actual detection effect, the regression positioning ability and confidence score of the LAH-YOLO algorithm are better than other detection algorithms in the YOLO series, and the detection effect on the rail surface defect dataset is outstanding.

[0090] The detection results of different detection algorithms for the peeling block caused by contact fatigue are as follows Figure 6 As shown, from left to right are the original image, YOLOv5-s, YOLOX-s, YOLOv6-s, YOLOv8-s, and LAH-YOLO. Figure 6It can be seen that when detecting peeled-off block defects, the YOLOX algorithm and the YOLOv6 algorithm also generated prediction boxes for some peeled parts that were not defects, resulting in false detections; the YOLOX algorithm also missed detections of large peeled parts, and the above two algorithms had poor detection accuracy on the rail surface defect dataset. The LAH-YOLO algorithm, the YOLOv5 algorithm, and the YOLOv8 algorithm accurately detected the peeled-off block parts, among which the two-processing confidence scores of the YOLOv5 algorithm were 0.90 and 0.88 respectively; the two-processing confidence scores of the YOLOv8 algorithm were 0.75 and 0.72 respectively; and the two-processing confidence scores of the LAH-YOLO algorithm were 0.95 and 0.92 respectively. From the actual detection effect, the LAH-YOLO algorithm has accurate positioning and the highest confidence score, and its detection performance on the rail surface defect dataset is better than that of other algorithms in the YOLO series.

[0091] Based on YOLOv5-s, the present invention uses the GRF-C3 module to replace the C3 module in the backbone network, and the feature fusion network uses the AFPN network. The performance improvement effect of each module is verified by ablation experiments. The experimental results are shown in Table 2. As can be seen from Table 2, the parameter amount of the YOLOv5-s algorithm using the C3 module and the PAN feature fusion network is 7.0M, the GFLOPs is 15.9, and the mAP 50-95 The mAP is 61.9%. 50 When the GhostC3 module is used to replace the C3 module, the number of parameters is reduced by 12.9%, GFLOPs is reduced by 16.9%, but mAP 50-95 0.5% reduction, mAP 50 When the GRF-C3 module is used to replace the C3 module, the number of parameters is reduced by 5.7%, GFLOPs is reduced by 7.5%, and the detection accuracy is reduced by 0.3%. 50-95 Improved by 1.2%, mAP 50 The GRF-C3 module significantly improves the feature extraction capability of the backbone network and well balances the performance improvement in terms of parameter quantity and accuracy. When the AFPN network is used to replace the PAN network in the feature fusion network, the parameter quantity is reduced by 10.6%, GFLOPs is reduced by 6.8%, and the detection accuracy mAP is reduced by 1.3%. 50-95 Improvement of 1%, mAP 50 The improvement is 0.3%, which verifies that the progressive fusion structure and C3 module in the AFPN fusion network have a good effect on improving the detection performance. The experimental results show that the GRF-C3 module and the AFPN feature fusion network have a positive effect on improving the detection performance.

[0092] Table 2 Ablation test results of each module of LAH-YOLO algorithm

[0093]

[0094] When designing GRF-BottleNeck, we designed a variety of convolution combinations with different sizes of convolution kernels based on the theoretical idea that increasing the convolution kernel can increase the receptive field and improve the detection accuracy. We also compared the performance improvement effects of each solution through comparative experiments. The experimental results are shown in Table 3. As can be seen from Table 3, when GhostBottleNeck uses a combination of 1×1 Ghost convolution module and 3×3 depthwise separable convolution, the algorithm has 6.1M parameters, 13.2 GFLOPs, and mAP. 50-95 The mAP is 61.2%. 50 When using the combination of 3×3 Ghost convolution modules and 5×5 depth-separable convolutions, the algorithm's parameters increase by 3.3%, GFLOPs increase by 6.1%, and mAP 50-95 Improved by 0.2%, mAP 50 Improved by 0.6%. When using the convolution combination of 5×5 Ghost convolution module and 7×7 depthwise separable convolution, the algorithm parameters continue to increase, and the detection accuracy improves very little. Experimental results show that increasing the convolution kernel has a certain effect on improving detection accuracy, but after the convolution kernel increases to a certain extent, the detection effect is no longer significantly improved. When the convolution kernel size is designed to be a combination of 3×3 and 5×5, the overall performance of the algorithm is optimal. When the 3×3CBR convolution layer replaces the 3×3 Ghost convolution module in the GRF-C3 module, the algorithm parameters increase by 4.8%, GFLOPs increases by 5.0%, and mAP 50-95 Improved by 1.1%, mAP 50 The improvement is 0.4%. From the experimental results, although the Ghost convolution module has a good lightweight effect, the detection accuracy is significantly reduced compared to when the CBR convolution layer is used. Based on the analysis of the algorithm parameters and detection accuracy, the convolution scheme combination used by the GRF-C3 module has the best performance.

[0095] Table 3 Comparative experimental results of GRF-BottleNeck convolution combination scheme

[0096] Convolutional combination scheme Parameter quantity GFLOPs <![CDATA[mAP 50-95 ]]> <![CDATA[mAP 50 ]]> 1×1Ghost convolution module + 3×3DWConv 6.1M 13.2 61.2％ 96.0％ 3×3Ghost convolution module + 5×5DWConv 6.3M 14.0 61.4％ 96.6％ 5×5Ghost convolution module + 7×7DWConv 6.8M 15.4 61.7％ 96.7％ 3×3CBR convolution layer + 5×5DWConv 6.6M 14.7 62.5％ 97％

[0097] The ECA attention mechanism can focus on the extraction of important features. In order to verify its performance improvement effect, an ablation experiment was conducted, and the experimental results are shown in Table 4. As can be seen from the table, when the ECA attention mechanism is used in GRF-BottleNeck, the number of parameters and GFLOPs do not change, but the mAP 50-95 Improved by 0.6%, mAP50 The improvement is 0.2%, which proves that using the ECA attention mechanism in the backbone network's GRF-BottleNeck can enhance the feature extraction capability and is beneficial to the accuracy improvement.

[0098] Table 4 ECA attention mechanism ablation experiment results

[0099] Module Name Parameter quantity GFLOPs <![CDATA[mAP 50-95 ]]> <![CDATA[mAP 50 ]]> No ECA attention mechanism 6.6M 14.7 62.5％ 97％ With ECA attention mechanism 6.6M 14.7 63.1％ 97.2％

[0100] The effect of Wise-IoU loss on the performance of the algorithm is compared through ablation experiments. The experimental results are shown in Table 5. As can be seen from Table 5, when using Wise-IoU loss, the number of parameters and GFLOPs of the algorithm do not change, but mAP 50-95 Improved by 0.2%, mAP 50 The improvement is 0.1%, which proves that Wise-IoU loss can reduce the impact of some low-quality true value frames on the overall detection accuracy, thereby improving the detection accuracy.

[0101] Table 5 Wise-IoU loss ablation experimental results

[0102] Loss function name Parameter quantity GFLOPs <![CDATA[mAP 50-95 ]]> <![CDATA[mAP 50 ]]> CIoULoss 5.9M 13.7 63.9％ 97.4％ Wise-IoULoss 5.9M 13.7 64.1％ 97.5％

[0103] The content of the present invention is not limited to the embodiments listed, and any equivalent changes made to the technical solution of the present invention by ordinary technicians in this field after reading the description of the present invention are covered by the claims of the present invention.

Claims

1. A rail surface defect detection method based on LAH-YOLO, characterized by: Acquire a rail surface image, and input the rail surface image into a constructed LAH-YOLO defect detection network; the LAH-YOLO defect detection network includes a backbone network, a feature fusion network, and a detection head; The backbone network includes a GRF-C3 module, and the feature fusion network is an AFPN feature fusion network.

2. The rail surface defect detection method based on LAH-YOLO according to claim 1, characterized in that: The GRF-C3 module first divides the input feature map into two parts according to half the number of channels, one part passes through the 1×1 CBR convolution layer and GRF-BottleNeck, and the other part directly passes through the 1×1 CBR convolution layer. The two parts are then concatenated according to the channel dimension and finally output after passing through the 1×1 CBR convolution layer.

3. The rail surface defect detection method based on LAH-YOLO according to claim 2, characterized in that: The GRF-BottleNeck uses 3×3 and 5×5 convolution kernels, uses a CBR convolution layer before a depth-wise separable convolution, and applies the ECA attention mechanism to the residual connection part in GRF-BottleNeck. The features at the residual connection need to be processed by the ECA attention mechanism to give different weights to each channel.

4. The rail surface defect detection method based on LAH-YOLO according to claim 3 is characterized in that: The AFPN fusion network first fuses the two adjacent layers of feature maps at the bottom layer, and then performs cross fusion after upsampling and downsampling; After feature fusion, four residual units are used to continue extracting features. Each residual unit includes two 3×3 standard convolutions. After the two layers of features are fused, the fused features are cross-fused with the top layer features. After being processed by the C3 module, the features are finally input into the detection head.

5. The rail surface defect detection method based on LAH-YOLO according to claim 4 is characterized in that: The fusion principle of the three-layer AFPN feature fusion network is as follows: y 3 =a 3 ·x 1→3 +b 3 ·x 2→3 +g 3 ·x 3→3 Among them, y 3 Represents the third layer features after fusion output; α 3 , β 3 , γ 3 They represent the weight factors of the input first-layer features, second-layer features, and third-layer features on the output third-layer features. Numerically, the sum of the weight factors is 1, i.e., α 3 +β 3 +γ 3 =1;x 1→3 、x 2→3 、x 3→3 They represent the first-layer features, second-layer features, and third-layer features of the input respectively.

6. The rail surface defect detection method based on LAH-YOLO according to claim 5, characterized in that: The detection head adopts an anchor-base detection mechanism. Three detection heads perform detection on feature maps of three different scales respectively, and output five physical quantities: predicted category, x and y coordinates of the center point of the predicted box, and width and height of the predicted box.

7. The rail surface defect detection method based on LAH-YOLO according to claim 6, characterized in that: The loss function of the target detection task includes three parts: classification loss function, bounding box loss function and confidence loss function; The bounding box loss function adopts the Wise-IoU loss function, and the expression of the Wise-IoU loss function is as follows: L WIoU =r·R WIoU ·L IoU Among them, L WIoU represents the Wise-IoU loss; r represents the gradient gain coefficient; R WIoU represents the distance attention coefficient, whose value range is between [1, e), which can significantly amplify the IoU loss of ordinary anchor boxes; L IoU represents the IoU loss.

8. The rail surface defect detection method based on LAH-YOLO according to claim 7, characterized in that: The gradient gain coefficient r is expressed as follows: Wherein β represents the clustering degree to describe the quality of the anchor frame, and its value range is between [0,). It is a dynamically adjusted value to adjust the gradient gain coefficient r to the best. α and δ represent two artificially set hyperparameters for controlling the clustering degree β. The present invention sets α and δ to 1.9 and 3 respectively.

9. The rail surface defect detection method based on LAH-YOLO according to claim 8, characterized in that: Distance attention coefficient R WIoU The expression is as follows: Where x and y represent the horizontal and vertical coordinates of the center point of the prediction box, respectively. gt ,y gt Respectively represent the horizontal and vertical coordinates of the center point of the truth box, W g , H g Respectively represent the width and height of the smallest bounding box of the prediction box and the true value box. The superscript * indicates that W g , H g Separating operations from the computation graph prevents R WIoU Produces gradients that hinder convergence.

10. The rail surface defect detection method based on LAH-YOLO according to claim 9, characterized in that: IoU loss L IoU The expression is as follows: Among them, S represents the area of ​​the prediction box, S gt Represents the area of ​​the true value box. The IoU loss is obtained by calculating the ratio of the intersection area of ​​the predicted box and the true value box to the union area.

Images