Track foreign matter detection method based on YOLO algorithm

Abstract

The invention belongs to the field of track foreign matter detection, and discloses a track foreign matter detection method based on a YOLO algorithm. Comprising the following steps: collecting track foreign matter images to form a data set, and performing sample expansion processing and label processing on each image in the data set so as to obtain a sample set; a channel attention mechanism and a direction perception mechanism are added into the YOLOv5, constraint conditions are set, and then an improved YOLOv5 algorithm is obtained and marked as a G-YOLOv5 model; a G-YOLOv5 model is pre-trained by using the sample set, and a trained G-YOLOv5 model is obtained; a real-time image of the track is collected and processed, the processed image is imported into the trained G-YOLOv5 model, and a real-time detection report is obtained; the obtained real-time detection report is arranged on a visual interface, and efficient track foreign matter detection based on the YOLO algorithm is achieved.

Description

Technical Field

[0001] The present invention relates to the field of track foreign object detection, and more specifically, to a track foreign object detection method based on the YOLO algorithm. Background Art

[0002] Combined with the current status of rail transit safety research, the YOLOv5 algorithm is a conventional technology for foreign object detection. However, in the use of the ordinary YOLOv5 algorithm, due to the randomness of the location of foreign objects intruding on the track, the scale of the foreign object target varies greatly, resulting in a decrease in the detection effect. In addition, there are multiple foreign objects, which leads to inaccurate recognition and misidentification during the recognition process. In complex backgrounds, the target recognition ability of the YOLOv5 algorithm is limited. In view of the problem of many interferences and low recognition of foreign objects in track detection under complex backgrounds, how to capture and focus on the precise location and reduce background information interference is the main problem faced.

[0003] In view of this, the present invention proposes a track foreign object detection method based on the YOLO algorithm to solve the above problem. Summary of the invention

[0004] In order to overcome the above defects of the prior art and to achieve the above objectives, the present invention provides the following technical solutions: A method for detecting foreign objects on a track based on the YOLO algorithm, comprising: Step S1: collecting images of foreign objects on a track to form a data set, performing sample expansion processing and label processing on each image in the data set, and then obtaining a sample set; Step S2: Add channel attention mechanism and direction perception mechanism to YOLOv5, and set constraints to obtain the improved YOLOv5 algorithm, which is recorded as G-YOLOv5 model; Step S3: Use the sample set to pre-train the G-YOLOv5 model to obtain the trained G-YOLOv5 model; Step S4: collect and process the real-time image of the track, import the processed image into the trained G-YOLOv5 model, and obtain a real-time detection report; Step S5: arranging the real-time detection report on the visual interface.

[0005] Preferably, the method for collecting images of track foreign matter to form a data set comprises: High-definition cameras are installed around the track to capture images at different time periods. Dynamic images are collected by extracting video frames, and static images are collected by acquiring single images. Images including foreign objects in static images and dynamic images are extracted as foreign object images, and the data set includes all types of foreign object images.

[0006] Preferably, the method of performing sample expansion processing and label processing on each image in the data set to obtain a sample set includes: The target box is obtained through the YOLOv5 algorithm, and the target box is defined as , randomly select crop ratio , according to the clipping formula , calculate the cropped image ; Among them, the cropped image The target foreign body is included. is the coordinate of the upper left corner of the target box, and represent the horizontal and vertical coordinates respectively, is the width of the target box, is the height of the target box, crop ratio , Indicates the image corresponding to the target box; After cropping the image Perform brightness adjustment and weather interference respectively to obtain the image after preliminary processing , use CutMix enhancement technology to combine any two images Mixing to generate new images , using Mosaic enhancement technology to transform any image larger than two , and Stitch them together to generate multiple target images , then the image , , and are all regarded as expanded sample images; For each sample image, use the pre-trained Faster R-CNN model to obtain a pre-annotated image, use the annotation tool to adjust the boundary and correct the category label of the target box of the pre-annotated image, and collect all the images after boundary adjustment and category label correction to form a sample set; The format of setting annotation is ,and It is the foreign body category label.

[0007] Preferably, the cropped image Perform brightness adjustment and weather interference respectively to obtain the image after preliminary processing , methods include: Use the improved formula to The brightness is adjusted, and the adjusted image is recorded as , improved formula , ; in, represents the brightness adjustment factor, and , represents the brightness offset, and , represents the contrast adjustment factor, and , represents the contrast offset, and , Representing images The contrast ratio, Represents the adjusted image The contrast value of Use multivariate formulas to group images The environment is adjusted to obtain images after rainy day simulation, foggy day simulation and snowy day simulation respectively. The adjusted images are also recorded as ; Among them, the multivariate formula group , Indicated in the image The function that simulates raindrops on Indicates the intensity of controlling raindrops. Represents the parameters that control the fog density, Represents the image Apply Gaussian blur, Represents a randomly generated snowflake effect. Indicates that the randomly generated snowflake effect follows the Bernoulli distribution. Represents the parameter that controls the density of snowflakes.

[0008] Preferably, the channel attention mechanism and the direction perception mechanism are added to YOLOv5, and constraints are set to obtain an improved YOLOv5 algorithm, which is recorded as a G-YOLOv5 model. The specific method includes: Set the input of the G-YOLOv5 model as the input feature map and transfer it to the backbone network; The channel attention mechanism is introduced into the coordinate attention module to generate adaptive channel weighted feature maps, and then the alienated feature maps at different resolutions are extracted through multi-scale feature fusion; Decode the alienated feature map, remove redundant boxes through non-maximum suppression, retain the box with the highest confidence, and output multi-scale bounding boxes and foreign object recognition results; The constraint condition is set to minimize the loss function. The constraint condition formula is expressed as: ; , and is the first loss function , the second loss function And the third loss function The adaptive weight of , , , then the weight adaptive adjustment functions are , , ; in, represents the true category label of the i-th foreign object in the image, represents the predicted category label of the i-th foreign body in the image, Represents the intersection area of ​​the predicted bounding box output by the model and the true bounding box, Represents the union area of ​​the predicted bounding box output by the model and the true bounding box, Represents the area of ​​the smallest enclosing rectangle between the predicted bounding box output by the model and the true bounding box, represents the square of the length of the diagonal in the minimum enclosing rectangle, represents the true confidence, is the prediction confidence of the model output; The optimization strategy is to use the Adam optimizer for training, set the initial learning rate to 0.001, and use cosine annealing for learning rate scheduling.

[0009] Preferably, the method of introducing a channel attention mechanism in a coordinate attention module to generate an adaptive channel weighted feature map includes: Step A1: Assume that the input feature map contains td channels, and the spatial dimension of each channel is , calculate each position in the input feature map The average value and maximum value of all channels are recorded as the single channel mean and the single channel maximum value respectively. The single channel mean and the single channel maximum value are spliced ​​in the spatial dimension of the channel to obtain the fused spatial feature map; Use a convolution layer with a convolution kernel size of 7×7 to process the spatial feature map, and generate spatial attention weights through the Sigmoid activation function. Apply the generated spatial attention weights to the input feature map, weight the features of each spatial position, and obtain a weighted feature map, which is recorded as the remade feature map; Perform global average pooling on the spatial dimension of the remade feature map, calculate the global average value of each channel, collect the global average values ​​of td channels to form a td-dimensional vector, recorded as global channel information; Step A2: Use a linear layer network with global channel information as input to predict the scaling factor , based on the predicted scaling factor and the preset original reduction rate, calculate the dynamic reduction rate ; in, represents the Sigmoid activation function, represents the weight matrix of the linear layer, Represents global channel information, represents the bias term, represents the original reduction rate; Step A3: The global context information is input into a two-layer MLP. The first layer of the MLP uses a weight matrix for dimensionality reduction and uses a ReLU activation function. The second layer of the MLP uses a weight matrix for dimensionality increase and outputs it. Then, the channel attention weight is generated through a Sigmoid activation function. The generated channel attention weight is element-wise multiplied with the input feature map to obtain a weighted channel feature map. Among them, the weight matrix used in the first layer is different from that used in the second layer.

[0010] Preferably, the method of extracting alienated feature maps at different resolutions by fusing multi-scale features includes: Step B1: Extract shallow features and deep features from the backbone network, add an enhancement path to the shallow features of the backbone network, and use the SE module to align the channels of the shallow and deep features; Step B2: Use a dynamic weight generator to adaptively adjust the fusion ratio of shallow and deep features, perform weighted fusion of shallow and deep features after channel alignment according to the generated weight, and add a residual term; Step B3: Use branches with different receptive fields to extract fused features, use the Transformer module or self-attention mechanism to capture long-distance dependency information, use the bidirectional feature pyramid network to fuse multi-scale features at different levels, and generate the final multi-scale feature map, which is recorded as the alienated feature map.

[0011] Preferably, the method of pre-training the G-YOLOv5 model using the sample set to obtain the trained G-YOLOv5 model includes: The input of the G-YOLOv5 model is an image, and the output is a detection report, which includes a bounding box and foreign object recognition results. The foreign object recognition results indicate the type and number of foreign objects recognized. Input each sample image in the sample set into the G-YOLOv5 model in turn for training, and perform enhancement processing using constraints and optimization strategies until all sample sets are input, and then obtain the trained G-YOLOv5 model.

[0012] Preferably, the real-time image of the track is collected and processed, and the processed image is imported into a trained G-YOLOv5 model to obtain a real-time detection report, and the method includes: Use high-definition cameras to collect real-time images of the track, including dynamic images and static images; Import the processed real-time image into the trained G-YOLOv5 model to obtain a real-time detection report, which is used to output the category and quantity of foreign objects in the real-time image; If the number of foreign objects is 0, it means that there are no foreign objects; otherwise, there are foreign objects.

[0013] Preferably, the visualization interface includes a terminal page of a computer or a display interface of a mobile device, and performs interface interaction with the user.

[0014] The technical effects and advantages of the track foreign body detection method based on the YOLO algorithm of the present invention are as follows: 1. Collection and expansion of data sets; By capturing images at different time periods, we ensured that the dataset contained images under different lighting conditions. We collected dynamic images by extracting video frames and static images by acquiring single images to ensure the diversity of the dataset.

[0015] The target frame is obtained through the YOLOv5 algorithm, the cropping ratio is randomly selected, and the cropped image is calculated according to the cropping formula. The brightness of the image is adjusted using the improved formula to increase the diversity of the data set. The environment of the image is adjusted using a multivariate formula group to simulate the rainy, foggy and snowy environments to improve the model's adaptability under different weather conditions. New images are generated using CutMix and Mosaic enhancement techniques to increase the size and diversity of the data set. Pre-annotated images are obtained using the pre-trained FasterR-CNN model, and the target frame of the pre-annotated image is adjusted and the category label is corrected using the annotation tool to ensure the accuracy of the annotation.

[0016] 2. Improved YOLOv5 model (G-YOLOv5); The channel attention mechanism is introduced into the coordinate attention module to generate an adaptive channel weighted feature map. Through multi-scale feature fusion, the alienated feature maps at different resolutions are extracted. The alienated feature maps are decoded, and redundant boxes are removed through non-maximum suppression, and the boxes with the highest confidence are retained. Multi-scale bounding boxes and foreign object recognition results are output.

[0017] Extract shallow and deep features from the backbone network to ensure feature diversity. Use the SE module to align the channels of shallow and deep features to ensure feature consistency. Use a dynamic weight generator to adaptively adjust the fusion ratio of shallow and deep features to ensure the rationality of feature fusion. Use branches with different receptive fields to extract fused features, use the Transformer module or self-attention mechanism to capture long-distance dependency information, and use a bidirectional feature pyramid network to further fuse multi-scale features at different levels to generate the final multi-scale feature map.

[0018] The constraint condition is set to minimize the loss function to ensure the accuracy of the model. The adaptive weights of the first loss function, the second loss function, and the third loss function are set to ensure the rationality of the loss function.

[0019] 3. Real-time detection and visualization; Place the test report on the computer's terminal page or the display interface of the mobile device to ensure that the user can intuitively view the test results. Interact with the user interface to ensure that the user can monitor and manage conveniently.

[0020] This track foreign body detection method based on the YOLO algorithm achieves efficient detection and identification of track foreign bodies through the diversity and expansion of data sets, improved YOLOv5 model, real-time detection and visualization and other technical means. This method has the advantages of high precision, high robustness and real-time performance, and can effectively improve the safety and reliability of track transportation. BRIEF DESCRIPTION OF THE DRAWINGS

[0021] Figure 1 It is a structural schematic diagram of a track foreign body detection method based on the YOLO algorithm of the present invention; Figure 2 The figure is a schematic diagram of the steps of a track foreign object detection method based on the YOLO algorithm of the present invention. DETAILED DESCRIPTION

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

[0023] Example 1, please refer to Figure 1 and Figure 2 As shown, the track foreign body detection method based on the YOLO algorithm described in this embodiment includes: Main research contents: Combined with the current status of rail transit safety research, this invention mainly studies the problem of foreign object intrusion on railway operation lines, and completes the application of YOLOv5 algorithm in the detection of foreign objects invading the track. Specifically, this project first preprocesses the visual image, and then delineates the obstacle area for the track image with too complex background. On the above basis, it plans to carry out research on track foreign object detection based on attention mechanism and track foreign object detection based on multi-scale feature fusion network. The two parts of the research content are interrelated and mutually reinforcing, solving the problem of intelligent detection of foreign object target objects from a technical level, achieving real-time monitoring and real-time detection, realizing the role of timely discovery and early response, ensuring the safe operation of trains, and preventing railway track transit accidents. The specific discussion of each part is as follows: (1) Research on track foreign object detection based on attention mechanism; Specifically, it is planned to add a coordinate attention module to the YOLOv5 feature extraction network, which mainly includes two steps: coordinate information embedding and coordinate attention generation. It can model the spatial information in the image into the channel, so that the network can effectively improve the model's attention to important feature areas while only increasing a small amount of parameters and calculations, thereby enhancing the algorithm's ability to pay attention to foreign objects invading the track.

[0024] (2) Research on track foreign body detection based on multi-scale feature fusion network; Due to the randomness of the location of foreign objects intruding into the track, the scale of foreign objects varies greatly, which leads to a decrease in the detection effect. Based on the above problems, it is planned to carry out research on foreign object detection in the track based on a multi-scale feature fusion network. Specifically, for the YOLOv5 network structure, a more effective multi-scale feature fusion network is adopted. By adding two different information fusion strategies between the main shallow network and the deep network, the model's target recognition ability in complex backgrounds is improved, thereby enhancing the algorithm's ability to detect foreign objects intruding at different locations on the track.

[0025] Technical indicators: (1) Research on track foreign object detection based on attention mechanism; It is planned to add a coordinate attention module to the YOLOv5 feature extraction network, which mainly includes two steps: coordinate information embedding and coordinate attention generation, which can effectively improve the model's attention to important feature areas, thereby enhancing the algorithm's ability to pay attention to track intrusions.

[0026] (2) Research on track foreign body detection based on multi-scale feature fusion network; It is planned to adopt a more effective multi-scale feature fusion network to improve the model's target recognition ability in complex backgrounds, so as to enhance the algorithm's ability to detect intruding foreign objects at different positions on the track.

[0027] Innovation: (1) Aiming at the YOLOv5 network structure, a more effective multi-scale feature fusion network EHPANet is adopted. By adding two different information fusion strategies between the shallow backbone network and the deep network, the model's target recognition ability in complex backgrounds is improved.

[0028] (2) In order to solve the problem of high interference and low recognition of track foreign object detection under complex background, it is proposed to add a coordinate attention module at the end of the YOLOv5 backbone network to improve the model's ability to capture precise position signals and focus on important features during feature extraction, thereby reducing background information interference.

[0029] In order to improve the model's ability to detect foreign objects on the track, the visual image data needs to be preprocessed.

[0030] Step S1: Collect images of foreign objects on the track to form a data set, perform sample expansion processing and label processing on each image in the data set, and then obtain a sample set; Image preprocessing module: collects images of foreign objects on the track to form a data set, and obtains a sample set through sample expansion processing and label processing; The methods for collecting orbital foreign object image formation datasets include: Install high-definition cameras around the track. It is best to choose industrial-grade cameras that can adapt to various weather conditions. These cameras can be fixed on both sides of the track, above the track or on the train. It is also necessary to ensure that the camera's field of view covers the track and its surrounding areas, and the camera's resolution should be high enough to capture the details of foreign objects. In order to ensure the diversity of data, images should be taken at different time periods, including images under weather conditions such as day, night, sunny days, rainy days, and haze. It also includes collecting images in different track environments (such as cities, suburbs, and mountainous areas). Dynamic images are collected by extracting video frames. The cameras installed during the train's travel collect videos, extract frames as images, and form dynamic image collection. Static images are collected by collecting single images. Static images refer to images of foreign objects on the track taken manually. Images including foreign objects in static images and dynamic images are extracted as foreign object images, and all types of foreign object images are included in the data set. The purpose of requiring all types to be included is to record the diversity of foreign objects. Foreign objects may be stones, branches, garbage, tools, animals, etc., and as many types of foreign objects as possible should be covered.

[0031] In addition, foreign object images can also be obtained by using public datasets or cooperation, by looking for public datasets related to track scenes, such as datasets for railway monitoring or transportation industries. Or cooperate with railway companies or relevant departments to obtain their surveillance videos or image data. These departments usually have track monitoring systems and may record real foreign object images.

[0032] The method of performing sample expansion and label processing on each image in the data set to obtain a sample set includes: The target box is obtained through the YOLOv5 algorithm, and the target box is defined as , randomly select crop ratio , according to the clipping formula , calculate the cropped image ; When randomly cropping, the integrity of the target foreign body is preserved to avoid loss of target information due to excessive cropping.

[0033] The specific process of obtaining the target frame through the YOLOv5 algorithm is as follows: Step F1, grid division: The YOLOv5 algorithm divides the input image into grid cells, each of which is responsible for predicting the target in the area; Step F2, bounding box prediction: Each grid cell predicts two bounding boxes, each bounding box is composed of the center point coordinates ,width ,high and confidence; Step F3, category prediction: Each bounding box also predicts a category probability distribution, indicating the possibility that the object in the box belongs to each category; Step F4, Anchor Boxes: To improve the accuracy of detection, YOLOv5 uses anchor boxes to help predict objects of different shapes and sizes; Step F5, loss function: YOLOv5 uses a specific loss function (such as mean square error, etc.) to train the network so that it can more accurately predict the location and size of the bounding box.

[0034] Among them, the cropped image The target foreign body is included. is the coordinate of the upper left corner of the target box, and represent the horizontal and vertical coordinates respectively, is the width of the target box, is the height of the target box, crop ratio , Indicates the image corresponding to the target box; The meaning of the target box is to ,use and Crops the width and height.

[0035] After cropping the image Perform brightness adjustment and weather interference respectively to obtain the image after preliminary processing , use CutMix enhancement technology to combine any two images Mixing to generate new images , using Mosaic enhancement technology to transform any image larger than two , and Stitch them together to generate multiple target images , then the image , , and are all regarded as expanded sample images; For each sample image, use the pre-trained Faster R-CNN model to obtain a pre-annotated image, use the annotation tool to adjust the boundary and correct the category label of the target box of the pre-annotated image, and collect all the images after boundary adjustment and category label correction to form a sample set; The format of setting annotation is ,and is the foreign object (target) category label.

[0036] Label processing mainly includes the pre-labeling of the Faster R-CNN model and the final manual correction. The steps to obtain pre-annotated images using the pre-trained Faster R-CNN model include: Load the Faster R-CNN pre-trained model, load the model code, and switch the model to inference mode; Read the image, convert it to RGB format, then convert it to Tensor format and increase the batch dimension; Use the pre-trained model to perform inference on the input image and generate initial bounding boxes, confidences, and categories.

[0037] The method of using the annotation tool for manual annotation is specifically to import the generated pre-annotated image into the annotation tool (such as LabelImg or CVAT) for manual correction, including automatically loading the initial boundary of the original target box, adjusting the boundary position and size, and modifying the category label (if necessary).

[0038] After manual correction, save the results as standardized annotation files (COCO or YOLO format).

[0039] After cropping the image Perform brightness adjustment and weather interference respectively to obtain the image after preliminary processing , methods include: Add contrast and saturation adjustments to simulate the track environment under different lighting conditions. Use the improved formula to adjust the image The brightness is adjusted, and the adjusted image is recorded as , improved formula , ; in, represents the brightness adjustment factor, and , used to control the brightness scaling. represents the brightness offset, and , used to increase or decrease the brightness. represents the contrast adjustment factor, and , used to control the contrast scaling. represents the contrast offset, and , used to increase or decrease the contrast. Representing images The contrast ratio, Represents the adjusted image The contrast value of Use multivariate formulas to group images The environment is adjusted to obtain images after rainy day simulation, foggy day simulation and snowy day simulation respectively. The adjusted images are also recorded as ; Among them, the multivariate formula group , Indicated in the image The function of simulating raindrops on the image can be implemented in many ways. For example, adding random noise to the image to simulate the effect of raindrops, or creating a linear optical flow effect to simulate raindrops, or combining random noise and linear optical flow to create a more complex raindrop effect. In these methods, Can be a custom function that generates images that simulate raindrop tracks. These tracks can be drawn by drawing white lines to simulate raindrops, and then applying Gaussian blur to simulate the blurring effect of raindrops. Indicates the intensity of controlling raindrops. Represents the parameter that controls the fog density, and its range is usually between [0,1]. Represents the image Apply Gaussian blur to simulate the effect of fog. Represents a randomly generated snowflake effect. Indicates that the randomly generated snowflake effect follows the Bernoulli distribution. Represents the parameter that controls the density of snowflakes.

[0040] Here are one way to do it: for Controlled rainy day simulation: Generates random noise to simulate the shape and distribution of raindrops.

[0041] Apply linear optical flow to simulate the direction and speed of falling raindrops.

[0042] Overlay the generated raindrop effect onto the original image.

[0043] for Controlled fog simulation: Applies a Gaussian blur to the original image, simulating the blurring effect of fog.

[0044] According to the fog concentration Adjust the blur level.

[0045] for Controlled snow simulation: Generate random white pixels to simulate snowflakes.

[0046] According to the density of snowflakes Controls the amount of snowflakes.

[0047] Overlays a snow effect onto the original image.

[0048] These methods can help the model learn the characteristics of images under different weather conditions and improve the performance and reliability of the model in practical applications. By simulating different weather disturbances, the model's adaptability to extreme weather conditions can be enhanced, allowing it to work stably in various environments.

[0049] YOLOv5 is an efficient object detection network, but its performance may be limited by the following issues: Insufficient global feature expression ability: CSPNet and PANet in the original YOLOv5 rely more on convolution operations in the feature extraction and fusion stages, lacking the ability to model global context. The importance of spatial and channel dimensions is not fully utilized.

[0050] Insufficient two-dimensional directional perception: Targets usually have directionality (such as horizontal, vertical or diagonal distribution). How to effectively capture the feature relationships in different directions is a key to improving detection performance.

[0051] The channel attention mechanism and direction perception mechanism are introduced into the backbone network and feature fusion module of YOLOv5 to enhance the network's sensitivity and expression ability to target features.

[0052] Step S2: Add channel attention mechanism and direction perception mechanism to YOLOv5, and set constraints to obtain the improved YOLOv5 algorithm, which is recorded as G-YOLOv5 model; The channel attention mechanism and direction perception mechanism are added to YOLOv5, and constraints are set to obtain the improved YOLOv5 algorithm, which is recorded as the G-YOLOv5 model. The specific methods include: Set the input of the G-YOLOv5 model as the input feature map and transfer it to the backbone network; The channel attention mechanism is introduced into the coordinate attention module to generate adaptive channel weighted feature maps, and then the alienated feature maps at different resolutions are extracted through multi-scale feature fusion; Decode the alienated feature map, remove redundant boxes through non-maximum suppression (NMS), retain the box with the highest confidence, and output multi-scale bounding boxes and foreign object recognition results; Decoding the alienated feature map is a key step in converting the feature map output by the model into actual bounding box coordinates and classification results. It is an existing technology and can be used directly.

[0053] The constraint condition is set to minimize the loss function. The constraint condition formula is expressed as: ; , and is the first loss function , the second loss function And the third loss function The adaptive weight of , , , then the weight adaptive adjustment functions are , , ; in, represents the true category label of the i-th foreign object in the image, represents the predicted category label of the i-th foreign body in the image, Represents the intersection area of ​​the predicted bounding box output by the model and the true bounding box, Represents the union area of ​​the predicted bounding box output by the model and the true bounding box, Represents the area of ​​the smallest enclosing rectangle between the predicted bounding box output by the model and the true bounding box, represents the square of the length of the diagonal in the minimum enclosing rectangle, represents the true confidence, is the prediction confidence of the model output; The optimization strategy is to use the Adam optimizer for training, set the initial learning rate to 0.001, and use cosine annealing for learning rate scheduling.

[0054] In addition, confidence is usually a value between 0 and 1, and it is necessary to ensure that the output layer of the model can generate values ​​within this range. And confidence is a probability value that indicates the model's certainty about the detection result, but does not directly indicate the accuracy of the detection result.

[0055] In the above process, the first loss function The second loss function is used to measure the difference between the predicted foreign object category and the true foreign object category. The third loss function is used to measure the difference between the predicted bounding box and the true bounding box. It is used to measure the difference between the predicted confidence and the true confidence. The larger the difference, the worse the training effect in this aspect. Then it will be readjusted to make this aspect occupy a larger weight, strengthen the approximation effect in the training process, and improve the overall constraint effect of the constraint conditions by focusing on the poor aspects.

[0056] This paper proposes an improved coordinate attention module (CA) to solve the problems of the original YOLOv5 CA module in complex scenes, such as the limitation of feature expression, the complexity of attention generation and the simplicity of feature remapping. The improved module effectively improves the model's detection ability for complex backgrounds and multi-scale targets through strategies such as multi-scale local information fusion, dynamic weight adaptation mechanism and nonlinear feature mapping enhancement, especially showing significant advantages in the task of track foreign object detection.

[0057] The channel attention mechanism is introduced into the backbone network, and the method of generating adaptively weighted channel feature maps includes: In the backbone network of YOLOv5, the CSP module is the key structure for extracting multi-scale features. In order to enhance the flexibility of feature extraction, an improved channel attention mechanism is introduced in each Bottleneck of CSP.

[0058] Step A1: global channel information extraction; Assume that the input feature map contains td channels, and the spatial dimension of each channel is , in order to extract global channel information, calculate each position in the input feature map The average and maximum values ​​of all channels are recorded as the single channel mean and single channel maximum value, respectively, to obtain a single channel feature map, highlighting the most significant features, and concatenate the single channel mean and single channel maximum value in the spatial dimension of the channel to obtain a fused spatial feature map; Global average pooling can capture the average activation intensity of all channels at each spatial position, reflecting the overall activation level of the position. Global maximum pooling can highlight the most significant feature activation at each spatial position. The concatenation operation can fuse these two different spatial information and provide richer spatial context.

[0059] Use a convolution layer with a convolution kernel size of 7×7 to process the spatial feature map, and generate spatial attention weights through the Sigmoid activation function. Apply the generated spatial attention weights to the input feature map, weight the features of each spatial position, and obtain a weighted feature map, which is recorded as the remade feature map; The role of the convolutional layer is to learn important areas in space. Areas with higher activation values ​​represent spatial locations that should be paid more attention to. The Sigmoid function ensures that the weight values ​​are between 0 and 1, which can be interpreted as the importance of each spatial location.

[0060] Perform global average pooling (GAP) on the spatial dimension of the reshaped feature map to calculate the global average value of each channel , collect the global average values ​​of td channels to form a td-dimensional vector, recorded as global channel information; in, Indicates channel The global average value of , the role of this step is to compress the spatial information of the original input feature map into the global statistical value of each channel, generating a td-dimensional vector. The spatial attention mechanism enables the network to focus on more important areas in the image and suppress the features of unimportant areas. This vector represents the global context information of each channel and captures the overall characteristics of each channel.

[0061] Step A2: Dynamic adjustment of reduction rate; Different feature maps may have different complexity and redundancy. For feature maps with less information or higher redundancy, more aggressive dimensionality reduction can be used; while for feature maps with rich information or more independent features, less dimensionality reduction should be used to retain more information. The goal of the dynamic reduction rate is to enable the network to learn and apply this adaptive dimensionality reduction strategy.

[0062] Use a linear layer network with global channel information as input to predict the scaling factor , based on the predicted scaling factor and the preset original reduction rate, calculate the dynamic reduction rate ; in, Represents the Sigmoid activation function, ensuring that the scaling factor is in the range (0,1), which can be interpreted as the adjustment ratio of the original reduction rate. represents the weight matrix of the linear layer, which is used to map the global context information to the scaling factor, Represents global channel information, Represents the bias term, which is used to adjust the value of the scaling factor. The function is used to ensure that the reduction ratio is an integer. represents the original reduction rate, To select 1 and The maximum value among ensures that the reduction rate does not at least lead to an increase in the number of channels.

[0063] Multiple scaling factors can be predicted as needed, for example, different scaling factors can be predicted for the first layer and the second layer of the MLP, or a single scaling factor can be used.

[0064] Step A3: MLP mapping and feature weighting; Use a shared two-layer ML to perform nonlinear mapping on the global channel information according to the dynamic reduction rate to generate channel attention weights; The global context information is input into a two-layer MLP. The first layer of the MLP uses a weight matrix for dimensionality reduction and uses a ReLU activation function. The second layer of the MLP uses a weight matrix for dimensionality increase and outputs it. Then, the channel attention weight is generated through a Sigmoid activation function. The generated channel attention weight is multiplied element-by-element with the input feature map to obtain a weighted channel feature map. The weight matrix used in the first layer is different from that in the second layer, and the structural parameters of the MLP (mainly the output dimension of the first layer and the input dimension of the second layer) are adjusted according to the dynamically calculated reduction rate.

[0065] Dynamic reduction rate can make the channel attention mechanism more intelligent and efficient. By introducing a prediction module, the network can adaptively adjust the degree of dimensionality reduction according to the input features, thereby better balancing information retention and computational efficiency. Using more appropriate dimensionality reduction rates for different feature maps may improve model performance. It reduces manual adjustments to fixed reduction rates and reduces the cost of hyperparameter search.

[0066] By fusing multi-scale features, we can extract alienation feature maps at different resolutions. The specific methods include: Step B1: Extract shallow features and deep features from the backbone network, add an enhancement path to the shallow features of the backbone network to process small objects, and use the SE module to align the channels of the shallow and deep features; Step B2: Use a dynamic weight generator to adaptively adjust the fusion ratio of shallow and deep features, perform weighted fusion of shallow and deep features after channel alignment according to the generated weight, and add a residual term; the residual term is used to capture the nonlinear relationship between shallow and deep features; Step B3: Use branches with different receptive fields to extract fused features, use the Transformer module or self-attention mechanism to capture long-distance dependency information, and use the bidirectional feature pyramid network (BiFPN) to fuse multi-scale features at different levels to generate the final multi-scale feature map, which is recorded as the alienated feature map.

[0067] The introduction of attention mechanism and dynamic weight generator enhances the fusion effect of shallow and deep features. The use of multi-branch structure, hole convolution and Transformer module significantly improves the model's ability to extract contextual information. Through shallow feature enhancement branches and dynamic feature pyramid adjustment, the detection ability of small targets and complex scenes is improved. It effectively solves the detection problem of multi-scale targets in track foreign object detection, and further optimizes the detection accuracy and robustness of the model.

[0068] Step S3: Use the sample set to pre-train the G-YOLOv5 model to obtain the trained G-YOLOv5 model; Model design module: used to design the G-YOLOv5 model and pre-train the G-YOLOv5 model; Use the sample set to pre-train the G-YOLOv5 model and obtain the trained G-YOLOv5 model. The method includes: The input of the G-YOLOv5 model is an image, and the output is a detection report, which includes a bounding box and foreign object recognition results. The foreign object recognition results indicate the type and number of foreign objects recognized. Input each sample image in the sample set into the G-YOLOv5 model in turn for training, and perform enhancement processing using constraints and optimization strategies until all sample sets are input, and then obtain the trained G-YOLOv5 model.

[0069] Step S4: collect and process the real-time image of the track, import the processed image into the trained G-YOLOv5 model, and obtain a real-time detection report; Collect and process real-time images of the track, import the processed images into the trained G-YOLOv5 model, and obtain real-time detection reports. The methods include: Use high-definition cameras to collect real-time images of the track, including dynamic images and static images; Import the processed real-time image into the trained G-YOLOv5 model to obtain a real-time detection report, which is used to output the category and quantity of foreign objects in the real-time image; If the number of foreign objects is 0, it means that there are no foreign objects; otherwise, there are foreign objects.

[0070] Step S5: arranging the real-time detection report on the visual interface; The visual interface includes the terminal page of the computer or the display interface of the mobile device, and interacts with the user.

[0071] The user clicks, queries, downloads and other interactive operations on the test report through the computer page (computer display screen, where the computer refers to a smart device, and the mobile device can be a mobile phone or tablet).

[0072] Real-time processing and visualization module: used for real-time image processing of tracks and layout of real-time detection reports on a visualization interface.

[0073] Example 2, please refer to Figure 1 As shown, the part not described in detail in this embodiment is described in Example 1, and a track foreign body detection system based on the YOLO algorithm is provided, including: Image preprocessing module: collects images of foreign objects on the track to form a data set, and obtains a sample set through sample expansion processing and label processing; Model design module: used to design the G-YOLOv5 model and pre-train the G-YOLOv5 model; Real-time processing and visualization module: used for real-time image processing of tracks and layout of real-time detection reports on a visualization interface.

[0074] Embodiment 3. This embodiment discloses an electronic device, including a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, the operation mode of the track foreign object detection method based on the YOLO algorithm provided above is implemented.

[0075] Since the electronic device introduced in this embodiment is an electronic device used to implement a track foreign object detection method based on the YOLO algorithm in the embodiment of the present application, based on the track foreign object detection method based on the YOLO algorithm introduced in the embodiment of the present application, the technical personnel of the field can understand the specific implementation of the electronic device of the present embodiment and its various variations, so how the electronic device implements the method in the embodiment of the present application is not introduced in detail here. As long as the technical personnel of the field implement the electronic device used in the track foreign object detection method based on the YOLO algorithm in the embodiment of the present application, it belongs to the scope of protection of this application.

[0076] The above formulas are all dimensionless and numerical calculations. The formula is a formula for the most recent real situation obtained by collecting a large amount of data and performing software simulation. The preset parameters and thresholds in the formula are set by technicians in this field according to actual conditions.

[0077] The above is only a preferred embodiment of the present invention, and the protection scope of the present invention is not limited to the above embodiments. All technical solutions under the concept of the present invention belong to the protection scope of the present invention. It should be pointed out that for ordinary technical users in this technical field, some improvements and modifications without departing from the principle of the present invention should also be regarded as the protection scope of the present invention.

Claims

1. A track foreign object detection method based on the YOLO algorithm, characterized in that: include: Step S1: Collect images of foreign objects on the track to form a data set, perform sample expansion processing and label processing on each image in the data set, and then obtain a sample set; Step S2: Add channel attention mechanism and direction perception mechanism to YOLOv5, and set constraints to obtain the improved YOLOv5 algorithm, which is recorded as G-YOLOv5 model; Step S3: Use the sample set to pre-train the G-YOLOv5 model to obtain the trained G-YOLOv5 model; Step S4: collect and process the real-time image of the track, import the processed image into the trained G-YOLOv5 model, and obtain a real-time detection report; Step S5: arranging the real-time detection report on the visual interface.

2. According to the YOLO algorithm-based track foreign body detection method of claim 1, it is characterized in that: The method for collecting images of track foreign matter to form a data set comprises: High-definition cameras are installed around the track to capture images at different time periods. Dynamic images are collected by extracting video frames, and static images are collected by acquiring single images. Images including foreign objects in static images and dynamic images are extracted as foreign object images, and the data set includes all types of foreign object images.

3. According to claim 2, a track foreign body detection method based on the YOLO algorithm is characterized in that: The method of performing sample expansion processing and label processing on each image in the data set to obtain a sample set includes: The target box is obtained through the YOLOv5 algorithm, and the target box is defined as , randomly select crop ratio , according to the clipping formula , calculate the cropped image ; Among them, the cropped image The target foreign body is included. is the coordinate of the upper left corner of the target box, and represent the horizontal and vertical coordinates respectively, is the width of the target box, is the height of the target box, crop ratio , Indicates the image corresponding to the target box; After cropping the image Perform brightness adjustment and weather interference respectively to obtain the image after preliminary processing , use CutMix enhancement technology to combine any two images Mixing to generate new images , using Mosaic enhancement technology to transform any image larger than two , and Stitch them together to generate multiple target images , then the image , , and are all regarded as expanded sample images; For each sample image, use the pre-trained Faster R-CNN model to obtain a pre-annotated image, use the annotation tool to adjust the boundary and correct the category label of the target box of the pre-annotated image, and collect all the images after boundary adjustment and category label correction to form a sample set; The format of setting annotation is ,and It is the foreign body category label.

4. According to the YOLO algorithm-based track foreign body detection method of claim 3, it is characterized in that: The cropped image Perform brightness adjustment and weather interference respectively to obtain the image after preliminary processing , methods include: Use the improved formula to The brightness is adjusted, and the adjusted image is recorded as , improved formula , ; in, represents the brightness adjustment factor, and , represents the brightness offset, and , represents the contrast adjustment factor, and , represents the contrast offset, and , Representing images The contrast ratio, Represents the adjusted image The contrast value of Use multivariate formulas to group images The environment is adjusted to obtain images after rainy day simulation, foggy day simulation and snowy day simulation respectively. The adjusted images are also recorded as ; Among them, the multivariate formula group , Indicated in the image The function that simulates raindrops on Indicates the intensity of controlling raindrops. Represents the parameters that control the fog density, Represents the image Apply Gaussian blur, Represents a randomly generated snowflake effect. Indicates that the randomly generated snowflake effect follows the Bernoulli distribution. Represents the parameter that controls the density of snowflakes.

5. According to the YOLO algorithm-based track foreign body detection method of claim 4, it is characterized in that: The channel attention mechanism and the direction perception mechanism are added to YOLOv5, and constraints are set to obtain an improved YOLOv5 algorithm, which is recorded as a G-YOLOv5 model. The specific method includes: Set the input of the G-YOLOv5 model as the input feature map and transfer it to the backbone network; The channel attention mechanism is introduced into the coordinate attention module to generate adaptive channel weighted feature maps, and then the alienated feature maps at different resolutions are extracted through multi-scale feature fusion; Decode the alienated feature map, remove redundant boxes through non-maximum suppression, retain the box with the highest confidence, and output multi-scale bounding boxes and foreign object recognition results; The constraint condition is set to minimize the loss function. The constraint condition formula is expressed as: ； , and is the first loss function , the second loss function And the third loss function The adaptive weight of , , , then the weight adaptive adjustment functions are , , ; in, represents the true category label of the i-th foreign object in the image, represents the predicted category label of the i-th foreign body in the image, Represents the intersection area of ​​the predicted bounding box output by the model and the true bounding box, Represents the union area of ​​the predicted bounding box output by the model and the true bounding box, Represents the area of ​​the smallest enclosing rectangle between the predicted bounding box output by the model and the true bounding box, represents the square of the length of the diagonal in the minimum enclosing rectangle, represents the true confidence, is the prediction confidence of the model output; The optimization strategy is to use the Adam optimizer for training, set the initial learning rate to 0.001, and use cosine annealing for learning rate scheduling.

6. A track foreign body detection method based on the YOLO algorithm according to claim 5, characterized in that: The method of introducing a channel attention mechanism into a coordinate attention module to generate an adaptive channel weighted feature map includes: Step A1: Assume that the input feature map contains td channels, and the spatial dimension of each channel is , calculate each position in the input feature map The average value and maximum value of all channels are recorded as the single channel mean and the single channel maximum value respectively. The single channel mean and the single channel maximum value are spliced ​​in the spatial dimension of the channel to obtain the fused spatial feature map; Use a convolution layer with a convolution kernel size of 7×7 to process the spatial feature map, and generate spatial attention weights through the Sigmoid activation function. Apply the generated spatial attention weights to the input feature map, weight the features of each spatial position, and obtain a weighted feature map, which is recorded as the remade feature map; Perform global average pooling on the spatial dimension of the remade feature map, calculate the global average value of each channel, collect the global average values ​​of td channels to form a td-dimensional vector, recorded as global channel information; Step A2: Use a linear layer network with global channel information as input to predict the scaling factor , based on the predicted scaling factor and the preset original reduction rate, calculate the dynamic reduction rate ; in, represents the Sigmoid activation function, represents the weight matrix of the linear layer, Represents global channel information, represents the bias term, represents the original reduction rate; Step A3: The global context information is input into a two-layer MLP. The first layer of the MLP uses a weight matrix for dimensionality reduction and uses a ReLU activation function. The second layer of the MLP uses a weight matrix for dimensionality increase and outputs it. Then, the channel attention weight is generated through a Sigmoid activation function. The generated channel attention weight is element-wise multiplied with the input feature map to obtain a weighted channel feature map. Among them, the weight matrix used in the first layer is different from that used in the second layer.

7. A track foreign body detection method based on the YOLO algorithm according to claim 6, characterized in that: The method of extracting alienation feature maps at different resolutions by fusing multi-scale features includes: Step B1: Extract shallow features and deep features from the backbone network, add an enhancement path to the shallow features of the backbone network, and use the SE module to align the channels of the shallow and deep features; Step B2: Use a dynamic weight generator to adaptively adjust the fusion ratio of shallow and deep features, perform weighted fusion of shallow and deep features after channel alignment according to the generated weight, and add a residual term; Step B3: Use branches with different receptive fields to extract fused features, use the Transformer module or self-attention mechanism to capture long-distance dependency information, use the bidirectional feature pyramid network to fuse multi-scale features at different levels, and generate the final multi-scale feature map, which is recorded as the alienated feature map.

8. The track foreign body detection method based on the YOLO algorithm according to claim 7, characterized in that: The method of pre-training the G-YOLOv5 model using the sample set to obtain the trained G-YOLOv5 model includes: The input of the G-YOLOv5 model is an image, and the output is a detection report, which includes a bounding box and foreign object recognition results. The foreign object recognition results indicate the type and number of foreign objects recognized. Input each sample image in the sample set into the G-YOLOv5 model in turn for training, and perform enhancement processing using constraints and optimization strategies until all sample sets are input, and then obtain the trained G-YOLOv5 model.

9. A track foreign body detection method based on the YOLO algorithm according to claim 8, characterized in that: The real-time image of the track is collected and processed, and the processed image is imported into the trained G-YOLOv5 model to obtain a real-time detection report. The method includes: Use high-definition cameras to collect real-time images of the track, including dynamic images and static images; Import the processed real-time image into the trained G-YOLOv5 model to obtain a real-time detection report, which is used to output the category and quantity of foreign objects in the real-time image; If the number of foreign objects is 0, it means that there are no foreign objects; otherwise, there are foreign objects.

10. A track foreign body detection method based on the YOLO algorithm according to claim 9, characterized in that: The visualization interface includes a terminal page of a computer or a display interface of a mobile device, and performs interface interaction with the user.

Images