Roof key component positioning method and system based on improved YOLOv5

Abstract

The application discloses a roof key component positioning method and system based on an improved YOLOv5, wherein the method uses a Ghost convolution to replace original convolution layers in a backbone network based on original YOLOv5 algorithm technology, so that the model is more lightweight; an SE attention mechanism is introduced to improve the attention degree to important features; and CIOU_Loss is used as a loss function to make a prediction frame more consistent with a real frame.The system comprises a data set construction module, an image preprocessing module, a model lightweight and feature extraction optimization module, a feature fusion module, a network structure model construction module, a model training module and an image detection module.The application has the advantages of high detection precision, fast detection speed and low dependence on roof image quality.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of traffic safety engineering, in particular to a roof key component positioning method and system based on improved YOLOv5. BACKGROUND

[0002] With the rapid development of China's railway system, the operation safety of trains is increasingly valued by people. In the through-type locomotive detection, the identification and detection of pantograph and other roof key components as an important part of locomotive detection have been widely concerned. There are a series of possible failures and hidden dangers in the train roof components, such as pantograph wear and deformation, lightning arrester rupture and air conditioner outdoor unit peeling, which may cause train operation failure. Therefore, the roof equipment state online detection system is of great significance to the safe operation of trains.

[0003] The roof equipment state detection methods include manual detection method and image processing-based detection technology. The manual detection method is to wait for the train to enter the vehicle depot, and after a series of operations such as parking, lowering the pantograph, and disconnecting the power, the staff climbs onto the roof with the help of simple auxiliary tools to check whether there are foreign matters on the roof surface and whether the key components have failed, etc. The manual detection method has the problems of low detection efficiency, secondary foreign matters, and certain safety hazards, which cannot meet the needs of today's safe operation.

[0004] The image processing-based detection technology acquires clear and complete roof image data, locates the key components of the roof, and then processes and analyzes them to meet the needs of roof equipment state detection and foreign matter detection. Wang Yongqiang of China Railway Design Group Co., Ltd. acquires roof images through 3 linear array cameras, locates the pantograph over the slide plate and porcelain bottle through Hough, and uses ECC registration, template matching, and structural similarity evaluation methods to detect defects of key components. The method has the problems of slow running speed and poor real-time performance. Wang Yechao of Tianjin constructs a data template library of roof key components by using HTM neural network algorithm to learn the structure of the key components of the train roof, and compares the key components in real time with the key components in the data template library to determine whether the key components are abnormal. The accuracy needs to be improved. SUMMARY

[0005] The purpose of the present application is to provide a roof key component positioning method and system based on improved YOLOv5 with high accuracy and fast detection rate.

[0006] The technical solution to achieve the purpose of the present application is: a roof key component positioning method based on improved YOLOv5, comprising the following steps:

[0007] Step 1, collect train roof images, label the key components of the train roof including pantograph slide, air conditioner outdoor unit, air conditioner fixed nut, air conditioner bolt, antenna, lightning arrester and porcelain bottle in the images, perform data enhancement on the collected train roof images, and construct a data set;

[0008] Step 2, using the Head part of YOLOv5 algorithm, pre-processing the un-labeled original train roof images collected in step 1, including Mosaic data enhancement, adaptive anchor frame calculation and adaptive picture scaling;

[0009] Step 3, optimizing the Backbone main network part of YOLOv5 algorithm, replacing the original CSP1_X module in Backbone with GhostBottleneck module to make the model lightweight; adding SE attention mechanism to the end of Backbone;

[0010] Step 4, using the Neck part of YOLOv5 algorithm, performing feature fusion;

[0011] Step 5, in the prediction part of YOLOv5 algorithm, using CIoU_Loss instead of GIOU_Loss as the positioning loss function, obtaining the network structure model of improved YOLOv5;

[0012] Step 6, using the data set labeled in step 1, training data and verifying performance through the network structure model composed of steps 2-5;

[0013] Step 7, using the best.pt weight file obtained after model training in step 6 to detect the test train roof image, and outputting the positioning detection result of the key components of the train roof.

[0014] A train roof key component positioning system based on improved YOLOv5, comprising a data set construction module, an image preprocessing module, a model lightweight and feature extraction optimization module, a feature fusion module, a network structure model construction module, a model training module and an image detection module, wherein:

[0015] The data set construction module is used to collect train roof images, label the key components of the train roof including pantograph slide, air conditioner outdoor unit, air conditioner fixed nut, air conditioner bolt, antenna, lightning arrester and porcelain bottle in the images, perform data enhancement on the collected train roof images, and construct a data set;

[0016] The image preprocessing module uses the Head part of YOLOv5 algorithm to pre-process the un-labeled original train roof images collected, including Mosaic data enhancement, adaptive anchor frame calculation and adaptive picture scaling;

[0017] The model lightening and feature extraction optimization module optimizes the Backbone main network part of the YOLOv5 algorithm, replaces the original CSP1_X module in the Backbone with a GhostBottleneck module, so that the model is lightened, and an SE attention mechanism is added to the end of the Backbone.

[0018] The feature fusion module adopts the Neck part of the YOLOv5 algorithm to perform feature fusion.

[0019] The network structure model construction module adopts CIoU_Loss instead of GIOU_Loss as a positioning loss function in the prediction part of the YOLOv5 algorithm to obtain an improved network structure model of YOLOv5.

[0020] The model training module trains data and verifies performance through the network structure model by using the data set labeled by the data set construction module.

[0021] The image detection module detects the to-be-detected roof image using the best.pt weight file obtained after model training, and outputs a roof key component positioning detection result.

[0022] A mobile terminal comprises a memory, a processor and a computer program stored on the memory and executable on the processor, and the processor implements the roof key component positioning method based on the improved YOLOv5 when executing the program.

[0023] A computer readable storage medium has a computer program stored thereon, and the program is executed by a processor to implement the steps in the roof key component positioning method based on the improved YOLOv5.

[0024] Compared with the prior art, the present application has the following advantages: (1) compared with the original YOLOv5 algorithm, the original convolution layer in the backbone network is replaced with a Ghost convolution, more features are generated with fewer parameters, and the detection speed is improved; (2) compared with the original YOLOv5 algorithm, the SE attention mechanism and CIOU_Loss are introduced as the loss function, and the detection accuracy is improved. BRIEF DESCRIPTION OF DRAWINGS

[0025] Figure 1 is a network structure diagram of the YOLOv5 algorithm.

[0026] Figure 2 is a network structure diagram of the improved YOLOv5 algorithm in the method of the present application.

[0027] Figure 3 is a flowchart of the method of the present application.

[0028] Figure 4is the mAP value column chart of each key component in the example of the application.

[0029] Figure 5 is the detection result chart of part of the method of the application.

[0030] Figure 6 is the detection result chart of another part of the method of the application. DETAILED DESCRIPTION

[0031] YOLOv5 is the fifth generation of YOLO series detection algorithm, which provides higher detection accuracy while being more lightweight, and has become the best algorithm in the series in terms of performance, and has achieved excellent results in engineering applications. Accurate positioning of key components on the roof is the key to realizing roof equipment state detection based on image processing. This paper mainly identifies and detects the key components on the roof of a high-speed train. In order to meet the real-time requirements of detection, a more lightweight YOLOv5s is selected as the benchmark model, which is convenient for subsequent application in embedded devices. The GhostBottleneck module and SE attention module are introduced to improve the model, realize the identification and positioning of the roof antenna, lightning arrester, porcelain bottle and pantograph, not only shorten the identification time, but also improve the identification accuracy, and the effect is good in practicality.

[0032] The application discloses a roof key component positioning method based on an improved YOLOv5.

[0033] Step 1, collecting train roof images, labeling the roof key components including pantograph slide, air conditioner outdoor unit, air conditioner fixed nut, air conditioner bolt, antenna, lightning arrester and porcelain bottle in the images, and performing data enhancement on the collected roof images to construct a data set;

[0034] Step 2, using the Head part of the YOLOv5 algorithm to pre-process the un-labeled original roof images collected in step 1, including Mosaic data enhancement, adaptive anchor frame calculation and adaptive picture scaling, enriching the data set, improving the running speed and reducing the loss function;

[0035] Step 3, optimizing the Backbone main network part of the YOLOv5 algorithm, using the GhostBottleneck module to replace the original CSP1_X module in the Backbone, so as to make the model lightweight; adding an SE attention mechanism to the end of the Backbone to improve the feature extraction capability of the model;

[0036] Step 4, using the Neck part of the YOLOv5 algorithm to perform feature fusion;

[0037] Step 5, in the prediction part of the YOLOv5 algorithm, CIoU_Loss is used instead of GIOU_Loss as the positioning loss function to improve the convergence speed and stability during training, and the improved network structure model of YOLOv5 is obtained;

[0038] Step 6, using the data set labeled in step 1, training the data through the network structure model composed of steps 2-5 and verifying the performance;

[0039] Step 7, using the best.pt weight file obtained after model training in step 6 to detect the test roof image and output the roof key component positioning detection result.

[0040] As a specific example, step 1 specifically includes:

[0041] Step 1.1, configure a high-definition industrial camera in the train detection area, and start the camera to continuously shoot the roof image when the train passes, to obtain high-definition images of the train roof. After manual selection, select images containing pantograph slide, air conditioner outdoor unit, air conditioner fixed nut, air conditioner bolt, antenna, lightning arrester and porcelain bottle;

[0042] Step 1.2, label the selected images using LabelImg tool to label the pantograph slide, air conditioner outdoor unit, air conditioner fixed nut, air conditioner bolt, antenna, lightning arrester and porcelain bottle in the image respectively;

[0043] Step 1.3, divide the labeled roof image into training set and validation set according to the ratio of 8:2.

[0044] As a specific example, step 3 specifically includes:

[0045] Step 3.1, introduce a lightweight network GhostNet to replace the residual module of CSP1 structure in the main network;

[0046] Step 3.1.1, GhostNet first uses standard convolution to obtain the original output, denoted as:

[0047] Y=X*f+b

[0048] Where the given input data X∈R (c×h×w) , * represents convolution operation, f is the convolution filter of this layer, b is the bias term, Y is the output, c is the number of input channels, h and w are the height and width of the input data respectively;

[0049] Step 3.1.2, generate m intrinsic feature maps using one convolution:

[0050] Y′=X′*f′

[0051] Where Y′∈R(h1×w1×m) is output, X' represents input data, f' represents the convolution filter of the layer, h1 and w1 are the height and width of the input data respectively, and m is the number of intrinsic feature maps;

[0052] Step 3.1.3, the output is converted into several groups of 5*5 size convolution kernels using a series of linear transformation operations, denoted as:

[0053] y ij = φ i,j (y i ′)

[0054] where φ i,j represents a linear operation, y i ' represents the i-th intrinsic feature map output in step 3.1.2, y ij represents a convolution kernel, i represents the intrinsic feature map number, and j represents the channel number;

[0055] Step 3.1.4, two fusion operations are performed, without changing the number of output feature maps;

[0056] Step 3.2, fuse SE attention mechanism on the basis of YOLOv5 algorithm;

[0057] Step 3.2.1, a convolution operation is performed on the input to generate a feature map, and then a compression operation is performed, which globally averages the feature map and compresses each channel H*W pixel value u into a real number to generate a 1*1*C vector Z c , the formula is as follows:

[0058]

[0059] where H and W are the height and width of the input data respectively, i and j are dimension variables, c is the number of channels, u c is the pixel value, Z c is a 1*1*C vector, and F sq is the compression formula;

[0060] Step 3.2.2, the vector Z c is processed through two fully connected layers;

[0061] Step 3.2.3, there is a hyperparameter r between the two fully connected layers, and the vector Z c is changed to 1*1*C / r in dimension after the first fully connected layer and is activated nonlinearly, and is changed to 1*1*C in dimension after the second fully connected layer, with a Sigmoid activation function, and finally s(1*1*C) is multiplied with the corresponding channel of the feature map (C*H*W) to output the result.

[0062] As a specific example, the step 5 specifically comprises:

[0063] CIOU_Loss is used as the loss function, on the basis of retaining the advantages of GIOU loss, DIOU loss increases the center point distance measurement, directly optimizes the distance between the centers of two boxes, and on the basis of DIOU, a CIOU_Loss is formed by increasing an item of tilt angle data, and the calculation formula is as follows:

[0064]

[0065]

[0066] Wherein d is the distance between the center points of the predicted box and the actual box B, c is the diagonal length of the minimum box enclosing the predicted box and the actual box, v is the width-height ratio measurement parameter of the predicted box and the actual box, CIOU is C-intersection ratio, IOU is intersection ratio, w B is the width of the B point, w A is the width of the A point, h B is the height of the B point, h A is the height of the A point, B point is the center point of the actual box, and A point is the center point of the predicted box.

[0067] As a specific example, the step 6 specifically comprises:

[0068] According to the training set and the verification set constructed in step 1, training parameters such as batch size and maximum iteration number are set, the training set and the verification set constructed in step 1 are input into the improved YOLOv5 network structure model for training, and training weights are obtained.

[0069] The application also provides a roof key component positioning system based on improved YOLOv5, which comprises a data set construction module, an image preprocessing module, a model lightweight and feature extraction optimization module, a feature fusion module, a network structure model construction module, a model training module and an image detection module, wherein:

[0070] The data set construction module is used for collecting train roof images, labeling the roof key components including pantograph slide, air conditioner outdoor unit, air conditioner fixed nut, air conditioner bolt, antenna, lightning arrester and porcelain bottle in the images, performing data enhancement on the collected roof images, and constructing a data set;

[0071] The image preprocessing module adopts the Head part of the YOLOv5 algorithm to preprocess the collected unlabeled original roof images, including Mosaic data enhancement, adaptive anchor box calculation and adaptive picture scaling, enriches the data set, improves the running speed and reduces the loss function;

[0072] The model lightening and feature extraction optimization module optimizes the Backbone main network part of the YOLOv5 algorithm, replaces the original CSP1_X module in the Backbone with a GhostBottleneck module to lighten the model, and adds an SE attention mechanism to the end of the Backbone to improve the feature extraction capability of the model.

[0073] The feature fusion module uses the Neck part of the YOLOv5 algorithm to perform feature fusion.

[0074] The network structure model construction module uses CIoU_Loss instead of GIOU_Loss as the positioning loss function in the prediction part of the YOLOv5 algorithm to improve the convergence speed and stability during training, and obtains the network structure model of the improved YOLOv5.

[0075] The model training module trains data and verifies the performance through the network structure model using the data set labeled by the data set construction module.

[0076] The image detection module detects the to-be-tested roof image using the best.pt weight file obtained after the model training, and outputs the roof key component positioning detection result.

[0077] The application also provides a mobile terminal comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the program to realize the roof key component positioning method based on the improved YOLOv5.

[0078] The application also provides a computer readable storage medium having a computer program stored thereon, wherein the program is executed by a processor to realize the steps in the roof key component positioning method based on the improved YOLOv5.

[0079] The application will be further described in detail below in combination with the drawings and specific embodiments.

[0080] Embodiment

[0081] The roof key component positioning method based on the improved YOLOv5 comprises the following steps: Figure 1 The YOLOv5 improved algorithm network structure diagram of the application is shown in Figure 2 The YOLOv5 improved algorithm network structure diagram of the application is shown in Figure 3 The YOLOv5 improved algorithm network structure diagram of the application is shown in

[0082] Step 1: Collect train roof images, label pantograph slide, air conditioner, air conditioner fixed nut, air conditioner bolt, antenna, lightning arrester and porcelain bottle in the image. Data augmentation is performed on the collected train roof images to construct the dataset.

[0083] Step 1.1: Configure high-definition industrial cameras in the train detection area. When the train passes, start the camera to continuously shoot the roof image and obtain high-definition images of the train roof. After manual selection, select images containing pantograph slide, air conditioner, air conditioner fixed nut, air conditioner bolt, antenna, lightning arrester and porcelain bottle.

[0084] Step 1.2: Label the pantograph slide, air conditioner, air conditioner fixed nut, air conditioner bolt, antenna, lightning arrester and porcelain bottle in the image using LabelImg tool.

[0085] Step 1.3: Divide the labeled train roof images into training set and validation set according to the ratio of 8:2. In this experiment, a total of 18356 train roof images are collected, of which the number of training set samples is 14685 and the number of validation set samples is 3671.

[0086] Step 2: Apply the Head part of the original YOLOv5 algorithm to preprocess the image, including Mosaic data enhancement, adaptive anchor box calculation and adaptive image scaling.

[0087] Step 3: Optimize and improve the Backbone main network part of YOLOv5 for feature extraction and add SE attention mechanism.

[0088] Step 3.1: Introduce lightweight network GhostNet to replace the residual module of CSP1 structure in the main network.

[0089] Step 3.1.1: GhostNet first uses standard convolution to obtain the original output, denoted as:

[0090] Y = X * f + b

[0091] where the given input data X e R (c×h×w) , * represents convolution operation, f is the convolution filter of this layer, b is the bias term, Y is the output, where c is the number of input channels, h and w are the height and width of the input data, respectively.

[0092] Step 3.1.2: Then generate m inherent feature maps using one convolution

[0093] Y' = X' * f'

[0094] where Y' e R (h1×w1×m)is the output, X' represents the input data, f' represents the convolution filter of this layer, hi and wi are the height and width of the input data respectively, and m is the number of intrinsic feature maps.

[0095] Step 3.1.3: The output is again transformed into several groups of 5*5 size convolution kernels using a series of linear transformation operations, denoted as:

[0096] y ij = Φ i,j (y i ′)

[0097] where Φ i,j represents a linear operation, y i ' represents the output of step 3.1.2, y ij represents the convolution kernel, i represents the intrinsic feature map number, and j represents the channel number.

[0098] Step 3.1.4: Finally, two fusion operations are performed, without changing the number of feature maps in the output.

[0099] Step 3.2: Fusion of SE attention mechanism based on YOLOv5 original algorithm.

[0100] Step 3.2.1: First, a convolution operation is performed on the input to generate a feature map, followed by a compression (Squeeze) operation, which performs global average pooling (Global Avg Pool) on the feature map. The H*W pixel values u of each channel are compressed into a real number, generating a 1*1*C vector Z c , as follows:

[0101]

[0102] where H and W are the height and width of the input data respectively, i and j are dimension variables, c is the number of channels, u c is the pixel value, Z c is a 1*1*C vector, and F sq is the compression formula

[0103] Step 3.2.2: Vector Z c is then processed through two fully connected layers.

[0104] Step 3.2.3: There is a hyperparameter r between the two fully connected layers. Vector Z c is first transformed into a 1*1*C / r dimension through the first fully connected layer with nonlinear activation, and then into a 1*1*C dimension through the second fully connected layer with a Sigmoid activation function. Finally, s(1*1*C) is multiplied by the corresponding channel of the feature map (C*H*W), and the output result is obtained.

[0105] Step 4: Apply the Neck part of the original YOLOv5 algorithm to perform feature fusion;

[0106] Step 5: Apply the prediction part of the original YOLOv5 algorithm, using the CIoU_Loss loss function. While retaining the advantages of GIOU loss, DIOU loss adds a center point distance metric, directly optimizing the distance between the centers of two bounding boxes. Furthermore, based on DIOU, an additional tilt angle data is added to form CIOU_Loss, the calculation formula of which is shown below.

[0107]

[0108]

[0109] Where d is the distance between the center points of the predicted bounding box and the actual bounding box B, c is the diagonal length of the smallest bounding box enclosing the predicted bounding box and the actual bounding box, v is the aspect ratio metric for the predicted bounding box and the actual bounding box, CIOU is the C-intersection over union ratio, IOU is the intersection over union ratio, and w B Let w be the width of point B. A Let h be the width of point A. B h is the height of point B. A Point A is the height of the bounding box, point B is the center point of the actual bounding box, and point A is the center point of the predicted bounding box.

[0110] Step 6: Based on the constructed training and validation sets, set the batch size to 16, the maximum number of iterations to 500, and the input image size to [640, 640], and train the improved YOLOv5 model. The trained PR is as follows: Figure 4 The AP values ​​for pantograph slider, air conditioner outdoor unit, air conditioner mounting nuts, air conditioner bolts, antenna, surge arrester, and ceramic insulator are 94.9%, 99.1%, 99.5%, 97.9%, 94.3%, 99.5%, and 87.8%, respectively. The average accuracy mAP reached 96.2%, and the detection of key components on the roof achieved the expected results.

[0111] Step 7: Use the best.pt weight file trained by the model to detect the roof image of the vehicle under test. Partial detection results are shown below. Figures 5-6 As shown.

[0112] Compared to other algorithms, the detection results for the unified dataset are shown in the table below. The improved network outperforms the YOLOv5 network in both accuracy and speed, with a 1.4% improvement in mAP and a 7.72 FPS increase in speed. This method fully meets the requirements for locating key components on the roof of trains in the field.

[0113] Table 1 Comparison of Model Detection Results

[0114]

[0115] The above description is only the preferred embodiment of the present application, but the protection scope of the present application is not limited to this. Any person skilled in the art can easily think of changes or replacements within the technical range disclosed by the present application, which should be covered in the protection scope of the present application.

Claims

1. A method for locating key components on the roof of a vehicle based on an improved YOLOv5, characterized in that, Includes the following steps: Step 1: Collect images of the train roof, label key roof components including pantograph plates, air conditioner outdoor units, air conditioner mounting nuts, air conditioner bolts, antennas, surge arresters, and porcelain insulators, perform data augmentation on the collected roof images, and construct a dataset. Step 2: Using the Head part of the YOLOv5 algorithm, preprocess the unlabeled original roof image acquired in Step 1, including Mosaic data augmentation, adaptive anchor box calculation, and adaptive image scaling. Step 3: Optimize the backbone network of the YOLOv5 algorithm by replacing the original CSP1_X module in the backbone with the GhostBottleneck module to make the model lighter; add the SE attention mechanism to the end of the backbone. Step 4: Perform feature fusion using the Neck part of the YOLOv5 algorithm; Step 5: In the prediction part of the YOLOv5 algorithm, CIoU_Loss is used instead of GIOU_Loss as the localization loss function to obtain the improved YOLOv5 network structure model. Step 6: Using the dataset labeled in Step 1, train the data using the network structure model constructed in Steps 2 to 5 and verify its performance. Step 7: Use the best.pt weight file obtained after model training in Step 6 to detect the roof image under test, and output the location detection results of key roof components; Step 3 specifically includes: Step 3.1: Introduce the lightweight network GhostNet to replace the residual modules of the CSP1 structure in the backbone network; Step 3.1.1: GhostNet first uses standard convolution to obtain the original output, represented as: ； Given input data , * indicates convolution operation, f is the convolution filter of the standard convolutional layer, b is the bias term, Y is the output, c is the number of input channels, and h and w are the height and width of the input data, respectively; Step 3.1.2: Generate m intrinsic feature maps using a single convolution: ； in It is the output. Indicates input data, represents the convolutional filter of a convolutional layer, h1 and w1 are the height and width of the input data respectively, and m is the number of intrinsic feature maps; Step 3.1.3: Then, apply a series of linear transformations to the output to convert it into several groups of 5x5 convolutional kernels, as follows: ； in Represents linear operations. This represents the i-th intrinsic feature map output in step 3.1.

2. represents the convolution kernel, i represents the intrinsic feature map index, and j represents the channel index; Step 3.1.4: Perform two more fusion operations without changing the number of output feature maps; Step 3.2: Integrate the SE attention mechanism based on the YOLOv5 algorithm; Step 3.2.1: Perform a convolution operation on the input to generate a feature map, followed by compression. Global average pooling is then applied to the feature map, compressing the H*W pixel values ​​u of each channel into a single real number, generating a 1*1*C vector. The formula is as follows: ； Where H and W are the height and width of the input data, respectively, i and j are size variables, and c is the number of channels. For pixel values, It is a 1*1*C vector. This is a compression formula; Step 3.2.2: Pair the vector through two fully connected layers. Process it; Step 3.2.3: There is a hyperparameter r between the two fully connected layers, a vector. After passing through the first fully connected layer, the dimension becomes 1*1*C / r, with non-linear activation; after passing through the second fully connected layer, the dimension becomes 1*1*C, with the activation function being Sigmoid. Finally, s (1*1*C) is multiplied by the corresponding channel of the feature map (C*H*W) to output the result.

2. The method for locating key components on the roof based on improved YOLOv5 according to claim 1, characterized in that, Step 1 specifically includes: Step 1.1: Configure a high-definition industrial camera in the train inspection section. When the train passes, start the camera to continuously take pictures of the roof of the train to obtain high-definition images of the roof. After manual selection, select images that include pantograph sliding plate, air conditioner outdoor unit, air conditioner fixing nut, air conditioner bolt, antenna, surge arrester and porcelain insulator. Step 1.2: For the selected images, use the LabelImg tool to label the pantograph plate, air conditioner outdoor unit, air conditioner fixing nut, air conditioner bolt, antenna, surge arrester and porcelain insulator in the images respectively; Step 1.3: Divide the labeled roof images into training and validation sets in an 8:2 ratio.

3. The method for locating key components on the roof based on improved YOLOv5 according to claim 1, characterized in that, Step 5 specifically includes: CIOU_Loss is used as the loss function. While retaining the advantages of GIOU loss, DIOU loss adds a center point distance metric to directly optimize the distance between the centers of two bounding boxes. Furthermore, based on DIOU, a tilt angle data is added to form CIOU_Loss, and the calculation formula is shown below: ； ； Where d is the distance between the center points of the predicted bounding box and the actual bounding box B, c is the diagonal length of the smallest bounding box that encloses both the predicted and actual bounding boxes, v is the aspect ratio metric for the predicted and actual bounding boxes, CIOU is the C-intersection over union ratio, and IOU is the intersection over union ratio. The width of point B, Let A be the width of point A. The altitude of point B, Point A is the height of the bounding box, point B is the center point of the actual bounding box, and point A is the center point of the predicted bounding box.

4. The method for locating key components on the roof based on improved YOLOv5 according to claim 1, characterized in that, Step 6 specifically includes: Based on the training and validation sets constructed in step 1, set training parameters such as batch size and maximum number of iterations, and input the training and validation sets constructed in step 1 into the improved YOLOv5 network structure model for training to obtain training weights.

5. A positioning system for key roof components based on an improved YOLOv5, characterized in that, This system is used to implement the method for locating key components on the roof based on improved YOLOv5 as described in any one of claims 1 to 4. The system includes a dataset construction module, an image preprocessing module, a model lightweighting and feature extraction optimization module, a feature fusion module, a network structure model construction module, a model training module, and an image detection module, wherein: The dataset construction module is used to collect images of the train roof, annotate key roof components in the images including pantograph plates, air conditioner outdoor units, air conditioner fixing nuts, air conditioner bolts, antennas, surge arresters, and porcelain insulators, perform data augmentation on the collected roof images, and construct the dataset. The image preprocessing module uses the Head part of the YOLOv5 algorithm to preprocess the acquired unlabeled original roof images, including Mosaic data augmentation, adaptive anchor box calculation, and adaptive image scaling. The model lightweighting and feature extraction optimization module optimizes the backbone network of the YOLOv5 algorithm by replacing the original CSP1_X module in the backbone with the GhostBottleneck module to make the model lighter; and adds the SE attention mechanism to the end of the backbone. The feature fusion module uses the Neck part of the YOLOv5 algorithm to perform feature fusion; In the network structure model construction module, CIoU_Loss is used instead of GIOU_Loss as the localization loss function in the prediction part of the YOLOv5 algorithm to obtain an improved YOLOv5 network structure model. The model training module uses the dataset labeled by the dataset construction module to train the data and verify the performance through the network structure model. The image detection module uses the best.pt weight file obtained after model training to detect the roof image of the vehicle under test and outputs the location detection results of key roof components.

6. A mobile terminal, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that, When the processor executes the program, it implements the method for locating key components on the roof based on the improved YOLOv5 as described in any one of claims 1 to 4.

7. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the program is executed by the processor, it implements the steps in the method for locating key roof components based on the improved YOLOv5 as described in any one of claims 1 to 4.

Images