Point switch part detection method and equipment based on YOLOv8 model and medium

Abstract

The invention discloses a point switch part detection method and device based on a YOLOv8 model and a medium, and the method comprises the steps: marking a target part in a point switch part image, and generating a target image; the method comprises the following steps: replacing an original backbone network darknet53 of a YOLOv8 model with a Swin transfer network, and introducing a residual attention module into the Swin transfer network; the method comprises the following steps of: constructing an improved YOLOv8 model by combining a spatial pyramid pooling layer on the basis of a Swin transformer network and replacing an original loss function CIoU of the YOLOv8 model with a new loss function Focal-EIoU at the same time; and inputting the target image into the improved YOLOv8 model, and generating a prediction result. And accurate detection of the target part is realized based on the detection result output by the improved YOLOv8 model.

Description

Technical Field

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

[0002] Switch machines are crucial in rail transit turnout switching equipment. They are used for train reversing and switching, and play an important role in train driving safety. Failure of switch machine parts will cause train delays and seriously affect the normal operation of trains. Traditional fault repair methods mainly rely on maintenance personnel to conduct manual inspection and repair through visual inspection and experience judgment. 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 scenes and then tracking and registering them, virtual objects are 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, and 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, device 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 a 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 introduces a residual attention module in 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 real aspect ratio corresponding to the real 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 real 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 real bounding box includes a real aspect ratio and a real 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 uses the Swin Transformer network to replace the original backbone network darknet53 of the original YOLOv8 model, and introduces the residual attention module in the Swin Transformer network. The residual attention module can focus on the target part more accurately, thereby realizing the global extraction and detail extraction of the features of the target part, and can enhance the representation ability of the target part and suppress the interference of redundant information, so that the feature map generated by the Swin transformer network is more accurate than the feature map generated by the original backbone network. At the same time, the feature map is combined with the spatial pyramid pooling layer to perform maximum pooling or average pooling to obtain the feature vector, which further enhances the extraction ability of the feature map. Then, by introducing a new loss function Focal-EIoU to replace the original loss function CIoU, because 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, and finally accurately positioning the target part 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 drawings required for use in the embodiments will be briefly introduced below. 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 paying creative work.

[0014] Figure 1 A schematic 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 The overall network structure diagram of an improved YOLOv8 model of 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 display diagram of 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 The overall structure 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 of 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 groups 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] Fig. 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] Fig.10 This is a comparison chart of other parts detection results 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 the present application to clearly and completely describe the technical solutions in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, not all of the embodiments. Based on the embodiments in the present application, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of this application.

[0025] With the rapid development of deep learning technology, researchers have begun to apply it to parts detection. Pechenin et al. matched 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 part detection algorithm based on multi-scale fusion, called MDSSD. The average accuracy on three mechanical parts data sets is 6.21% higher than that of SSD. Guo Fei et al. proposed an improved FastR-CNN part detection algorithm, introduced the Inception structure into Fast R-CNN, and verified by a self-made mechanical parts data set that the detection accuracy of the algorithm is better than that of the Fast R-CNN algorithm. Liu Xueping et al. of Tsinghua University combined the designed adaptive edge error function with the improved PSO algorithm and proposed the YOLOv3-AEEF part detection algorithm, which is more robust than YOLOv3 under the interference of multi-target image parts. The method based on deep learning is more robust in the case of large differences in part size, simple texture, and illumination changes. Among them, the single-stage detection method has faster detection speed, lower equipment requirements, and is easy to deploy in industrial scenarios compared with the two-stage detection method. In view of this, the switch machine parts are detected based on the single-stage target detection method. Xing Jinchao et al. proposed an improved YOLOv5 algorithm, which improves the channel domain of the CBAM attention mechanism to improve the problem of channel information loss caused by dimensionality reduction, and adds the improved CBAM to the backbone network of YOLOv5s to enhance its feature extraction ability. However, due to the limited receptive field range in the spatial attention module, the target feature information is not fully extracted, and it is easily disturbed by complex background information, thus affecting the detection accuracy. In order to make the model have a larger receptive field and better capture global features, many experts and scholars have tried to apply Transformer to the field of computer vision. Since the encoder-decoder architecture has great advantages in feature extraction, the Vision Transformer (ViT) pioneered the application of the Transformer architecture to image classification to extract global features of the image, but its self-attention mechanism has a very large number of parameters in terms of 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 effort. 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, the introduction of convolutional neural network structure in the early stage of Vision Transformer (ViT) helps to enhance the robustness of the model to hyperparameters and the convergence speed, and has a beneficial effect on improving the accuracy of the model. In convolutional neural networks, convolution operations are good at extracting local features, but have certain limitations in capturing global feature representations. Although the cascaded self-attention module in ViT can capture long-distance feature dependencies, it still ignores the details of local features, resulting in insufficient ability to extract features for switch parts. In addition to the insufficient ability to extract features for switch parts, the YOLOv8 model also has the problem of inaccurate output of detection results for switch parts.

[0027] In order to make the above-mentioned objects, features and advantages of the present application more obvious and easy to understand, the present application is further described in detail below in conjunction with 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 boundary box for positioning the target part.

[0030] In some embodiments, marking the target parts in the switch machine parts image specifically includes: screening all target parts larger than a preset minimum size among all target parts for marking; marking the visible parts of each target part by drawing frames close to the edges of the target parts according to the tangent; if there are occluded parts in the marked target parts, the occluded parts of the target parts are completed according to the shape of the target parts before marking; discarding abnormal target parts in the switch machine parts image, and drawing frames for each of the remaining target parts 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 parts image; the target frame is the real boundary frame.

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

[0032] Acquiring a data set, the data set comprising an initial switch machine part image;

[0033] The initial switch machine part image is normalized and preprocessed to obtain the switch machine part image.

[0034] In practical applications, in order to verify the effectiveness of the algorithm and accurately detect the parts of the switch machine, a self-made 3D-printed ZD6 DC electric switch machine was used and the images of the switch machine parts were collected using a mobile phone camera as a data set. Then, 2,000 images containing 8 types of targets were obtained through flipping, rotating, cropping, and scaling operations. Finally, labelImg software was used for annotation. After the annotation was completed, the corresponding labels were made, and the target images of each type were randomly assigned in a ratio of 7:1.5:1.5 as the training set, validation set, and test set of the improved YOLOv8 model.

[0035] Step 102: Replace the original backbone network darknet53 of the YOLOv8 model with the Swin transformer network, and introduce the residual attention module (ResidualAttention Module, ResAM) 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 real aspect ratio corresponding to the real 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, based on the shift window multi-head self-attention (SW-MSA), the part information can be better divided and the key information can be retained, thereby improving the feature extraction ability of the network and solving the problem of lack of information exchange between non-overlapping windows. In view of this, according to the distribution characteristics of switch machine parts, the Swin Transformer network is used to improve darknet53 and combined with the spatial pyramid pooling layer (SPPF) to construct a new feature extraction network, that is, the combination of the Swin Transformer network and the spatial pyramid pooling layer (SPPF), to more accurately detect switch machine parts.

[0037] In some embodiments, 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.

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

[0039] In some embodiments, the Swintransformer 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.

[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 displacement window-based hierarchical visual transformer (SwinTransformer Block) generally appears in a two-stage series structure, and different self-attention calculation methods are selected according to whether the current displacement window-based hierarchical visual transformer (SwinTransformerBlock) is an odd or even number. Then, in each stage after the first stage, it is composed of an image block merging module (PatchMerging) and several displacement window-based hierarchical visual transformers (SwinTransformerBlock). The role 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 downsampling operations. Assuming that the target image of the input image block merging module (PatchMerging) is 4×4, the image block merging module (Patch Merging) will divide each 2×2 adjacent pixels into an image block patch (patch), and then splice the pixels at the same position in the image block patch to obtain four feature maps, and then embed and splice them through the layer normalization layer (LayerNorm). At this time, the length and width become 1 / 2 of the original, and the channel dimension is 4 times the original. Finally, a fully connected layer adjusts the channel dimension to 2 times the original, and the channel dimension finally becomes 2C. In other words, after the data passes through the image block merging module (PatchMerging), the length and width become 1 / 2 of the original, and the depth becomes 2 times the original.

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

[0043] Reference Figure 4 The structure of the shift window based hierarchical visual transformer (SwinTransformerBlock) is similar to that of the Transformer block. SwinTransformerBlock replaces the multi-head self-attention layer (MSA) and multi-layer perceptron block (MLP) in TransformerBlock with window multi-head self-attention (WindowMulti-Self-Attention, W-MSA) and shift window multi-head self-attention (ShiftWindowMulti-HeadSelfAttention, SW-MSA), where Figure 4 Here, Z^l represents the output features of the moving window multi-head self-attention module of the Lth block; zl represents the output features of the multi-layer perceptron module of the Lth block; Z^l+1 represents the output features of the moving window multi-head self-attention module of the L+1th block; Zl+1 represents the output features of the multi-layer perceptron module of the L+1th block; Zl-1 represents the output features of the multi-layer perceptron module of the L-1th block.

[0044] W-MSA and SW-MSA are performed alternately in consecutive SwinTransformerBlocks. For the nth layer SwinTransformerBlock, the input format is (h,w,C) for the flattened feature map, where h represents the height of the prediction box; w represents the width of the prediction box, and C represents the dimension. Here, the size of the flattened feature map is obtained by flattening and concatenating the overlapping feature maps. The computational complexity of MSA is shown in formula (1) and the computational complexity of 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] Among them, Ω(MSA) represents the computational complexity of MSA; Ω(W-MSA) represents the computational complexity of W-MSA; M represents the window size of each feature map in the input flattened feature map. The previous layer Swin Transformer Block passes the input feature map through the hierarchical feature downsampling module and flattens and concatenates it. Figure 5 The self-attention calculation mechanism of the shifted window divides the new feature map size and window, where Figure 5 The X in the figure represents the feature of the target part. Then the SwinTransformer Block window of the next layer is moved so that the non-overlapping windows of the previous layer are added to the calculation again, which greatly increases the receptive field of the model.

[0048] Among them, the residual attention module (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 robustness of the model to hyperparameters and the convergence speed, showing great potential. In convolutional neural network operations, convolution is better at extracting local features and has certain limitations in the ability to extract global features. The Transformer structure, on the contrary, ignores 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, the present 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 (ie, four structures) of the Swin Transformer network.

[0050] like Figure 6 As shown in the figure, the input features are first 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, a 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-wise 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 the consumption of computing resources and the problem of model overfitting, there will also be gradient vanishing / explosion problems, which will cause the shallow network parameters to be unable to be updated, so the residual connection structure is used to optimize the overall training process and improve network performance.

[0052] like Figure 7As 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 famous 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. Activated neurons usually inhibit 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, 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 gtrepresents the center point of the ground truth box; h gt represents the height of the ground truth box.

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

[0060]

[0061] Among them, L CIoU represents the difference between the predicted aspect ratio and the ground truth aspect ratio corresponding to the YOLOv8 model; ρ represents the Euclidean distance between the predicted center point coordinates of the predicted bounding box and the ground truth center point coordinates; c represents the diagonal length of the minimum enclosing box of the predicted bounding box and the ground truth bounding box; v represents calculating the aspect ratio consistency of the predicted bounding box and the ground truth bounding box, measured by the tan angle value; α represents a balance coefficient, and the priority of α is assigned according to the set IoU threshold. When 0 < IoU < 0.5, the coefficient of α is 0. When IoU ≥ 0.5, the larger the IoU threshold, the larger the coefficient of α.

[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 large impact on the regression loss, it will make it more difficult for high-quality regression samples to be further optimized, resulting in the problem of imbalance between positive and negative samples. Therefore, this application uses the new loss function Focal-EIoU to enable the improved YOLOv8 model to recalculate the scale information of the predicted bounding box of each target part to replace the scale information of the predicted bounding box obtained by the YOLOv8 model with the original loss function CIoU, and at the same time combines L (Focal-EIoU) to alleviate the problem of imbalance between positive and negative samples. Finally, the accurate prediction of the aspect ratio of the target part is achieved, and the accurate detection of the target part is finally realized based on the detection results output by the improved YOLOv8 model.

[0063] Step 103: Input the target image into the improved YOLOv8 model to generate a prediction result; the prediction result includes the scale information corresponding to the predicted bounding box of the target part and the scale information corresponding to the ground truth bounding box; among them, the scale information corresponding to the predicted bounding box includes the predicted aspect ratio and the predicted center point coordinates; the scale information corresponding to the ground truth bounding box includes the ground truth aspect ratio and the ground truth center point coordinates.

[0064] To quantitatively analyze the detection performance of the YOLO-SRF algorithm, ablation experiments are carried out with YOLOv8n under the YOLOv8 model as the benchmark model, as follows:

[0065] The first step: Experimental analysis of the backbone network.

[0066] To verify the effectiveness of the improved YOLOv8 model, the experiment compares YOLOv8n, YOLOv8x, and the improved SwinTransformer network based on the original YOLOv8n. The changes in recall R and mAP50 when using ResAM and not using ResAM are compared. The comparison results are shown in Table 1.

[0067] Table 1. Ablation experiment results of backbone network

[0068]

[0069]

[0070] According to Table 1, compared with the original YOLOv8n network, the recall rate R of Swin Transformer increased by 2% and the mAP50 increased by 1.2%; compared with the YOLOv8x network, the recall rate R increased by 2.4% and the mAP50 increased by 0.7%; after adding the ResAM module, compared with the original YOLOv8n network, the recall rate R increased by 1.3% and the mAP50 increased by 0.3%; compared with the YOLOv8x network, the recall rate R increased by 1.7% and the mAP50 decreased by 0.2%. This is because after adding the ResAM module to YOLOv8n, the number of network layers, model size and algorithm complexity are much smaller than YOLOv8x; after Swin Transformer and ResAM construct a new feature extraction network, the recall rate R and mAP50 values ​​are higher than YOLOv8n and YOLOv8x, which are increased by 2.4, 2.4 and 3.0, 1.9 percentage points respectively, verifying the superiority of the proposed network.

[0071] Step 2: ResAM experimental analysis.

[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 increased by 1.2% and mAP50-95 is increased 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 than other mainstream attention modules, and the accuracy is better. The results prove that the ResAM module effectively improves the performance of the model and is better than 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 increases the recall rate R by 0.3% and the 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 benchmark model, and representative images of each category in the test set were selected for qualitative analysis. Figure 8 , Fig. 9 and Fig.10 As shown, 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 8As shown, the motor pointed by the white arrow in the first group of left pictures is missing part of the body. After multiple downsampling operations of the original backbone network 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 learn feature representation more effectively. At the same time, ResAM effectively restores the detail information of the target parts, thus avoiding missed detection. 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 detection. The method of the present application is more accurate than the prediction box in the left picture (indicated by the white arrow) while detecting all parts. As shown Fig. 9 , the predicted bounding box in the prediction result obtained by YOLO-SRF is more accurate, and the scale information of the target part can be better defined, proving the effectiveness of the new loss function Focal-EIoU. Therefore, YOLO-SRF, which is the improved YOLOv8 model, can better locate the position of the target parts in the switch machine and suppress the interference of redundant information compared with the YOLOv8 model, and has a significant improvement in detection accuracy and detection capability.

[0084] Furthermore, in order to verify that the YOLO-SRF algorithm proposed in this application has better target detection capabilities than other algorithms, this application conducts comparative experiments with the current mainstream single-stage target detection algorithms SSD, YOLOv5, YOLOv8n, and YOLOv8x under YOLOv8. Using recall rate R and mAP50 as evaluation criteria, the detection results of eight categories of displacement contactors, automatic switches, reducers, spindles, indicating rods, operating rods, motors, and rack blocks in the self-made track parts dataset by different target detection methods are shown in Table 4 (unit: %), where All represents the average accuracy of the eight categories.

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

[0086]

[0087] According to Table 4, the algorithm of YOLO-SRF in this application is better than other single-stage mainstream algorithms in the duplicate checking rate R and mAP50 values ​​of eight types of parts. Compared with SSD, YOLOv5, YOLOv8n, and YOLOv8x under YOLOv8, the average duplicate checking rate R of eight types of parts increased by 3.2, 1.0, 2.9, and 3.3 respectively; the average precision mAP50 values ​​of eight types of parts increased by 6.0, 2.9, 2.7, and 2.2 respectively. Among them, the mAP50 values ​​of five types of parts, such as shift contactors, automatic openers and closes, reducers, indicator rods, and operating rods, are higher than those of other algorithms; the mAP50 value of the spindle part is higher than that of SSD and is the same as that of YOLOv5, YOLOv8n, and YOLOv8x; the mAP50 value of the motor part is higher than that of SSD and YOLOv8x and is the same as that of YOLOv5 and YOLOv8n.

[0088] This experiment shows that the YOLO-SRF algorithm has better robustness in the case of complex background, mixed targets, occlusion, etc. in the switch machine target. The experimental results show that compared with the above classic algorithms, the YOLO-SRF algorithm significantly improves the detection accuracy of target parts in the switch machine image, and can locate the target parts more accurately, with the best overall target detection capability.

[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 of 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 real bounding box of the target part are used as output, and the YOLOv8 model is trained until the difference between the predicted aspect ratio corresponding to the predicted bounding box and the real aspect ratio corresponding to the real 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 real center point coordinates of the real bounding box is less than a second preset threshold, then the training is stopped.

[0091] This application proposes a switch machine parts detection method based on the YOLOv8 model. According to 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, which is combined with the improved feature extraction network. 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, wherein a computer program is stored in the memory, and the processor implements the steps in the above-mentioned method embodiments when executing the computer program.

[0093] In an exemplary embodiment, a computer-readable storage medium is provided, storing a computer program, and when the computer program is executed by a processor, the steps in the above 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 can understand that all or part of the processes in the above-mentioned embodiments can be completed by instructing the relevant hardware through a computer program, and 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. Among them, any reference to the memory, database or other medium used in the embodiments provided in the present application can include at least one of non-volatile and volatile memory. Non-volatile memory can include read-only memory (ROM), magnetic tape, floppy disk, flash memory, optical memory, high-density embedded non-volatile memory, resistive memory (ReRdM), magnetic memory (MRdM), ferroelectric memory (FRdM), phase change memory (PCM), graphene memory, etc. Volatile memory can 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 database involved in each embodiment provided in this application may include at least one of a relational database and a non-relational database. The non-relational database may include a distributed database based on blockchain, etc., but is not limited thereto. The processor involved in each embodiment provided in this application may be a general-purpose processor, a central processing unit, a graphics processor, a digital signal processor, a programmable logic device, a data processing logic device based on quantum computing, etc., but is not limited thereto.

[0097] The technical features of the above embodiments may be combined arbitrarily. 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 article uses specific examples to illustrate the principles and implementation methods of this application. The description of the above embodiments is only used to help understand the method and core ideas of this application. At the same time, for those skilled in the art, according to the ideas of this application, there will 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 introduces a residual attention module in 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 real aspect ratio corresponding to the real 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 real 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 real bounding box includes a real aspect ratio and a real 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 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.

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 real 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 of 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 real bounding box of the target part are used as output, and the YOLOv8 model is trained until the difference between the predicted aspect ratio corresponding to the predicted bounding box and the real aspect ratio corresponding to the real 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 real center point coordinates of the real bounding box is less than a second preset threshold, then 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 to 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 shall be completed according to the shape of the target part before marking; Abnormal target parts in the switch machine parts 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 parts image; the target frame is the real 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 is also included: Acquiring a data set, the data set comprising an initial switch machine part image; The initial switch machine part image is normalized and preprocessed 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 as described in 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 method for detecting switch machine parts based on the YOLOv8 model described in any one of claims 1 to 8 is implemented.

Images