Method for identifying railway on-track tool by using neural network model

Abstract

The invention relates to a method for identifying a railway on-track tool by using a neural network model. The method comprises the following steps: acquiring a tool image of the railway on-track tool; inputting the tool image into the backbone network, and extracting features of the tool image to obtain image feature information; inputting the image feature information into a CBAM module, and performing image feature enhancement on the image feature information to obtain an attention-enhanced feature map; inputting the attention enhancement feature map into the neck network and the head network in sequence to obtain the position and the category of the railway on-track tool in the tool image; according to the method, the CBAM module is introduced into the neural network model, so that the small target identification capability of the network is improved, the identification process is simple, and the risk of missing detection and false detection is reduced.

Description

Technical Field

[0001] The present disclosure relates to the technical field of image recognition, and in particular to a method for identifying railway on-track tools by using a neural network model. Background Art

[0002] In view of the current situation of my country's railway system and its future development needs, it is necessary to strengthen the safety management mechanism of the railway system. However, there are currently various factors that restrict the development of the safety management mechanism of the railway system, resulting in major problems in the safety management of railway track workers. In serious cases, it may lead to major railway accidents and cause significant losses of life and property.

[0003] With the vigorous development of railway industry, the types and quantity of railway on-road tools are also increasing. Considering the complexity of maintenance site and the diversity of on-road tools needed by operators, if the on-road tools are left on the railway line due to improper operation of maintenance personnel, the consequences will be very serious. Therefore, the research on railway on-road tool management system is particularly important. At present, railway practitioners mainly use tool boxes to properly store and manage the on-road tools they carry.

[0004] Generally, the identification methods of railway on-road tools mostly rely on special detection belts and security inspection channel doors. Workers need to put the on-road tools on the detection belt and pass through the security inspection channel door and other equipment to ensure that the on-road tools are not missed. These huge equipment not only occupy a large area, but also are difficult to move, and require a special room to place the detection equipment. In addition, special personnel are required to participate in the detection, which increases the operating cost. Workers are required to paste special detection patches on the on-road tools. These detection patches contain positioning modules, maintenance identification codes (QR codes), individual numbers, etc., and these patches are used as targets for detection. Each on-road tool needs to be pasted. There will be occasional situations where a large number of tools are needed during on-site work. When these situations occur, pasting detection patches will reduce work efficiency. If the detection patch is lost, even if the on-road tool has been taken out, it will be identified as the on-road tool being lost. Summary of the invention

[0005] The present invention provides a method for identifying railway on-track tools by using a neural network model, so as to solve the problems of complicated procedures and high risks of missed detection and false detection in existing railway on-track tool identification methods.

[0006] In a first aspect, the present disclosure provides a method for identifying a railway on-road tool using a neural network model, wherein the neural network model includes a backbone network, an attention module, a neck network, and a head network, including:

[0007] Acquire tool images of railway on-track tools;

[0008] Inputting the tool image into the backbone network, extracting features of the tool image, and obtaining image feature information;

[0009] Inputting the image feature information into the attention module, performing image feature enhancement on the image feature information, and obtaining an attention enhancement feature map;

[0010] The attention enhancement feature map is input into the neck network and the head network in sequence to obtain the position and category of the railway on-road tool in the tool image.

[0011] In some embodiments, the attention module includes a channel attention module and a spatial attention module;

[0012] The steps of inputting the image feature information into the attention module, performing image feature enhancement on the image feature information, and obtaining an attention enhanced feature map include:

[0013] Inputting the image feature information into the channel attention module to obtain a channel attention feature map;

[0014] Inputting the image feature information into the spatial attention module to obtain a spatial attention feature map;

[0015] Based on the channel attention feature map and the spatial attention feature map, the attention enhancement feature map is obtained.

[0016] In some embodiments, the step of inputting the image feature information into a channel attention module to obtain a channel attention feature map includes:

[0017] Inputting the image feature information into a channel attention module to perform maximum pooling operation and average pooling operation to obtain two corresponding first intermediate feature maps;

[0018] Input the corresponding two first intermediate feature maps into the multilayer perceptron to obtain a first feature map;

[0019] Input the first feature map into an activation function to obtain a channel attention vector;

[0020] Multiplying the channel attention vector by the tool image to obtain a channel attention feature map;

[0021] The step of inputting the image feature information into a spatial attention module to obtain a spatial attention feature map comprises:

[0022] Inputting the image feature information into the spatial attention module to perform maximum pooling operation and average pooling operation to obtain two corresponding second intermediate feature maps;

[0023] Connect the corresponding two second intermediate feature maps along the channel dimension and input them into the convolution layer to obtain the second feature map;

[0024] Input the second feature map into an activation function to obtain a spatial attention vector;

[0025] The spatial attention vector is multiplied by the tool image to obtain a spatial attention feature map.

[0026] In some embodiments, the training process of the neural network model includes:

[0027] Annotating a tool image of a railway on-road tool to obtain an annotation frame of the tool image and its true category, wherein the annotation frame is used to indicate a true position of the railway on-road tool;

[0028] Inputting a tool image into the backbone network, extracting features of the tool image, and obtaining image feature information;

[0029] Inputting the image feature information into the attention module, performing image feature enhancement on the image feature information, and obtaining an attention enhancement feature map;

[0030] Inputting the attention-enhanced feature map into the neck network and the head network in sequence to obtain a prediction box of the tool image and its prediction category, wherein the prediction box is used to indicate a predicted position of the tool on the railway;

[0031] Calculating the loss value between the predicted box and its predicted category and the marked box and its true category;

[0032] The parameters of the backbone network, CBAM module, neck network and head network are updated based on the loss value.

[0033] In some embodiments, the step of inputting the attention enhancement feature map into the neck network and the head network in sequence, and obtaining the position and category of the railway on-road tool in the tool image comprises:

[0034] Inputting the attention enhancement feature map into the neck network, integrating the attention enhancement feature map at different scales, and obtaining an integrated attention enhancement feature map;

[0035] The integrated attention-enhanced feature map is input into the head network, and feature recognition is performed on the integrated attention-enhanced feature map to obtain the position and category of the railway on-track tools in the tool image.

[0036] In some embodiments, before the tool image is input into the backbone network, the tool image is preprocessed by one or more preprocessing operations of scaling, rotating or enhancing brightness contrast to obtain a preprocessed tool image.

[0037] In a second aspect, the present disclosure provides a device for identifying railway on-road tools using a neural network model, comprising:

[0038] An image acquisition module, used to acquire tool images of railway on-track tools;

[0039] A feature information acquisition module, used to input the tool image into the backbone network, extract features of the tool image, and obtain image feature information;

[0040] A feature map acquisition module, used to input the image feature information into the attention module, perform image feature enhancement on the image feature information, and obtain an attention enhancement feature map;

[0041] The recognition module is used to input the attention enhancement feature map into the neck network and the head network in sequence to obtain the position and category of the railway on-road tool in the tool image.

[0042] In a third aspect, the present disclosure provides a computer device, including a memory, a processor, and a computer program stored in the memory, wherein the processor executes the computer program to implement the steps of the method described in the above aspects.

[0043] In a fourth aspect, the present disclosure provides a computer-readable storage medium having a computer program stored thereon, which implements the steps of the method described in the above aspects when executed by a processor.

[0044] In a fifth aspect, the present disclosure provides a computer program product, including a computer program, which implements the steps of the method described in the above aspects when executed by a processor.

[0045] The beneficial effects of the present invention are as follows:

[0046] The present invention provides a method for identifying railway on-road tools by using a neural network model, which comprises the following steps: obtaining a tool image of the railway on-road tool; inputting the tool image into a backbone network, extracting features of the tool image, and obtaining image feature information; inputting the image feature information into an attention module, performing image feature enhancement on the image feature information, and obtaining an attention-enhanced feature map; inputting the attention-enhanced feature map into a neck network and a head network in sequence, and obtaining the position and category of the railway on-road tool in the tool image; the present invention improves the network's recognition ability for small targets by introducing an attention module into a neural network model, simplifies the recognition process, and reduces the risk of missed detection or false detection. BRIEF DESCRIPTION OF THE DRAWINGS

[0047] The present disclosure will be described in more detail below based on embodiments and with reference to the accompanying drawings:

[0048] Figure 1 A schematic flow chart of a method for identifying railway on-track tools using a neural network model provided in an embodiment of the present disclosure;

[0049] Figure 2 A detailed flow chart of a method for identifying railway on-track tools using a neural network model provided in an embodiment of the present disclosure;

[0050] Figure 3 A schematic diagram of a tool image of a railway on-track tool provided in an embodiment of the present disclosure;

[0051] Figure 4 A flowchart of step S101 in a method for identifying a railway on-track tool using a neural network model provided in an embodiment of the present disclosure;

[0052] Figure 5 A schematic diagram of the marking of the railway on-track tool provided in the embodiment of the present disclosure;

[0053] Figure 6 A schematic diagram of the Focuson-YOLOv5s network model structure provided in an embodiment of the present disclosure;

[0054] Figure 7 A schematic diagram of the attention module processing flow provided in an embodiment of the present disclosure;

[0055] Figure 8 A schematic diagram of the training process of the neural network model provided in the embodiment of the present disclosure;

[0056] Fig. 9 A schematic diagram of the structure of a device for identifying railway track tools using a neural network model provided in an embodiment of the present disclosure;

[0057] Fig.10 A schematic diagram of the computer device structure provided in the embodiment of the present disclosure;

[0058] In the drawings, the same reference numerals are used for the same components and the drawings are not drawn to scale. DETAILED DESCRIPTION

[0059] In order to enable those skilled in the art to better understand the technical solution of the present disclosure, and to fully understand and implement how the present disclosure applies technical means to solve technical problems and achieve the corresponding technical effects, the technical solution in the embodiments of the present disclosure will be clearly and completely described below in conjunction with the drawings in the embodiments of the present disclosure. Obviously, the described embodiments are only embodiments of a part of the present disclosure, not all of the embodiments. The embodiments of the present disclosure and the various features in the embodiments can be combined with each other without conflict, and the technical solutions formed are all within the scope of protection of the present disclosure. Based on the embodiments in the present disclosure, all other embodiments obtained by ordinary technicians in this field without making creative work should fall within the scope of protection of the present disclosure.

[0060] It should be noted that the terms "first", "second", etc. in the specification and claims of the present disclosure and the above-mentioned drawings are used to distinguish similar objects, and are not necessarily used to describe a specific order or sequence. It should be understood that the data used in this way can be interchangeable where appropriate, so that the embodiments of the present disclosure described herein can be implemented in an order other than those illustrated or described herein. In addition, the terms "including" and "having" and any variations thereof are intended to cover non-exclusive inclusions, for example, a process, method, system, product, or device that includes a series of steps or units is not necessarily limited to those steps or units that are clearly listed, but may include other steps or units that are not clearly listed or inherent to these processes, methods, products, or devices.

[0061] It should be noted that the steps shown in the flowcharts of the accompanying drawings can be executed in a computer system such as a set of computer executable instructions, and that, although a logical order is shown in the flowcharts, in some cases, the steps shown or described can be executed in an order different from that shown here.

[0062] In view of the current situation of my country's railway system and its future development needs, it is necessary to strengthen the safety management mechanism of the railway system. However, there are currently various factors that restrict the development of the safety management mechanism of the railway system, resulting in major problems in the safety management of railway track workers. In serious cases, it may lead to major railway accidents and cause significant losses of life and property.

[0063] With the vigorous development of the railway industry, the types and quantity of railway tools are also increasing. Considering the complexity of the maintenance site and the diversity of railway tools that operators need to use, if the railway tools are left on the railway line due to improper operation by maintenance personnel, the consequences will be very serious. Therefore, the research on the railway tool management system is particularly important. At present, railway practitioners mainly use tool boxes to properly store and manage the railway tools they carry.

[0064] Most of the methods for identifying railway road tools rely on special detection belts and security inspection channel doors. Workers need to put the railway road tools on the detection belt and pass through the security inspection channel door and other equipment to ensure that the railway road tools are not missed. These huge equipment not only occupy a large area, but also are difficult to move, and require a special room to place the detection equipment. In addition, special personnel are required to participate in the detection, which increases the operating cost. Workers are required to paste special detection patches on the railway road tools. These detection patches contain positioning modules, maintenance identification codes (QR codes), individual numbers, etc., and these patches are used as targets for detection. Each railway road tool needs to be pasted. There will be occasional situations where a large number of railway road tools are needed during on-site work. When these situations occur, pasting detection patches will reduce work efficiency. If the detection patch is lost, even if the railway road tool has been taken out, it will be identified as the railway road tool being lost.

[0065] Therefore, the existing railway on-track tool identification methods have the problems of large equipment involved, complicated processes, and high risk of missed detection and false detection.

[0066] Embodiment 1

[0067] Figure 1 A schematic diagram of a process for identifying railway on-road tools using a neural network model provided in an embodiment of the present disclosure.

[0068] like Figure 1 As shown, an embodiment of the present invention provides a method for identifying railway on-road tools using a neural network model, wherein the neural network model includes a backbone network, an attention module, a neck network, and a head network, including:

[0069] Acquire tool images of railway on-track tools;

[0070] Inputting the tool image into the backbone network, extracting features of the tool image, and obtaining image feature information;

[0071] Input the image feature information into the attention module, perform image feature enhancement on the image feature information, and obtain an attention enhanced feature map;

[0072] The attention-enhanced feature map is sequentially input into the neck network and the head network to obtain the location and category of the railway on-track tools in the tool image.

[0073] Figure 2 Based on Figure 1 An embodiment provides a method for identifying railway on-road tools using a neural network model.

[0074] like Figure 2As shown, an embodiment of the present invention provides a method for identifying railway on-road tools using a neural network model. In this embodiment, the neural network model is a Focuson-YOLOv5s network model, and the Focuson-YOLOv5s network model includes a backbone network, a CBAM module 602, a neck network, and a head network;

[0075] Methods include:

[0076] S201: Collect photos of railway on-road tools to obtain tool images of railway on-road tools;

[0077] S202: Introduce the CBAM module 602 into the YOLOv5 network to build a Focuson-YOLOv5s network model;

[0078] S203: Input the tool image into the Focuson-YOLOv5s network model for iterative training to obtain a trained Focuson-YOLOv5s network model;

[0079] S204: Input the tool image of the railway tool to be identified into the trained Focuson-YOLOv5s network model to obtain the position and category of the railway tool to be identified.

[0080] In an optional embodiment of the present invention, step S201 includes:

[0081] Collect photos of the railway tools that railway workers bring with them for on-track maintenance, and obtain tool images of the railway tools, such as Figure 3 Shown is a schematic diagram of a tool image of a railway on-track tool.

[0082] Figure 4 A flowchart of step S201 in the railway on-track tool identification method provided in an embodiment of the present disclosure.

[0083] like Figure 4 As shown, in an optional embodiment of the present invention, collecting photos of railway on-road tools and obtaining tool images of railway on-road tools include:

[0084] S401: storing the photo of the railway on-road tool into the Images file to obtain the original tool image;

[0085] S402: Based on the LabelImg labeling tool, the original tool image is labeled to obtain the labeled tool image, that is, the labeling box of the tool image and its real category. The labeling box is used to indicate the real position of the tool on the railway. That is, the labeling content includes the position and category of the tool on the railway, such as Figure 5 The figure shows a schematic diagram of marking a railway on-track tool provided by an embodiment of the present disclosure;

[0086] S403: Convert the format of the labeled tool image and save it in the Labels folder to obtain the tool image.

[0087] Specifically, the LabelImg annotation tool is an open source image annotation tool used to create and edit object bounding boxes on images. It is usually used for data annotation for target detection tasks. LabelImg has a user-friendly interface and supports a variety of common image formats, such as JPEG, PNG, etc., and can save annotation information in common annotation formats, such as PASCALVOC, YOLO, etc.

[0088] Furthermore, in an optional embodiment of the present invention, a labeling tool such as labelme, MakeSense or CVAT may be used to label the original tool image.

[0089] Further, in an optional embodiment of the present invention, before the tool image is input into the backbone network, the tool image is preprocessed by one or more preprocessing operations of scaling, rotating or enhancing brightness contrast to obtain a preprocessed tool image;

[0090] Specifically, the tool image is preprocessed to obtain a preprocessed tool image; since the size, placement direction, and lighting conditions of the railway on-track tools used for on-track maintenance are random, the tool image is preprocessed through one or more preprocessing operations such as scaling, rotating, or enhancing brightness and darkness contrast to expand the scene of the tool image.

[0091] Figure 6 A schematic diagram of the Focuson-YOLOv5s network model structure provided in an embodiment of the present disclosure.

[0092] like Figure 6 As shown, in this embodiment, the backbone network is a backbone network 601, the CBAM module 602, the neck network is a neck network 603, and the head network is a head output terminal 604.

[0093] The backbone network 601 mainly uses spatial pyramid pooling (SPP) and cross stage partial network (CSP) structures to obtain image features. The SPP module performs maximum pooling (Max Pooling) on ​​the feature information obtained by the convolution layer at different scales, and then splices the pooled feature maps together to retain significant feature information while removing redundant feature information.

[0094] The CBAM module 602 (Convolut iona l Block Attention Module, attention mechanism module) is a module used to enhance the performance of a convolutional neural network (CNN). By introducing an attention mechanism, it can dynamically learn the relationship between different spatial positions and channels, thereby improving the network's ability to represent image features.

[0095] After the CBAM module 602 is added to the backbone network 601, the present invention performs comprehensive screening on the complex image features obtained in the backbone network 601 to play a role in feature weight adjustment.

[0096] In an optional embodiment of the present invention, inputting the tool image into the Focuson-YOLOv5s network model for iterative training, and obtaining the trained Focuson-YOLOv5s network model includes:

[0097] The preprocessed tool image is input into the backbone network 601, the features of the tool image are extracted to obtain image feature information, the image feature information is input into the CBAM module 602, the image feature information is enhanced to obtain an attention enhancement feature map, the attention enhancement feature map is input into the neck network 603, the attention enhancement feature map is integrated at different scales to obtain an integrated attention enhancement feature map, the integrated attention enhancement feature map is input into the head output terminal 604, the integrated attention enhancement feature map is feature recognized to obtain the position and category of the railway on-road tools in the tool image; all tool images are input into the Focuson-YOLOv5s network model in turn for training, and a C-IoU Loss loss function is constructed, the weights of each module in the Focuson-YOLOv5s network model are updated to obtain a trained Focuson-YOLOv5s network model.

[0098] Specifically, the preprocessed tool image is input into the backbone network 602, and the features of the tool image are extracted to obtain image feature information including:

[0099] The railway on-road tool image with a size of 3×640×640 in the preprocessed tool image is input into the backbone network 601, and the Focus module in the backbone network 601 performs a slicing operation, which is divided into 4 small blocks of 320×320 and re-spliced ​​to obtain a 12×320×320 feature map; the preliminary feature map is sequentially subjected to spatial pyramid pooling SPP and cross-stage local network CSP to obtain image feature information.

[0100] Furthermore, the attention module includes a channel attention module and a spatial attention module;

[0101] The steps of inputting the image feature information into the attention module, performing image feature enhancement on the image feature information, and obtaining the attention enhanced feature map include:

[0102] Input the image feature information into the channel attention module to obtain the channel attention feature map;

[0103] Inputting the image feature information into the spatial attention module to obtain a spatial attention feature map;

[0104] Based on the channel attention feature map and the spatial attention feature map, an attention enhancement feature map is obtained.

[0105] The steps of inputting the image feature information into the channel attention module and obtaining the channel attention feature map include:

[0106] Input the image feature information into the channel attention module to perform maximum pooling operation and average pooling operation to obtain the corresponding two first intermediate feature maps;

[0107] Input the corresponding two first intermediate feature maps into the multilayer perceptron to obtain a first feature map;

[0108] Input the first feature map into the activation function to obtain the channel attention vector;

[0109] Multiplying the channel attention vector by the tool image to obtain a channel attention feature map;

[0110] The steps of inputting the image feature information into the spatial attention module and obtaining the spatial attention feature map include:

[0111] The image feature information is input into the spatial attention module to perform maximum pooling and average pooling operations to obtain two corresponding second intermediate feature maps;

[0112] Connect the corresponding two second intermediate feature maps along the channel dimension and input them into the convolution layer to obtain the second feature map;

[0113] Input the second feature map into the activation function to obtain the spatial attention vector;

[0114] The spatial attention vector is multiplied by the tool image to obtain a spatial attention feature map.

[0115] Specifically, the CBAM module 602 is a lightweight module including a channel attention module 705 (CAM attention module) and a spatial attention module 706 (SAM attention module);

[0116] Among them, Figure 7As shown, the CAM attention module includes a first maximum pooling operation (MaxPooling) module 707, a first average pooling operation (AveragePooling) module 708, a shared multi-layer perceptron (Share MLP) module 709 and a first activation function module 710. The function of the CAM attention module is to make the network pay more attention to important feature information; the SAM attention module includes a second maximum pooling operation module 711, a second average pooling operation module 712 and a second activation function module 713. The function of the SAM attention module is to make the network pay more attention to the location information of the object;

[0117] The steps of inputting the image feature information into the CBAM module 602, performing image feature enhancement on the image feature information, and obtaining an attention enhancement feature map include:

[0118] The image feature information is input into the first maximum pooling operation module 707 and the first average pooling operation module 708 in the CAM attention module, and the maximum pooling operation and the average pooling operation are performed to obtain two corresponding first intermediate feature maps; the corresponding two first intermediate feature maps are input into the shared multi-layer perceptron module 709 for shared multi-level operation, the number of channels is compressed and then expanded to the original number of channels, and the first feature map is obtained; the first feature map is input into the first activation function module 710, and the channel attention vector is obtained through the activation function, and the channel attention vector is multiplied by the tool image to obtain a channel attention feature map, so as to weight and enhance useful features and suppress useless features;

[0119] The image feature information is input into the second maximum pooling operation module 711 and the second average pooling operation module 712 in the SAM attention module, and the maximum pooling operation and the average pooling operation are performed to obtain two corresponding second intermediate feature maps, and the convolution operation is performed to convert the corresponding two second intermediate feature maps into a second feature map of a channel, and the second feature map is input into the second activation function module 713, and the spatial attention vector is obtained through the activation function, and the spatial attention vector is multiplied by the tool image to obtain the spatial attention feature map;

[0120] Multiply the channel attention feature map and the spatial attention feature map element by element to obtain the attention enhancement feature map;

[0121] In an optional embodiment of the present invention, the activation function is a sigmoid activation function.

[0122] Furthermore, the attention-enhanced feature map is sequentially input into the neck network and the head network, and the steps of obtaining the position and category of the railway on-track tool in the tool image include:

[0123] Input the attention enhancement feature map into the neck network, integrate the attention enhancement feature map at different scales, and obtain the integrated attention enhancement feature map;

[0124] The integrated attention-enhanced feature map is input into the head network, and feature recognition is performed on the integrated attention-enhanced feature map to obtain the position and category of the railway on-track tools in the tool image.

[0125] Specifically, the attention enhancement feature map is input into the neck network 603, and the attention enhancement feature map is integrated at different scales to obtain an integrated attention enhancement feature map. The integrated attention enhancement feature map is input into the head output terminal 604, and feature recognition is performed on the integrated attention enhancement feature map to obtain the position and category of the railway on-road tool in the tool image, including:

[0126] The attention enhancement feature map is input into the neck network 603, and the attention enhancement feature map is integrated at different scales through the feature pyramid network FPN and the path aggregation network PAN to obtain an integrated attention enhancement feature map, thereby improving the network expression ability; the integrated attention enhancement feature map is input into the head output terminal 604, and the integrated attention enhancement feature map is divided into grids of 8×8, 16×16, and 32×32 pixels in the head output terminal 604, which correspond to three prediction maps respectively. The three prediction maps of different scales can be used to predict targets of different sizes; the position and length and width values ​​of the center point of the prediction box are obtained through the convolution layer. After obtaining all the required bounding boxes, the algorithm will screen them through confidence and non-maximum suppression (NMS), remove the prediction boxes with confidence less than 0.5 to obtain the position and number of the final prediction boxes, and if the IoU between the same two prediction boxes is greater than a certain threshold, remove the prediction box with a smaller confidence; obtain the prediction box of the tool image and its prediction category, and the prediction box is used to indicate the predicted position of the tool on the railway.

[0127] In an optional embodiment of the present invention, the training process of the neural network model includes:

[0128] Annotate the tool image of the railway on-track tool to obtain the annotation box of the tool image and its real category, and the annotation box is used to indicate the real position of the railway on-track tool;

[0129] Inputting the tool image into the backbone network, extracting features of the tool image, and obtaining image feature information;

[0130] Input the image feature information into the attention module, perform image feature enhancement on the image feature information, and obtain an attention enhanced feature map;

[0131] The attention enhanced feature map is sequentially input into the neck network and the head network to obtain a prediction box of the tool image and its prediction category, wherein the prediction box is used to indicate a predicted position of the tool on the railway;

[0132] Calculate the loss value between the predicted box and its predicted category, and the labeled box and its true category;

[0133] Update the parameters of the backbone network, CBAM module, neck network, and head network based on the loss value.

[0134] Specifically, the steps of constructing the C-IoU Loss loss function, updating the weights of each module in the Focuson-YOLOv5s network model, and obtaining the trained Focuson-YOLOv5s network model include:

[0135] Construct the C-IoU Loss loss function and use it to calculate the loss value between the predicted box and its predicted category and the labeled box and its true category;

[0136] Update the parameters of the backbone network, CBAM module, neck network and head network based on the loss value;

[0137] Update the weights of the backbone network, CBAM module, neck network and head output in the Focuson-YOLOv5s network model to obtain the trained Focuson-YOLOv5s network model.

[0138] The function expression of C-IoU Loss loss function is:

[0139]

[0140] Among them, L CIoU represents the value of C-IoU Loss; x gt ,y gt represents the coordinates of the center point of the true value of the annotated annotation box, x and y represent the coordinates of the center point of the predicted box, c represents the diagonal distance of the minimum enclosing rectangle, and W g Indicates the width of the minimum enclosing rectangle, H g represents the height of the minimum bounding rectangle, α is a parameter used to balance the ratio, γ is used to measure the proportional consistency between the predicted box and the labeled box, and is used to further adjust the loss function, considering the shape and direction of the labeled box, IoU represents the intersection over union ratio of the predicted box and the labeled box, w represents the width and height of the predicted box, h represents the width and height of the predicted box, w gt Indicates the width of the annotation box, h gt Indicates the height of the label box.

[0141] C-IoU Loss not only pays attention to the position consistency and size consistency between the predicted box and the annotated box, but also pays attention to the aspect ratio difference. For different repair tools, the aspect ratio is a very important feature. Therefore, the introduction of C-IoULoss can effectively improve the optimization effect of the model during training.

[0142] like Figure 8 The figure shows a training process diagram of the neural network model provided by the embodiment of the present disclosure. During the training process, the pre-processed tool image is input into the backbone network 601, the features of the tool image are extracted, and the image feature information is obtained;

[0143] The image feature information is input into the CBAM module 602, and the image feature information is enhanced to obtain an attention enhancement feature map;

[0144] The attention enhanced feature map is sequentially input into the neck network 603 and the head output terminal 604, and the prediction box of the tool image and its prediction category are obtained by forward reasoning calculation. The prediction box is used to indicate the predicted position of the tool on the railway.

[0145] The predicted box and the labeled box are input into the C_IoU Loss loss function module 814, and the loss value between the predicted box and its predicted category, the labeled box and its true category is calculated by the C_IoU Loss loss function, and the back-propagation gradient value is obtained. The weights / parameters of the backbone network 601, the CBAM module 602, the neck network 603 and the head output terminal 604 are reversely updated according to the back-propagation gradient value.

[0146] Input the image of the tool to be identified on the railway into the trained Focuson-YOLOv5s network model to obtain the location and category of the tool to be identified on the railway, including:

[0147] The image of the railway tool to be identified is input into the trained Focuson-YOLOv5s network model, and the features of the tool image are extracted through the backbone network 601 to obtain the image feature information. The image feature information is sequentially input into the CBAM module 602, the neck network 603 and the head output terminal 604 to obtain the position and category of the railway tool in the tool image.

[0148] Specifically, the railway on-road tool identification method can be applied on a mobile phone, and the specific detection process includes:

[0149] Select the image of the railway tool to be identified, upload the image of the railway tool to be identified to the Focuson-YOLOv5s network model on the mobile phone, select one of the recognition buttons in the CPU or GPU to perform tool recognition, obtain the recognition result of the railway tool to be identified, and display the recognition result on the mobile phone display control to display the tool name.

[0150] The present invention obtains the Focuson-YOLOv5s network model through training. Referring to Table 1, compared with the original YOLOv5s network and the network with other attention mechanisms, the detection accuracy is improved, especially the problem of missed detection caused by occlusion of objects is solved.

[0151] Table 1. Comparison of experimental results

[0152]

[0153] In the above table, mAP@0.5 means that when IoU is set to 0.5, the AP (AveragePrecision, which is the average precision of a single category) of all images in each category is calculated, and then the average of all categories is calculated;

[0154] mAP@0.95 means that when IoU is set to 0.95, the AP of all images in each category is calculated, and then the average of all categories is calculated.

[0155] mAP is the average precision across all categories.

[0156] The present invention discloses a method for identifying railway on-road tools by using a neural network model, which comprises the following steps: obtaining a tool image of the railway on-road tools; inputting the tool image into a backbone network, extracting features of the tool image, and obtaining image feature information; inputting the image feature information into a CBAM module, performing image feature enhancement on the image feature information, and obtaining an attention enhancement feature map; inputting the attention enhancement feature map into a neck network and a head network in sequence, and obtaining the position and category of the railway on-road tools in the tool image; the present invention improves the network's recognition ability for small targets by introducing a CBAM module into a YOLOv5s network, and the recognition process is simple, thereby reducing the risk of missed detection or false detection; the present invention can be used on a mobile terminal, and a method for realizing tool detection and recognition can be realized only through simple mobile terminal operations; the equipment space required for detection is greatly reduced, and the operator's hands required for detection are saved; and a photo of the tools in a tool box can be taken on the upper track and the lower track through the tool recognition function of the mobile terminal so as to realize tool management.

[0157] Embodiment 2

[0158] Based on the above embodiments, the CBAM module 501 in this embodiment can be replaced by other attention mechanisms, such as an SE attention mechanism, a self-attention mechanism, an ECA attention mechanism, or a CA attention mechanism.

[0159] Embodiment 3

[0160] Based on the above embodiment, in this embodiment, the YOLOv5 network can be replaced by the HOG algorithm.

[0161] Embodiment 4

[0162] Based on the above embodiments, Fig. 9 As shown, this embodiment provides a device for identifying railway on-road tools using a neural network model, including:

[0163] An image acquisition module 901 is used to acquire a tool image of a railway on-track tool;

[0164] A feature information acquisition module 902 is used to input the tool image into the backbone network, extract features of the tool image, and obtain image feature information;

[0165] A feature map acquisition module 903 is used to input the image feature information into the CBAM module, perform image feature enhancement on the image feature information, and obtain an attention enhancement feature map;

[0166] The recognition module 904 is used to input the attention enhancement feature map into the neck network and the head network in sequence to obtain the position and category of the railway on-track tool in the tool image.

[0167] Embodiment 5

[0168] Based on the above embodiments, Fig.10 As shown, this embodiment provides a computer device, including a processor 10, a memory 11, a communication bus 12 and a communication interface 13, and may also include a computer program stored in the memory 11 and executable on the processor 10, and the processor 10 executes the computer program to implement the steps of the method described in the above embodiment.

[0169] In some embodiments, the processor 10 may be composed of an integrated circuit, for example, a single packaged integrated circuit, or a plurality of packaged integrated circuits with the same or different functions, including one or more central processing units (CPUs), microprocessors, digital processing chips, graphics processors, and combinations of various control chips, etc. The processor 10 is the control core (ControlUnit) of the computer device, and uses various interfaces and lines to connect various components of the entire computer device, and executes or executes programs or modules stored in the memory 11 (for example, executing a converter noise analysis program based on an agent model, etc.), and calls data stored in the memory 11 to execute various functions of the computer device and process data.

[0170] The memory 11 includes at least one type of medium, including flash memory, mobile hard disk, multimedia card, card-type memory (for example: SD or DX memory, etc.), magnetic memory, disk, optical disk, etc. The memory 11 may be an internal storage unit of a computer device in some embodiments, such as a mobile hard disk of the computer device. The memory 11 may also be an external storage computer device of a computer device in other embodiments, such as a plug-in mobile hard disk, a smart memory card (Smart Media Card, SMC), a secure digital (Secure Digital, SD) card, a flash card (Flash Card), etc. equipped on the computer device. Further, the memory 11 may also include both an internal storage unit of a computer device and an external storage computer device. The memory 11 may not only be used to store application software and various types of data installed in the computer device, such as the code of a railway on-road tool identification program, but may also be used to temporarily store data that has been output or is to be output.

[0171] The communication bus 12 may be a peripheral component interconnect (PCI) bus or an extended industry standard architecture (EISA) bus, etc. The bus may be divided into an address bus, a data bus, a control bus, etc. The bus is configured to realize connection and communication between the memory 11 and at least one processor 10, etc.

[0172] The communication interface 13 is used for communication between the above-mentioned computer device and other computer devices, including a network interface and a user interface. Optionally, the network interface may include a wired interface and / or a wireless interface (such as a WI-FI interface, a Bluetooth interface, etc.), which is generally used to establish a communication connection between the computer device and other computer devices. The user interface may be a display (Display), an input unit (such as a keyboard (Keyboard)), and optionally, the user interface may also be a standard wired interface, a wireless interface. Optionally, in some embodiments, the display may be an LED display, a liquid crystal display, a touch-sensitive liquid crystal display, and an OLED (Organic Light-Emitting Diode, organic light-emitting diode) touch device, etc. Among them, the display may also be appropriately referred to as a display screen or a display unit, which is used to display information processed in the computer device and to display a visual user interface.

[0173] The figure only shows a computer device with components. Those skilled in the art will understand that the structure shown in the figure does not constitute a limitation on the computer device, and may include fewer or more components than shown in the figure, or combine certain components, or arrange the components differently.

[0174] For example, although not shown, the computer device may also include a power source (such as a battery) for supplying power to various components. Preferably, the power source may be logically connected to the at least one processor 10 through a power management device, so that the power management device can realize functions such as charging management, discharging management, and power consumption management. The power source may also include any components such as one or more DC or AC power sources, recharging devices, power failure detection circuits, power converters or inverters, power status indicators, etc. The computer device may also include a variety of sensors, Bluetooth modules, Wi-Fi modules, etc., which will not be repeated here.

[0175] It should be understood that the embodiment is for illustration only and the scope of the patent application is not limited to this structure.

[0176] The railway on-road tool identification method program stored in the memory 11 in the computer device is a combination of multiple instructions. When running in the processor 10, it can achieve:

[0177] Acquire tool images of railway on-track tools;

[0178] The tool image is input into a pre-built Focuson-YOLOv5s network model, and the backbone network in the pre-built Focuson-YOLOv5s network model is used to extract the features of the tool image to obtain image feature information. The image feature information is enhanced by using the CBAM module in the pre-built Focuson-YOLOv5s network model to obtain an attention enhanced feature map. The attention enhanced feature map is integrated at different scales by using the neck network in the pre-built Focuson-YOLOv5s network model to obtain an integrated attention enhanced feature map. The integrated attention enhanced feature map is subjected to feature recognition by using the head output end in the pre-built Focuson-YOLOv5s network model to obtain the position and category of the railway on-road tool in the tool image.

[0179] Specifically, the specific implementation method of the processor 10 for the above instructions can refer to the description of the relevant steps in the corresponding embodiment of the accompanying drawings, which will not be repeated here.

[0180] Furthermore, if the module / unit integrated in the computer device is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. The computer-readable storage medium can be volatile or non-volatile. For example, the computer-readable medium can include: any entity or device capable of carrying the computer program code, a recording medium, a USB flash drive, a mobile hard disk, a magnetic disk, an optical disk, a computer memory, and a read-only memory (ROM).

[0181] In some implementations of this embodiment, a computer-readable storage medium is provided, on which a computer program is stored, characterized in that when the computer program is executed by a processor, the steps of the method described in the above embodiment are implemented.

[0182] These program codes can also be loaded onto a computer or other programmable data processing device so that a series of operation steps are executed on the computer or other programmable device to produce a computer-implemented process, thereby providing instructions executed on the computer or other programmable device for implementing the process. Figure 1 The steps of a specified function in a process or multiple processes.

[0183] Computer-readable storage media include permanent and non-permanent, removable and non-removable media, and can be implemented by any method or technology to store information. Information can be computer-readable instructions, data structures, modules of programs or other data. Examples of computer-readable storage media can include, but are not limited to, phase change memory (PRAM), static random access memory (SRAM), dynamic random access memory (DRAM), other types of random access memory (RAM), read-only memory (ROM), electrically erasable programmable read-only memory (EEPROM), flash memory or other memory technology, compact disk read-only memory (CD-ROM), digital versatile disk (DVD) or other optical storage, magnetic cassettes, magnetic disk storage or other magnetic storage devices or any other non-transmission medium that can be used to store information that can be accessed by a computing device.

[0184] In some implementations of this embodiment, a computer program product is provided, including a computer program, characterized in that when the computer program is executed by a processor, the steps of the method described in the above embodiment are implemented.

[0185] In the embodiments provided in the present disclosure, it should be understood that the disclosed devices and methods can also be implemented in other ways. The device embodiments described above are merely schematic. For example, the flowcharts and block diagrams in the accompanying drawings show the possible architecture, functions and operations of the devices, methods and computer program products according to multiple embodiments of the present disclosure. In this regard, each box in the flowchart or block diagram can represent a module, a program segment or a part of a code, and the above-mentioned module, program segment or a part of the code contains one or more executable instructions for implementing the specified logical function. It should also be noted that in some alternative implementations, the functions marked in the box can also occur in a different order from the order marked in the accompanying drawings. For example, two consecutive boxes can actually be executed substantially in parallel, and they can sometimes be executed in the opposite order, depending on the functions involved. It should also be noted that each box in the block diagram and / or flowchart, and the combination of boxes in the block diagram and / or flowchart can be implemented with a dedicated hardware-based system that performs a specified function or action, or can be implemented with a combination of dedicated hardware and computer instructions.

[0186] It should be noted that in the present disclosure, the terms "include", "comprises" or any other variations thereof are intended to cover non-exclusive inclusion, so that a process, method, article or device including a series of elements includes not only those elements, but also other elements not explicitly listed, or also includes elements inherent to such process, method, article or device. In the absence of further restrictions, an element limited by the sentence "includes a ..." does not exclude the existence of other identical elements in the process, method, article or device including the element.

[0187] Although the embodiments disclosed in the present disclosure are as above, the above contents are only embodiments adopted for facilitating the understanding of the present disclosure and are not intended to limit the present disclosure. Any technician in the technical field to which the present disclosure belongs can make any modifications and changes in the form and details of the implementation without departing from the spirit and scope disclosed in the present disclosure, but the scope of patent protection of the present disclosure shall still be subject to the scope defined in the attached claims.

Claims

1. A method for identifying railway on-road tools using a neural network model, wherein the neural network model includes a backbone network, an attention module, a neck network and a head network, characterized in that: include: Acquire tool images of railway on-track tools; Inputting the tool image into the backbone network, extracting features of the tool image, and obtaining image feature information; Inputting the image feature information into the attention module, performing image feature enhancement on the image feature information, and obtaining an attention enhancement feature map; The attention enhancement feature map is input into the neck network and the head network in sequence to obtain the position and category of the railway on-road tool in the tool image.

2. The method according to claim 1, characterized in that The attention module includes a channel attention module and a spatial attention module; The steps of inputting the image feature information into the attention module, performing image feature enhancement on the image feature information, and obtaining an attention enhanced feature map include: Inputting the image feature information into the channel attention module to obtain a channel attention feature map; Inputting the image feature information into the spatial attention module to obtain a spatial attention feature map; Based on the channel attention feature map and the spatial attention feature map, the attention enhancement feature map is obtained.

3. The method according to claim 2, characterized in that The step of inputting the image feature information into a channel attention module to obtain a channel attention feature map comprises: Inputting the image feature information into a channel attention module to perform maximum pooling operation and average pooling operation to obtain two corresponding first intermediate feature maps; Input the corresponding two first intermediate feature maps into the multilayer perceptron to obtain a first feature map; Input the first feature map into an activation function to obtain a channel attention vector; Multiplying the channel attention vector by the tool image to obtain a channel attention feature map; The step of inputting the image feature information into a spatial attention module to obtain a spatial attention feature map comprises: Inputting the image feature information into the spatial attention module to perform maximum pooling operation and average pooling operation to obtain two corresponding second intermediate feature maps; Connect the corresponding two second intermediate feature maps along the channel dimension and input them into the convolution layer to obtain the second feature map; Input the second feature map into an activation function to obtain a spatial attention vector; The spatial attention vector is multiplied by the tool image to obtain a spatial attention feature map.

4. The method according to claim 1, characterized in that The training process of the neural network model includes: Annotating a tool image of a railway on-road tool to obtain an annotation frame of the tool image and its true category, wherein the annotation frame is used to indicate a true position of the railway on-road tool; Inputting a tool image into the backbone network, extracting features of the tool image, and obtaining image feature information; Inputting the image feature information into the attention module, performing image feature enhancement on the image feature information, and obtaining an attention enhancement feature map; Inputting the attention-enhanced feature map into the neck network and the head network in sequence to obtain a prediction box of the tool image and its prediction category, wherein the prediction box is used to indicate a predicted position of the tool on the railway; Calculating the loss value between the predicted box and its predicted category and the marked box and its true category; The parameters of the backbone network, CBAM module, neck network and head network are updated based on the loss value.

5. The method according to claim 1, characterized in that The step of inputting the attention enhancement feature map into the neck network and the head network in sequence to obtain the position and category of the railway on-road tool in the tool image comprises: Inputting the attention enhancement feature map into the neck network, integrating the attention enhancement feature map at different scales, and obtaining an integrated attention enhancement feature map; The integrated attention-enhanced feature map is input into the head network, and feature recognition is performed on the integrated attention-enhanced feature map to obtain the position and category of the railway on-track tools in the tool image.

6. The method according to any one of claims 1 to 5, characterized in that Before inputting the tool image into the backbone network, the tool image is preprocessed by one or more preprocessing operations of scaling, rotating or enhancing brightness contrast to obtain a preprocessed tool image.

7. A device for identifying railway on-road tools using a neural network model, characterized in that: include: An image acquisition module, used to acquire tool images of railway on-track tools; A feature information acquisition module, used to input the tool image into the backbone network, extract features of the tool image, and obtain image feature information; A feature map acquisition module, used to input the image feature information into the attention module, perform image feature enhancement on the image feature information, and obtain an attention enhancement feature map; The recognition module is used to input the attention enhancement feature map into the neck network and the head network in sequence to obtain the position and category of the railway on-road tool in the tool image.

8. A computer device comprising a memory, a processor and a computer program stored in the memory, characterized in that: The processor executes the computer program to implement the steps of the method according to any one of claims 1 to 6.

9. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 6 are implemented.

10. A computer program product, comprising a computer program, characterized in that When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 6 are implemented.

Images