A turnout machine part detection method, device and medium based on a YOLOv8 model

Abstract

The application discloses a kind of based on YOLOv8 model's switch machine part detection method, equipment and medium, the method includes: labeling target part in switch machine part image, generates target image;Swin transformer network replaces the original main network darknet53 of YOLOv8 model, residual attention module is introduced in Swin transformer network;Based on Swin transformer network combines spatial pyramid pooling layer, simultaneously using new loss function Focal-EIoU replaces the original loss function CIoU of YOLOv8 model, constructs improved YOLOv8 model;The target image is input to the improved YOLOv8 model, generates prediction result.Based on the detection result of improved YOLOv8 model output realizes the accurate detection of target part.

Description

Technical Field

[0001] The present application relates to the field of detection technology, and in particular to a method, device, and medium for detecting switch machine parts based on a YOLOv8 model. Background Art

[0002] Switches are crucial in rail transit turnout switching equipment. They are used for train reversing and shunting, and play an important role in train driving safety. Failure of switch machine parts can cause train delays and seriously affect the normal operation of the train. Traditional fault repair methods mainly rely on maintenance personnel's visual inspection and experience judgment to perform manual inspection and repair. This requires a lot of experience accumulation, high manual inspection costs, poor real-time performance, low fault detection efficiency, and lack of robustness to changes in feature diversity. It can no longer meet the needs of today's industrial inspection. At present, three-dimensional tracking and registration technology is becoming increasingly popular. By detecting and locating objects in real maintenance scenarios and then tracking and registering them, the virtual objects can be superimposed on the real maintenance environment and visualized, which can greatly improve the maintenance efficiency of maintenance personnel and reduce maintenance costs.

[0003] In order to accurately track and register and achieve efficient fault repair, accurate detection and positioning of switch machine parts becomes the key. The detection accuracy will directly affect the efficiency of fault repair. Therefore, in order to achieve accurate detection of switch machine parts, a method for accurate detection of switch machine parts is urgently needed. Summary of the Invention

[0004] The purpose of this application is to provide a switch machine parts detection method, equipment and medium based on the YOLOv8 model to solve the problem of low accuracy in switch machine parts detection.

[0005] To achieve the above objectives, this application provides the following solutions:

[0006] In a first aspect, the present application provides a switch machine parts detection method based on the YOLOv8 model, comprising: marking a target part in a switch machine parts image to generate a target image; the target image includes a real bounding box for locating the target part;

[0007] The Swin transformer network replaces the original backbone network darknet53 of the YOLOv8 model, and a residual attention module is introduced into the Swin transformer network; based on the Swin transformer network combined with the spatial pyramid pooling layer, a new loss function Focal-EIoU is used to replace the original loss function CIoU of the YOLOv8 model to construct an improved YOLOv8 model; the residual attention module is used to focus on the features corresponding to the target part so that the Swin transformer network generates a feature map corresponding to the target part; the spatial pyramid pooling layer is used to perform maximum pooling or average pooling on the feature map to obtain a feature vector; the new loss function Focal-EIoU is used to calculate the difference between the predicted aspect ratio corresponding to the predicted bounding box of the target part and the actual aspect ratio corresponding to the actual bounding box of the target part;

[0008] The target image is input into the improved YOLOv8 model to generate a prediction result; the prediction result includes scale information corresponding to the predicted bounding box of the target part and scale information corresponding to the true bounding box; wherein the scale information corresponding to the predicted bounding box includes a predicted aspect ratio and a predicted center point coordinate; the scale information corresponding to the true bounding box includes a true aspect ratio and a true center point coordinate.

[0009] In a second aspect, the present application provides a computer device comprising: a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the steps of the switch machine parts detection method based on the YOLOv8 model.

[0010] In a third aspect, the present application provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the steps of the switch machine parts detection method based on the YOLOv8 model.

[0011] According to the specific embodiments provided in this application, this application discloses the following technical effects:

[0012] This application first replaces the original YOLOv8 model's original backbone network, darknet53, with the Swin Transformer network. It then introduces a residual attention module into the Swin Transformer network. This module can more accurately focus on the target part, thereby achieving global and detailed feature extraction of the target part. This module can enhance the target part's representation capabilities and suppress the interference of redundant information, making the feature maps generated by the Swin transformer network more accurate than those generated by the original backbone network. Furthermore, by combining the spatial pyramid pooling layer with the feature map to perform maximum pooling or average pooling to obtain feature vectors, the extraction capability of the feature map is further enhanced. Then, by introducing a new loss function Focal-EIoU to replace the original loss function CIoU, since the original loss function CIoU is not accurate in predicting the aspect ratio of the target part, this application adopts the new loss function Focal-EIoU to enable the improved YOLOv8 model to recalculate the scale information of the prediction box of each target part to replace the scale information of the prediction box obtained by the YOLOv8 model of the original loss function CIoU, thereby achieving accurate prediction of the aspect ratio of the target part. Finally, the target part is accurately located based on the detection results output by the improved YOLOv8 model, thereby achieving accurate detection of the target part. BRIEF DESCRIPTION OF THE DRAWINGS

[0013] In order to more clearly illustrate the embodiments of the present application or the technical solutions in the prior art, the following briefly introduces the drawings required for use in the embodiments. Obviously, the drawings described below are only some embodiments of the present application. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work.

[0014] Figure 1 A schematic diagram of a flow chart of a switch machine parts detection method based on a YOLOv8 model provided in an embodiment of the present application;

[0015] Figure 2 A Swintransformer network structure diagram of a switch machine parts detection method based on a YOLOv8 model provided in an embodiment of the present application;

[0016] Figure 3 This is a diagram of the overall network structure of an improved YOLOv8 model for a switch machine parts detection method based on a YOLOv8 model provided in an embodiment of the present application;

[0017] Figure 4 A structural diagram of a layered visual converter based on a displacement window for a switch machine parts detection method based on a YOLOv8 model provided in an embodiment of the present application;

[0018] Figure 5 A diagram showing the shift window self-attention calculation of a switch machine parts detection method based on the YOLOv8 model provided in an embodiment of the present application;

[0019] Figure 6 This is an overall structural diagram of the residual attention module of a switch machine parts detection method based on the YOLOv8 model provided in an embodiment of the present application;

[0020] Figure 7 A schematic diagram of a parameter-free attention module for a switch machine parts detection method based on a YOLOv8 model provided in an embodiment of the present application;

[0021] Figure 8 A comparison diagram of two sets of motor detection and prediction frames before and after improvement of a switch machine parts detection method based on the YOLOv8 model provided in an embodiment of the present application;

[0022] Figure 9 A comparison diagram of the motor prediction frames before and after the improvement of a switch machine parts detection method based on the YOLOv8 model provided in an embodiment of the present application;

[0023] Figure 10 This is a comparison chart of the detection results of other parts before and after the improvement of a switch machine parts detection method based on the YOLOv8 model provided in an embodiment of the present application. DETAILED DESCRIPTION

[0024] The following will be combined with the drawings in the embodiments of this application to clearly and completely describe the technical solutions in the embodiments of this application. Obviously, the embodiments described are only part of the embodiments of this application, not all of the embodiments. Based on the embodiments in this application, all other embodiments obtained by ordinary technicians in this field without making creative efforts are within the scope of protection of this application.

[0025] With the rapid development of deep learning technology, researchers have begun applying it to parts inspection. Pechenin et al. combined convolutional neural networks with data classification and trained multiple convolutional neural networks to achieve accurate classification and detection of industrial parts. Based on the SSD model, Chen Guanqi et al. proposed a multi-scale fusion part detection algorithm, called MDSSD. On three mechanical parts datasets, the average accuracy was 6.21% higher than that of SSD. Guo Fei et al. proposed an improved Fast R-CNN part detection algorithm, incorporating the Inception architecture into Fast R-CNN. Verified on a self-developed mechanical parts dataset, the algorithm achieved superior detection accuracy to Fast R-CNN. Liu Xueping et al. from Tsinghua University combined a designed adaptive edge error function with an improved PSO algorithm to propose the YOLOv3-AEEF part detection algorithm, which outperforms YOLOv3 in the presence of multiple target parts. Deep learning-based methods are more robust to large variations in part size, simple textures, and varying lighting conditions. Single-stage detection methods offer faster detection speeds, lower equipment requirements, and ease of deployment in industrial scenarios compared to dual-stage methods. In light of this, a single-stage target detection method was used to detect switch machine parts. Xing Jinchao et al. proposed an improved YOLOv5 algorithm. This algorithm improves the channel domain of the CBAM attention mechanism to address the problem of missing channel information caused by dimensionality reduction. The improved CBAM was added to the YOLOv5 backbone network, enhancing its feature extraction capabilities. However, due to the limited receptive field of the spatial attention module, target feature information is not fully extracted and it is easily disturbed by complex background information, which affects detection accuracy. To give the model a larger receptive field and better capture global features, many experts and scholars have attempted to apply the Transformer to the field of computer vision. Due to the significant advantages of the encoder-decoder architecture in feature extraction, the Vision Transformer (ViT) pioneered the application of the Transformer architecture to image classification to extract global image features. However, its self-attention mechanism has a very large number of parameters when it comes to image size and lacks scalability. Swin Transformer is based on the idea of ​​the ViT model and pioneered the introduction of a sliding window mechanism, allowing the model to learn cross-window information. At the same time, through the downsampling layer, the model can process super-resolution images, saving computational complexity. Its performance is better than that of convolutional neural networks. It has replaced the classic convolutional neural network architecture and has become a common backbone network in the field of computer vision.

[0026] According to current research, introducing a convolutional neural network structure early in the Vision Transformer (ViT) model helps enhance the model's robustness to hyperparameters and convergence speed, which has a beneficial effect on improving model accuracy. In convolutional neural networks, convolution operations excel at extracting local features, but have certain limitations in capturing global feature representations. Although the cascaded self-attention module in ViT can capture long-range feature dependencies, it still ignores local feature details, resulting in insufficient feature extraction for switch parts. Furthermore, the YOLOv8 model not only lacks the ability to extract features for switch parts, but also suffers from inaccurate output of detection results for switch parts.

[0027] In order to make the above-mentioned purposes, features and advantages of the present application more obvious and easy to understand, the present application is further described in detail below with reference to the accompanying drawings and specific implementation methods.

[0028] like Figure 1 As shown, the present application provides a switch machine parts detection method based on the YOLOv8 model, including:

[0029] Step 101: annotate a target part in a switch machine part image to generate a target image; the target image includes a real bounding box for locating the target part.

[0030] In some embodiments, labeling the target parts in the switch machine part image specifically includes: screening all target parts that are larger than a preset minimum size among all target parts for labeling; for visible parts of each target part, drawing frames closely along the tangent to the edge of the target part; if there are occluded parts in the labeled target parts, the occluded parts of the target parts are completed according to the shape of the target parts before labeling; discarding abnormal target parts in the switch machine part image, and drawing frames for each remaining target part separately to generate target frames; the abnormal target parts include blurred, dark and overexposed target parts, and the frame edge of the target frame cannot coincide with the boundary of the switch machine part image; the target frame is the real boundary frame.

[0031] In some embodiments, before labeling the target part in the switch machine part image, the method further includes:

[0032] Collecting a data set, the data set including an initial switch machine part image;

[0033] Normalization preprocessing is performed on the initial switch machine part image to obtain the switch machine part image.

[0034] In practical applications, to verify the effectiveness of the algorithm and accurately detect switch machine parts, a self-made 3D-printed ZD6 DC electric switch machine was used. Images of the switch machine parts were captured using a mobile phone camera as a dataset. A total of 2,000 images containing eight target categories were generated through flipping, rotation, cropping, and scaling. Finally, the images were annotated using labelImg software. After annotation, corresponding labels were generated and the target images were randomly assigned to each category in a ratio of 7:1.5:1.5 to form the training, validation, and test sets for the improved YOLOv8 model.

[0035] Step 102: The original backbone network darknet53 of the YOLOv8 model is replaced by the Swin transformer network, and the residual attention module (ResAM) is introduced into the Swin transformer network; based on the Swin transformer network combined with the spatial pyramid pooling layer, the new loss function Focal-EIoU is used to replace the original loss function CIoU of the YOLOv8 model to construct an improved YOLOv8 model; the residual attention module is used to focus on the features corresponding to the target part so that the Swin transformer network generates a feature map corresponding to the target part; the spatial pyramid pooling layer is used to perform maximum pooling or average pooling on the feature map to obtain a feature vector; the new loss function Focal-EIoU is used to calculate the difference between the predicted aspect ratio corresponding to the predicted bounding box of the target part and the actual aspect ratio corresponding to the actual bounding box of the target part;

[0036] Specifically, in the detection of switch machine parts, due to the inherent distribution pattern of switch machine parts, shift window multi-head self-attention (SW-MSA) can better segment part information and retain key information, thereby improving the network's feature extraction capabilities and solving the problem of lack of information exchange between non-overlapping windows. In view of this, based on the distribution characteristics of switch machine parts, the Swin Transformer network is used to improve darknet53 and combine it with the spatial pyramid pooling layer (SPPF). A new feature extraction network is constructed, namely the combination of the Swin Transformer network and the spatial pyramid pooling layer (SPPF), which can more accurately detect switch machine parts.

[0037] In some embodiments, the improved YOLOv8 model further includes an image block segmentation module; the image block segmentation module is used to segment the switch machine part image into a plurality of image blocks; the Swin transformer network is used to generate a feature map of the target part in the switch machine part image based on the image blocks.

[0038] like Figure 2 As shown in the figure, at the beginning of the input, the PatchPartition convolutional layer of the image segmentation module with a patch size of 4 divides the image data into several non-overlapping image patches, converting the feature dimension of each image patch to 4×4×3=48. In practice, the Patch Partition module and the Linear Embedding module are implemented through a two-dimensional convolutional layer (the output channel is the Embedding dimension, the convolution kernel size is the image patch size, and the stride size is the image patch size).

[0039] In some embodiments, the Swin transformer network includes: a first structure, a second structure, a third structure and a fourth structure; wherein the first structure includes: a linear embedding module and a hierarchical visual converter based on a displacement window; the second structure includes an image block merging module and a hierarchical visual converter based on a displacement window; the third structure and the fourth structure have the same structure as the second structure.

[0040] In some embodiments, the residual attention module is respectively arranged between the linear embedding module of the first structure and the hierarchical visual converter based on the displacement window, between the image block merging module of the second structure and the hierarchical visual converter based on the displacement window, between the image block merging module of the third structure and the hierarchical visual converter based on the displacement window, and between the image block merging module of the fourth structure and the hierarchical visual converter based on the displacement window.

[0041] Specifically, in the first stage (i.e., the first structure), the linear embedding module (Linear Embedding) is used to adjust the number of channels to C. The subsequent shift window-based hierarchical visual transformer (Swin Transformer Block) generally appears in a two-stage series structure, and different self-attention calculation methods are selected according to whether the current shift window-based hierarchical visual transformer (Swin Transformer Block) is an odd or even number. Then, each stage after the first stage is composed of an image block merging module (Patch Merging) and several shift window-based hierarchical visual transformers (Swin Transformer Blocks). The function of the image block merging module Patch Merging layer is to reduce the resolution of the target image at the beginning of each stage and perform a downsampling operation. Assuming the target image input to the Patch Merging module is 4×4, the Patch Merging module will divide each 2×2 adjacent pixel into an image patch. Pixels at the same position in the image patch are then spliced ​​together to obtain four feature maps. After embedding and splicing, the image is passed through a layer normalization layer (LayerNorm). At this point, the length and width are reduced to halve, and the channel dimension is quadrupled. Finally, a fully connected layer adjusts the channel dimension to double, ultimately reducing the channel dimension to 2C. In other words, after passing through the Patch Merging module, the length and width of the data are reduced to halve, and the depth is doubled.

[0042] Reference Figure 4 The feature extraction network, a combination of a Swin transformer network and a spatial pyramid pooling layer (SPPF), replaces the original backbone network and is responsible for extracting the features of the target part from the input target image. The neck network further processes these features, performing methods such as feature fusion and upsampling. Finally, the head network performs the final prediction based on these processed features and outputs the predicted result.

[0043] Reference Figure 4, the structure of the displacement window-based hierarchical visual transformer (Swin Transformer Block) is similar to that of the Transformer block, the Swin Transformer Block replaces the multi-head self-attention layer (MSA) and the multi-layer perception block (MLP) in the Transformer Block with a window multi-head self-attention (W-MSA) and a shift window multi-head self-attention (SW-MSA), wherein, Figure 4 represents the output feature of the shift window multi-head self-attention module of the Lth block; z l represents the output feature of the multi-layer perception module of the Lth block; represents the output feature of the shift window multi-head self-attention module of the L+1th block; Z l+1 represents the output feature of the multi-layer perception module of the L+1th block; Z l-1 represents the output feature of the multi-layer perception module of the L-1th block.

[0044] The W-MSA and the SW-MSA are alternately performed in the consecutive Swin Transformer Blocks, for the flattened feature map with an input format of (h, w, C) of the nth layer Swin Transformer Block, h represents the height of the prediction box; w represents the width of the prediction box, and C represents the dimension. The size of the flattened feature map here is obtained by flattening and splicing the overlapped feature maps. The computational complexity of the MSA is shown in formula (1) and the computational complexity of the W-MSA is shown in formula (2):

[0045] Ω(MSA) = 4hwC 2 + 2 (hw) 2 C (1)

[0046] Ω(W-MSA) = 4hwC 2 + 2M 2 hwC (2)

[0047] wherein, Ω(MSA) represents the computational complexity of the MSA; Ω(W-MSA) represents the computational complexity of the W-MSA; M represents the window size of each feature map in the input flattened feature map. The input feature map of the previous layer Swin Transformer Block is divided into a new feature map size and a window by the shift window self-attention mechanism of the hierarchical level feature downsampling module and Figure 7 Figure 7 ​​The X in the image represents the features of the target part. The SwinTransformer Block window of the next layer is then moved, allowing the non-overlapping windows of the previous layer to be included in the calculation again, greatly increasing the receptive field of the model.

[0048] Among them, the Residual Attention Module (ResAM) includes a parameter-free attention module and a convolution operator.

[0049] Specifically, the introduction of convolutional structure in ViT can enhance the model's robustness to hyperparameters and convergence speed, showing great potential. In convolutional neural network operations, convolution is better at extracting local features, and has certain limitations in its ability to extract global features. The Transformer structure, on the contrary, will ignore local feature details. Therefore, adding ResAM to the Swin Transformer network has a better and more efficient receptive field than the single structure of the Swin Transformer network, and can better balance the extraction of local and global feature details. In target detection, the attention mechanism can help the model pay more attention to the information of the target parts in the target image, enhance the representation ability of the target parts, and thus improve the accuracy of detecting switch parts. Therefore, this application uses a residual attention module (Residual Attention Module, ResAM) with a parameter-free attention module and a convolution operator as basic components, and then embeds it into the four stages (i.e., four structures) of the Swin Transformer network.

[0050] like Figure 6 As shown in the figure, first, the input features are passed through a 1×1 convolution, then through the MBConv3×3 convolution with a stride of 2 and the self-improving attention mechanism SimAM module

[23] , and then through a 1×1 convolution for channel dimension increase. Finally, the residual connection structure is introduced to optimize the overall training process.

[0051] Specifically, MBConv combines two key concepts: inverted residual structure and lightweight depth-wise separable convolution. The inverted residual structure includes the steps of first reducing the dimension, then performing nonlinear transformation, and finally increasing the dimension, which helps to learn more effective feature representation. Lightweight depth-wise separable convolution reduces the computational cost by decomposing the standard convolution into depth-wise convolution and point-by-point convolution. The main advantage of MBConv is that it can improve the expressiveness of the model while maintaining a low computational cost. This is very beneficial for using deep learning models in resource-constrained environments such as mobile devices and embedded systems. When the network is very deep, in addition to increasing computing resource consumption and model overfitting problems, gradient vanishing / exploding problems will also occur, resulting in the inability to update shallow network parameters. Therefore, the residual connection structure is used to optimize the overall training process and improve network performance.

[0052] like Figure 7 As shown in the figure, the Self-Improving Attention Mechanism (SimAM) is a simple and very effective attention module. Unlike the existing channel / spatial attention modules, this module can derive 3D attention weights for feature maps without additional parameters. Specifically, based on the well-known neuroscience theory, an optimized energy function is proposed to explore the importance of neurons, and a fast analytical solution is further derived for the energy function. By calculating the attention weights and focusing the model's attention on the most relevant features, the robustness and generalization ability of the model can be effectively improved, thereby improving the model's detection accuracy for target parts. In order to better achieve attention, it is necessary to evaluate the importance of each neuron. Activating neurons usually inhibits surrounding neurons, that is, spatial inhibition. The energy function used to evaluate the importance of different neurons in SimAM is shown in formula (3):

[0053]

[0054] In the formula, is the minimum energy, t is the target neuron, are the mean and variance of all neurons except t, and λ is a hyperparameter. The lower the energy of , the more different the target neuron t is from the surrounding neurons and the more important it is. The importance of a neuron is expressed by get.

[0055] In some embodiments, it is characterized in that the new loss function Focal-EIoU is:

[0056]

[0057] L (Focal-EIoU) =IoU γ L EIoU..... (5)

[0058] Among them, L EIoU represents the difference between the predicted aspect ratio and the true aspect ratio corresponding to the improved YOLOv8 model; h c represents the height of the minimum bounding rectangle of the predicted bounding box and the true bounding box; w c Indicates the width of the minimum bounding rectangle of the predicted bounding box and the true bounding box; L (Focal-EIoU) represents the loss of the improved YOLOv8 model; γ represents the hyperparameter that controls the degree of negative sample suppression; ρ represents the Euclidean distance between the predicted center point coordinates and the true center point coordinates of the predicted bounding box; IoU represents the set threshold; w represents the width of the predicted box; w gt represents the width of the real box; h represents the height of the predicted box; b represents the center point of the predicted box; b gt Indicates the center point of the real box; h gt Indicates the height of the true box.

[0059] Among them, the calculation formula of the original loss function CIoU of the YOLOv8 model is shown in formula (6):

[0060]

[0061] Among them, L CIoU It represents the difference between the predicted aspect ratio and the true aspect ratio corresponding to the YOLOv8 model; ρ represents the Euclidean distance between the predicted center point coordinates of the predicted bounding box and the true center point coordinates; c represents the minimum diagonal length of the predicted bounding box and the true bounding box; υ represents the aspect ratio consistency of the predicted bounding box and the true bounding box, measured by the tan angle value; α represents a balance coefficient, which is given priority according to the set IoU threshold. When 0 <IoU<0.5时,α的系数为0,当IoU≥0.5时,IoU的阈值越大,α的系数越大。

[0062] Specifically, the original loss function CIoU is not accurate in predicting the aspect ratio of the target part. When low-quality regression samples have a greater impact on the regression loss, it will make it difficult to further optimize the regression samples with good quality, resulting in an imbalance between positive and negative samples. Therefore, this application adopts a new loss function Focal-EIoU to enable the improved YOLOv8 model to recalculate the scale information of the prediction box of each target part instead of the scale information of the prediction box obtained by the YOLOv8 model of the original loss function CIoU, and at the same time combine L (Focal-EIoU) To alleviate the problem of imbalance between positive and negative samples, we finally achieved accurate prediction of the aspect ratio of the target part and finally achieved accurate detection of the target part based on the detection results output by the improved YOLOv8 model.

[0063] Step 103: inputting the target image into the improved YOLOv8 model to generate a prediction result; the prediction result includes scale information corresponding to a predicted bounding box of the target part and scale information corresponding to a real bounding box; wherein the scale information corresponding to the predicted bounding box includes a predicted aspect ratio and a predicted center point coordinate; and the scale information corresponding to the real bounding box includes a real aspect ratio and a real center point coordinate.

[0064] To quantitatively analyze the detection performance of the YOLO-SRF algorithm, an ablation experiment was performed using YOLOv8n under the YOLOv8 model as the benchmark model, as follows:

[0065] First step: analysis of the backbone network experiment.

[0066] To verify the effectiveness of the improved YOLOv8 model, the experiment compared YOLOv8n, YOLOv8x, and the improved Swin Transformer network based on the original YOLOv8n, and compared the recall R and mAP50 changes with and without the use of ResAM. The comparison results are shown in Table 1.

[0067] Table 1: Ablation experiment results of the backbone network

[0068]

[0069]

[0070] According to Table 1, the Swin Transformer improved the recall R by 2% and the mAP50 by 1.2% compared to the original YOLOv8n network; the recall R improved by 2.4% and the mAP50 improved by 0.7% compared to the YOLOv8x network; with the addition of the ResAM module, the recall R improved by 1.3% and the mAP50 improved by 0.3% compared to the original YOLOv8n network; the recall R improved by 1.7% and the mAP50 decreased by 0.2% compared to the YOLOv8x network, which is because after the addition of the ResAM module to the YOLOv8n, the network layer, model size, and algorithm complexity are much smaller than those of the YOLOv8x; after the construction of the new feature extraction network by the Swin Transformer and the ResAM, the recall R and the mAP50 values are higher than those of the YOLOv8n and the YOLOv8x, with an increase of 2.4, 2.4 and 3.0, 1.9 percentage points, respectively, which verifies the superiority of the proposed network.

[0071] Second step: analysis of the ResAM experiment.

[0072] In order to analyze the superiority of the proposed ResAM, ResAM is added to the YOLOv8-Swin Transformer and other mainstream attention modules SE and CBAM are experimentally analyzed. The results are shown in Table 2.

[0073] Table 2 ResAM experimental results

[0074]

[0075] According to Table 2, adding the SimAM attention module to ResAM has the highest detection accuracy. Compared with the YOLOv8-SwinTransformer network, mAP50 is improved by 1.2% and mAP50-95 is improved by 4.6%. It can be seen that ResAM containing the SimAM attention module has significantly improved various indicators compared to replacing other mainstream attention modules, and the number of parameters is smaller and the accuracy is better than other mainstream attention modules. The results prove that the ResAM module effectively improves the performance of the model and is superior to other mainstream attention modules.

[0076] Step 3: Experimental analysis of loss function.

[0077] In order to analyze and verify the effectiveness of the new loss function Focal-EIoU, the loss functions of CIoU, EIoU and Focal-EIoU are compared based on YOLOv8-Swin TransformerResAM. The results are shown in Table 3.

[0078] Table 3 Loss function experimental results

[0079]

[0080] According to Table 3, compared with the original loss function CIoU, the new loss function Focal-EIoU improves the recall rate R by 0.3% and mAP50 by 0.3%. The accuracy is better than the other two mainstream loss functions, which verifies the effectiveness of the new loss function Focal-EIoU in improving accuracy.

[0081] Step 4: Comparative experimental analysis.

[0082] The YOLOv8 model and the improved YOLOv8 model (YOLOv8-Swin Transformer ResAMFocal-EIoU, YOLO-SRF) were selected to test the part images. YOLOv8n was used as the baseline model, and representative images of each category in the test set were selected for qualitative analysis. Figure 8 、 Figure 9 and Figure 10As shown in the figure, the left side is the detection result diagram of the YOLOv8 model, and the right side is the detection result diagram of the improved YOLOv8 model.

[0083] like Figure 8 As shown, the motor pointed by the white arrow in the first group of left pictures lacks part of the fuselage. After the original backbone network has been down-sampled multiple times to extract features, the target detail information is severely lost, resulting in missed detection. However, YOLO-SRF uses the improved backbone network Swin Transformer, which can more effectively learn feature representations. At the same time, ResAM effectively restores the detail information of the target parts, thereby avoiding missed detections. The motor pointed by the white arrow in the second group of left pictures has unclear features due to the complex background and occlusion of screws, resulting in missed detections. The method of this application is more accurate than the prediction box in the left picture (indicated by the white arrow) while detecting all parts. As shown Figure 9 The predicted bounding boxes obtained by YOLO-SRF are more accurate and better define the scale information of the target parts, demonstrating the effectiveness of the new loss function Focal-EIoU. Therefore, YOLO-SRF, which is an improved YOLOv8 model, can better locate the position of target parts in the switch machine and suppress the interference of redundant information, significantly improving detection accuracy and detection capabilities.

[0084] Furthermore, to verify that the proposed YOLO-SRF algorithm outperforms other algorithms in object detection, we conducted comparative experiments with the currently mainstream single-stage object detection algorithms SSD, YOLOv5, YOLOv8n, and YOLOv8x. Using recall R and mAP50 as evaluation criteria, we demonstrate the detection results of different object detection methods on eight categories of track component datasets: shift contactors, automatic switches, speed reducers, spindles, indicator rods, operating levers, motors, and rack blocks. The experimental results are shown in Table 4 (unit: %), where All represents the average precision of the eight categories.

[0085] Table 4 Detection results of different target detection algorithms

[0086]

[0087] Table 4 shows that the YOLO-SRF algorithm of this application outperforms other single-stage mainstream algorithms in terms of duplicate detection rate (R) and mAP50 values ​​for eight parts categories. Compared to SSD, YOLOv5, YOLOv8, YOLOv8n, and YOLOv8x, the average duplicate detection rate (R) for the eight parts categories increased by 3.2, 1.0, 2.9, and 3.3, respectively; and the average precision (mAP50) for the eight parts categories increased by 6.0, 2.9, 2.7, and 2.2, respectively. Among them, the mAP50 values ​​for five parts categories—shift contactors, automatic switches, speed reducers, indicator levers, and operating levers—were higher than those of other algorithms. The mAP50 value for the spindle part was higher than that of SSD and on par with YOLOv5, YOLOv8n, and YOLOv8x. The mAP50 value for the motor part was higher than that of SSD and YOLOv8x and on par with YOLOv5 and YOLOv8n.

[0088] This experiment demonstrates that the YOLO-SRF algorithm is more robust in situations where complex backgrounds, cluttered objects, and occlusions are present in switch machine images. The experimental results demonstrate that, compared to the aforementioned classic algorithms, the YOLO-SRF algorithm significantly improves the detection accuracy of target parts in switch machine images and can more accurately locate target parts, achieving the best overall target detection capabilities.

[0089] In some embodiments, the training process of the improved YOLOv8 model specifically includes:

[0090] A preset batch of target images is used as input to the improved YOLOv8 model, and scale information corresponding to the predicted bounding box of the target part in the target image and scale information corresponding to the true bounding box of the target part are used as output. The YOLOv8 model is trained until the difference between the predicted aspect ratio corresponding to the predicted bounding box and the true aspect ratio corresponding to the true bounding box is less than a first preset threshold, and the distance between the predicted center point coordinates of the predicted bounding box and the true center point coordinates of the true bounding box is less than a second preset threshold, and the training is stopped.

[0091] This application proposes a switch machine parts detection method based on the YOLOv8 model. In view of the inherent characteristics of switch machine parts, Swin Transformer is combined with SPPF to construct a new feature extraction network, and a new module ResAM is proposed. It is combined with the improved feature extraction network, and finally the new loss function Focal-EIoU is used to improve the original loss function CIoU. The improved YOLOv8 model is trained on a self-made rail transit dataset and then put into the test set for testing. The experimental results show that the improved YOLOv8 model (YOLO-SRF) has good target detection capabilities and excellent detection accuracy. And from the comparative experimental analysis, it is found that the detection accuracy of the YOLO-SRF algorithm is higher than the original algorithm and other classic target detection algorithms.

[0092] In an exemplary embodiment, a computer device is further provided, including a memory and a processor. The memory stores a computer program, and the processor implements the steps in the above method embodiments when executing the computer program.

[0093] In an exemplary embodiment, a computer-readable storage medium is provided, storing a computer program. When the computer program is executed by a processor, the steps in the above-mentioned method embodiments are implemented.

[0094] It should be noted that the user information (including but not limited to user device information, user personal information, etc.) and data (including but not limited to data used for analysis, stored data, displayed data, etc.) involved in this application are all information and data authorized by the user or fully authorized by all parties, and the collection, use and processing of relevant data must comply with relevant regulations.

[0095] Those skilled in the art will appreciate that all or part of the processes in the above-mentioned embodiments can be implemented by instructing the relevant hardware through a computer program. The computer program can be stored in a non-volatile computer-readable storage medium. When the computer program is executed, it can include the processes of the embodiments of the above-mentioned methods. In particular, any reference to memory, database or other media used in the embodiments provided in this application may include at least one of non-volatile and volatile memory. Non-volatile memory may include read-only memory (ROM), magnetic tape, floppy disk, flash memory, optical memory, high-density embedded non-volatile memory, resistive random access memory (ReRdM), magnetic random access memory (MRdM), ferroelectric random access memory (FRdM), phase change memory (PCM), graphene memory, etc. Volatile memory may include random access memory (RdM) or external cache memory, etc. By way of illustration and not limitation, RdM may be in various forms, such as static random access memory (SRdM) or dynamic random access memory (DRdM).

[0096] The databases involved in the various embodiments provided herein may include at least one of a relational database and a non-relational database. Non-relational databases may include, but are not limited to, distributed databases based on blockchains. The processors involved in the various embodiments provided herein may include, but are not limited to, general-purpose processors, central processing units, graphics processing units, digital signal processors, programmable logic units, data processing logic units based on quantum computing, and the like.

[0097] The technical features of the above embodiments can be combined arbitrarily. In order to make the description concise, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.

[0098] This document uses specific examples to illustrate the principles and implementation methods of this application. The description of the above examples is only intended to help understand the method and core concept of this application. At the same time, for those skilled in the art, based on the concept of this application, there may be changes in the specific implementation methods and application scope. In summary, the content of this specification should not be understood as limiting this application.

Claims

1. A switch machine parts detection method based on the YOLOv8 model, characterized in that: The method comprises: Annotating a target part in the switch machine part image to generate a target image; the target image includes a real bounding box for locating the target part; The Swin transformer network replaces the original backbone network darknet53 of the YOLOv8 model, and a residual attention module is introduced into the Swin transformer network; based on the Swin transformer network combined with the spatial pyramid pooling layer, a new loss function Focal-EIoU is used to replace the original loss function CIoU of the YOLOv8 model to construct an improved YOLOv8 model; the residual attention module is used to focus on the features corresponding to the target part so that the Swin transformer network generates a feature map corresponding to the target part; the spatial pyramid pooling layer is used to perform maximum pooling or average pooling on the feature map to obtain a feature vector; the new loss function Focal-EIoU is used to calculate the difference between the predicted aspect ratio corresponding to the predicted bounding box of the target part and the actual aspect ratio corresponding to the actual bounding box of the target part; The target image is input into the improved YOLOv8 model to generate a prediction result; the prediction result includes scale information corresponding to the predicted bounding box of the target part and scale information corresponding to the true bounding box; wherein the scale information corresponding to the predicted bounding box includes a predicted aspect ratio and a predicted center point coordinate; the scale information corresponding to the true bounding box includes a true aspect ratio and a true center point coordinate.

2. The switch machine parts detection method based on the YOLOv8 model according to claim 1 is characterized in that: The Swin transformer network includes: a first structure, a second structure, a third structure and a fourth structure; wherein the first structure includes: a linear embedding module and a hierarchical visual converter based on a displacement window; the second structure includes an image block merging module and a hierarchical visual converter based on a displacement window; the third structure and the fourth structure are the same as the second structure.

3. The switch machine parts detection method based on the YOLOv8 model according to claim 2 is characterized in that: The residual attention modules are respectively arranged between the linear embedding module and the displacement window-based hierarchical visual converter of the first structure, between the image block merging module and the displacement window-based hierarchical visual converter of the second structure, between the image block merging module and the displacement window-based hierarchical visual converter of the third structure, and between the image block merging module and the displacement window-based hierarchical visual converter of the fourth structure.

4. The switch machine parts detection method based on the YOLOv8 model according to claim 1 is characterized in that: The improved YOLOv8 model also includes an image block segmentation module; the image block segmentation module is used to segment the switch machine part image into a plurality of image blocks; the Swin transformer network is used to generate a feature map of the target part in the switch machine part image based on the image blocks.

5. The switch machine parts detection method based on the YOLOv8 model according to claim 1 is characterized in that: The new loss function Focal-EIoU is: L (Focal-EIoU) =IoU γ L EIoU Among them, L EIoU represents the difference between the predicted aspect ratio and the true aspect ratio corresponding to the improved YOLOv8 model; h c represents the height of the minimum bounding rectangle of the predicted bounding box and the true bounding box; w c Indicates the width of the minimum bounding rectangle of the predicted bounding box and the true bounding box; L (Focal-EIoU) represents the loss of the improved YOLOv8 model; γ represents the hyperparameter that controls the degree of negative sample suppression; ρ represents the Euclidean distance between the predicted center point coordinates and the true center point coordinates of the predicted bounding box; IoU represents the set threshold; w represents the width of the predicted box; w gt represents the width of the real box; h represents the height of the predicted box; b represents the center point of the predicted box; b gt Indicates the center point of the real box; h gt Indicates the height of the true box.

6. The switch machine parts detection method based on the YOLOv8 model according to claim 1 is characterized in that: The training process of the improved YOLOv8 model specifically includes: A preset batch of target images is used as input to the improved YOLOv8 model, and scale information corresponding to the predicted bounding box of the target part in the target image and scale information corresponding to the true bounding box of the target part are used as output. The YOLOv8 model is trained until the difference between the predicted aspect ratio corresponding to the predicted bounding box and the true aspect ratio corresponding to the true bounding box is less than a first preset threshold, and the distance between the predicted center point coordinates of the predicted bounding box and the true center point coordinates of the true bounding box is less than a second preset threshold, and the training is stopped.

7. The switch machine parts detection method based on the YOLOv8 model according to claim 1 is characterized in that: Marking the target parts in the switch machine parts image specifically includes: Filter all target parts that are larger than the preset minimum size for marking; For the visible parts of various target parts, draw frames closely along the edges of the target parts according to the tangent lines; If there is an obstructed part in the target part to be marked, the obstructed part of the target part should be completed according to the shape of the target part before marking; Abnormal target parts in the switch machine part image are discarded, and each remaining target part is framed separately to generate a target frame; the abnormal target parts include blurred, dark, and overexposed target parts, and the frame edge of the target frame cannot overlap with the boundary of the switch machine part image; the target frame is the true boundary frame.

8. The switch machine parts detection method based on the YOLOv8 model according to claim 1 is characterized in that: Before marking the target parts in the switch machine part image, the following steps are also included: Collecting a data set, the data set including an initial switch machine part image; Normalization preprocessing is performed on the initial switch machine part image to obtain the switch machine part image.

9. A computer device comprising: A memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the switch machine parts detection method based on the YOLOv8 model according to any one of claims 1 to 8.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the switch machine parts detection method based on the YOLOv8 model described in any one of claims 1 to 8 is implemented.

Images