Helmet wearing detection method based on YOLOv5 target detection model

Abstract

The present invention discloses a helmet wearing detection method based on a YOLOv5 target detection model. On the basis of the initial YOLOv5 target detection model, a feature extraction part of YOLOv5 is used to extract a feature map of a larger size, and feature fusion is performed in a feature fusion part. The center loss is used as a regularization term to constrain the confidence loss and classification loss in the YOLOv5 target detection model, so that the improved YOLOv5 target detection model can more accurately detect the category of small targets, wherein the feature map of a larger size has a smaller receptive field, and is therefore more conducive to the YOLOv5 target detection model detecting small targets. The improved YOLOv5 model proposed by the present invention has excellent performance in the detection task of helmets in construction scenarios, and some small targets that YOLOv5 cannot detect are detected by the improved YOLOv5.

Description

Technical Field

[0001] The present invention belongs to the technical field of target detection in computer vision, and in particular relates to a method for detecting helmet wearing based on a YOLOv5 target detection model. Background Art

[0002] In the current railway construction environment, the complex construction environment and unsafe behaviors of construction workers can easily lead to safety accidents, which may even threaten the lives of construction workers. Among them, head injuries of construction workers are the most fatal, which is also the main cause of safety accidents in construction scenes. Therefore, wearing a safety helmet is very necessary, and it is also the most basic requirement in the construction behavior code. However, in actual working conditions, people often fail to wear safety helmets. Therefore, in order to reduce safety hazards, it is very necessary to detect the wearing of safety helmets by people during construction.

[0003] With the continuous development of computer vision technology, the field of target detection has been widely used in industrial production, among which helmet detection technology is an important application field. Early detection algorithms based on manual features mainly include Viola-Jones detector, Histogram of Oriented Gradients (HOG), etc. With the development of convolutional neural networks, target detection algorithms have also turned to detection algorithms based on deep neural networks. Target detection algorithms based on deep neural networks can be divided into two categories, namely, "single-stage" target detection models and "double-stage" target detection models. Among them, the models belonging to "single-stage" target detection mainly include SSD and YOLO series. And the models belonging to "double-stage" target detection mainly include RCNN, Faster RCNN, etc. The present invention is mainly based on the "single-stage" YOLOv5 model to improve the detection of helmets.

[0004] The YOLOv5 model is improved on the basis of the previous generation YOLOv4. Compared with YOLOv4, YOLOv5 has a faster training speed, and YOLOv5s has a smaller model size, which is conducive to the rapid deployment of the model. When performing helmet detection, the distance of the shooting scene and the complex construction environment produce a large number of small targets. However, in practical applications, the YOLOv5 model still has certain deficiencies in the detection of small targets. Therefore, the present invention improves the YOLOv5 model to improve the detection accuracy of small targets while ensuring the detection speed. Summary of the invention

[0005] The purpose of the present invention is to provide a helmet wearing detection method based on the YOLOv5 target detection model, aiming to improve the detection capability of small targets in construction scenarios.

[0006] The technical solution adopted by the present invention is: a helmet wearing detection method based on the YOLOv5 target detection model, and the specific operation steps are as follows:

[0007] Step 1, obtain the helmet detection data set SHWD and generate the corresponding label file; divide the data set into a training set, a validation set and a test set; the helmet detection data set includes a helmet image and a helmet label file, and the label file includes a label box, the position information of the label box and the real category information of the content marked by the label box;

[0008] Step 2: Input the input image into the original YOLOv5 target detection model, and extract four feature maps of different sizes through the feature extraction part of the original YOLOv5 target detection model; then perform feature fusion on the four feature maps of different sizes, and finally arrive at the improved YOLOv5 target detection model;

[0009] Step 3, the loss function of the improved YOLOv5 target detection model is divided into a positioning loss function, a classification loss function, and a confidence loss function; the center loss function in metric learning is used as a regularization term to constrain the classification loss and the confidence loss, thereby forming a new objective function; the stochastic gradient descent method is used to optimize the objective function, thereby achieving the purpose of training the improved YOLOv5 target detection model;

[0010] Step 4: Input the training set images in the helmet detection data set into the improved YOLOv5 model in step 3 for training, and save the weight parameters of the improved YOLOv5 model with the highest detection accuracy on the validation set during the training process, and name the weight file best.pt;

[0011] Step 5. Load the weight file best.pt described in step 4 into the improved YOLOv5 target detection model, input the images in the test set into the improved YOLOv5 target detection model, and obtain the detection results of the model in the test set. The detected results include two parts, namely, the safety helmet worn in the construction scene and the normal head without wearing it.

[0012] The present invention is also characterized in that:

[0013] The data set in step 1 is a picture set, which is divided into two categories: head and helmet. Corresponding txt label files are created for pictures in different categories.

[0014] In step 2, the model structure of YOLOv5 is improved so that YOLOv5 can predict the bounding box of the detection target at four scales; for an input image with an input size of 640×640 pixels, the feature extraction part of YOLOv5 extracts feature maps of four sizes: 160×160, 80×80, 40×40, and 20×20; the feature fusion part fuses the feature maps of these four sizes; in the detection part of the YOLOv5 model, a 3D prediction tensor is output at each scale; the information encoded in the prediction tensor includes: the location information of the prediction box, the confidence that the prediction box contains the target to be detected, and the category of the target contained in the prediction box.

[0015] The output target categories are set to be normal heads without helmets and helmets; therefore, in the detection part of the YOLOv5 model, for feature maps with a scale of N×N, the size of the prediction tensor generated by the YOLOv5 model is N×N×[4*(1+4+2)], where 4 represents the tensor predicted by the improved YOLOv5 model at four scales, 1 represents the prediction of confidence, 4 represents the prediction box position information, and 2 represents the target category is two classes. Therefore, the four scale prediction tensors output by the improved YOLOv5 target detection model are 160×160×28, 80×80×28, 40×40×28, and 20×20×28, respectively.

[0016] In step 3, the loss function of the initial YOLOv5 target detection model consists of three parts: positioning loss, confidence loss, and classification loss, expressed as:

[0017]

[0018] Wherein λ1 is the weight coefficient of the positioning loss, λ2 is the weight coefficient of the confidence loss, and λ3 is the weight coefficient of the classification loss; is the positioning loss function, and its expression is: Among them, GIoU represents the degree of overlap between the target box and the real box when locating the target; and They are confidence loss and classification loss respectively. YOLOv5 uses BCELogits(x,y) as and Loss function; assuming that the given data set There are N samples x i and the corresponding label y i , BCELogits(x,y) is:

[0019]

[0020] Where σ(·) is the sigmoid function, which is expressed as For input x i , its output σ(x i ) has a value range of (0,1);

[0021] The present invention adds a central loss function as a constraint after the confidence loss and classification loss of YOLOv5. The expression of this function is: For each sample x i and the corresponding class label y i , Representative class y i After adding the center loss function, the expressions of confidence loss and classification loss are as follows:

[0022]

[0023]

[0024] in, and They represent the confidence loss and classification loss function constrained by the center loss function, and α and β are the corresponding weight coefficients. Finally, the three loss functions are combined to obtain the improved YOLOv5 loss function:

[0025]

[0026] In step 4, the size of the input images in the training set is set to 640×640, and the loss function is quantized using mini-batch stochastic gradient descent with a batch size of 8. For optimization, the initial learning rate is 0.01, the momentum value is set to 0.937, and the weight decay value is set to 5e-4; during the training process, the cosine annealing method is used to reduce the learning rate, thereby accelerating the model to converge to the global optimum;

[0027] The training set is used to train the model, and the validation set is used to provide feedback on the model training status. When the model has the highest detection accuracy on the validation set, the weight parameters of the model are saved to the weight file best.pt. The test set is mainly used for the final evaluation of the trained model.

[0028] In step 5, the non-maximum suppression method is used to screen the detection results in the test set, where the detected helmet is framed by a green frame, and the detected head is framed by a red frame.

[0029] The beneficial effects of the present invention are as follows: the helmet wearing detection method based on the YOLOv5 target detection model of the present invention utilizes the feature extraction part of YOLOv5 to extract a feature map of a larger size on the basis of the initial YOLOv5 target detection model, and performs feature fusion in the feature fusion part.

[0030] The feature map with a larger size has a smaller receptive field, so it is more conducive to the YOLOv5 target detection model to detect small targets. In addition, the present invention also uses the center loss as a regularization term to constrain the confidence loss and classification loss in the YOLOv5 target detection model, so that the improved YOLOv5 target detection model can more accurately detect the category of small targets.

[0031] Experiments have shown that the improved YOLOv5 model proposed in the present invention has excellent performance in the detection task of safety helmets in construction scenarios, and some small targets that cannot be detected by YOLOv5 are detected by the improved YOLOv5. BRIEF DESCRIPTION OF THE DRAWINGS

[0032] Figure 1 It is a flow chart of a helmet wearing detection method based on the YOLOv5 target detection model of the present invention;

[0033] Figure 2 It is a detection result diagram of the helmet wearing detection method based on the YOLOv5 target detection model of the present invention. DETAILED DESCRIPTION

[0034] Reference Figure 1 The specific implementation steps of the helmet wearing detection method based on the YOLOv5 target detection model of the present invention are as follows:

[0035] Step 1: Get the Safety Helmet Wearing Dataset (SHWD). The present invention sets two classes: "head" and "helmet". The dataset contains 7,581 images, including 9,044 images of people wearing helmets and 111,514 images of normal heads.

[0036] Step 1.1. Create a label file corresponding to each class in the dataset. According to the annotation file in the above dataset, convert it into a txt label file in the format required for YOLOv5 training. The dataset is divided into training set, validation set and test set. The training set is used to train the model, and the validation set is used to provide feedback on the training of the model. The weight parameters of the model are saved based on the performance of the model on the validation set, and are used to load during testing and reasoning detection. The test set is mainly used for the final evaluation of the trained model.

[0037] Step 2: For the input image of 640×640, the original YOLOv5 target detection model extracts three feature maps of different sizes through the feature extraction part, which are 80×80, 40×40, and 20×20. The feature map with a larger size has a smaller receptive field, so it can capture more image details, which is conducive to the positioning of the target and also conducive to the model to detect small targets; the feature map with a smaller size has a larger receptive field, so it can extract more semantic information and is more conducive to the classification of targets. When detecting safety helmets in construction scenes, there are many small targets due to the distance of shooting and the complexity of the construction scene. In order to further improve the detection effect of small targets, this scheme improves the multi-scale prediction of YOLOv5 and adopts feature maps of four scales of 160×160, 80×80, 40×40, and 20×20. The target categories set in this paper are "head" without a safety helmet and "safety helmet". Therefore, in the detection part of the YOLOv5 model, for the feature map with a scale of N×N, the size of the prediction tensor generated by the YOLOv5 model is N×N×[4*(1+4+2)], where 4 represents the tensor of four scales predicted by the improved YOLOv5 model, 1 represents the prediction of the confidence, 4 represents the predicted box position information, and 2 represents the target category of two classes. Therefore, the four scales of prediction tensors output by the improved YOLOv5 target detection model are 160×160×28, 80×80×28, 40×40×28 and 20×20×28 respectively.

[0038] Step 3: The central loss function in metric learning is introduced to constrain the confidence loss and classification loss functions in YOLOv5 to obtain a new loss function

[0039] Step 4: Input the training set images in the helmet detection data set into the improved YOLOv5 model in step 3 for training, and save the weight parameters of the improved YOLOv5 model with the highest detection accuracy on the validation set during the training process, and name the weight file best.pt;

[0040] Step 4.1: During specific training, the size of the images in the training set is set to 640×640. The loss function is calibrated using a mini-batch stochastic gradient descent method with a batch size of 8. Optimize. The initial learning rate is 0.01, the momentum value is set to 0.937, and the weight decay value is 5e-4. During the training process, the cosine annealing method is used to reduce the learning rate so that the model is as close to the local minimum as possible.

[0041] Step 4.2. In addition, due to and In It is obtained through learning as a parameter, so the present invention uses the stochastic gradient descent method to learn it, and the initial learning rate is 1e-3.

[0042] Step 4.3: The training process is carried out for 50 rounds in total. The model loads the weight file YOLOv5m.pt of the YOLOv5 model trained on the COCO dataset as the initial weight. In the training process, the weight parameters of the model with the best performance on the validation set are saved as the weight file best.pt, and the weight parameters of the model trained to the end are saved as the weight file last.pt.

[0043] Step 4.4: After training, load the best.pt file with the best weight parameters in the training process into the improved YOLOv5, and then use it to test the helmets in the railway construction scene. During the test, set the image size to 640×640 and the batch size to 32.

[0044] Step 4.5: The training process is performed on a graphics card equipped with a GeForce RTX 2080 Ti. And and The values ​​of weight coefficients α and β are set to α=1e-4 and β=1e-6 respectively.

[0045] Step 4.6. Precision, Recall and mAP are used in the experiment to measure the performance of the model. The specific calculation of Precision and Recall depends on IoU. First, a threshold is set. If the IoU between the predicted box and the true box is greater than the threshold, the prediction is called a "true positive" (TP), otherwise the prediction is called a "false positive" (FP). In addition, the target missed by the model is called a "false negative" (FN). Therefore, Precision measures the percentage of all correct predictions, while Recall is the percentage of all targets in the true box that are predicted as positive. Specifically, the calculation formulas for Precision and Recall are as follows:

[0046]

[0047]

[0048] Step 4.7, based on the precision and recall rate, the average precision AP of each class can be calculated, and the average AP of all categories is mAP. And mAP@.5 means the size of the mAP of the model when the threshold is 0.5. mAP@.5:.95: represents the size of the average mAP when the threshold value is between 0.5 and 0.95 and the step size is 0.05. The test of the improved YOLOv5 target detection model proposed in the present invention is shown in the following table.

[0049] The improved YOLOv5 training results are shown in Table 1 below

[0050] Table 1

[0051]

[0052] Step 5: Load the weight file best.pt described in step 4 into the improved YOLOv5 target detection model, input the images in the test set into the improved YOLOv5 target detection model, and obtain the detection results of the model in the test set. The detected results include two parts, namely, the safety helmet worn in the construction scene and the normal head without wearing (such as Figure 2 shown).

[0053] In summary, according to the test results and the results in the table, it can be seen that the improvement of the present invention is effective. The scheme proposed by the present invention effectively improves the detection accuracy of the safety helmet detection in the railway construction scene. In addition, for the smaller targets in the scene that are difficult to detect and distinguish, the improved YOLOv5 proposed by the present invention can also effectively detect them.

Claims

1. A helmet wearing detection method based on the YOLOv5 target detection model, characterized in that: Step 1: Obtain the helmet detection dataset SHWD and generate the corresponding label file; The dataset is divided into training set, validation set and test set; Step 2: Input the input image into the original YOLOv5 target detection model, and extract four feature maps of different sizes through the feature extraction part of the original YOLOv5 target detection model; then perform feature fusion on the four feature maps of different sizes, and finally arrive at the improved YOLOv5 target detection model; In step 2, the model structure of YOLOv5 is improved so that YOLOv5 can predict the bounding box of the detection target at four scales; for an input image with an input size of 640×640 pixels, the feature extraction part of YOLOv5 extracts feature maps of four sizes: 160×160, 80×80, 40×40, and 20×20; the feature fusion part fuses the feature maps of these four sizes; in the detection part of the YOLOv5 model, a 3D prediction tensor is output at each scale; the information encoded in the prediction tensor includes: the location information of the prediction box, the confidence that the prediction box contains the target to be detected, and the category of the target contained in the prediction box; The output target categories are set to be normal heads without helmets and helmets; therefore, in the detection part of the YOLOv5 model, for feature maps with a scale of N×N, the size of the prediction tensor generated by the YOLOv5 model is N×N×[4*(1+4+2)], where 4 represents the tensors of four scales predicted by the improved YOLOv5 model, 1 represents the prediction of confidence, 4 represents the prediction box position information, and 2 represents the target categories of two classes. Therefore, the four scale prediction tensors output by the improved YOLOv5 target detection model are 160×160×28, 80×80×28, 40×40×28, and 20×20×28 respectively; Step 3, the loss function of the improved YOLOv5 target detection model is divided into a positioning loss function, a classification loss function, and a confidence loss function; the center loss function in metric learning is used as a regularization term to constrain the classification loss and the confidence loss, thereby forming a new objective function; the stochastic gradient descent method is used to optimize the objective function, thereby achieving the purpose of training the improved YOLOv5 target detection model; In step 3, the loss function of the initial YOLOv5 target detection model consists of three parts: positioning loss, confidence loss, and classification loss, expressed as: in is the weight coefficient of the positioning loss, is the weight coefficient of confidence loss, is the weight coefficient of classification loss; is the positioning loss function, and its expression is: in Represents the degree of overlap between the target frame and the real frame when locating the target; and They are confidence loss and classification loss, YOLOv5 uses As and Loss function; assuming that the given data set have Samples and the corresponding labels , for: in is the sigmoid function, and its expression is , for input , whose output The value range of is (0, 1); Add the center loss function as a constraint after the confidence loss and classification loss of YOLOv5. The expression of this function is: , for each sample and the corresponding class labels , Representative Class After adding the center loss function, the expressions of confidence loss and classification loss are as follows: in, and Respectively represent the confidence loss and classification loss function constrained by the center loss function, and is the corresponding weight coefficient; finally, the three parts of the loss function are combined to obtain the improved YOLOv5 loss function: Step 4: Input the training set images in the helmet detection data set into the improved YOLOv5 model in step 3 for training, and save the weight parameters of the improved YOLOv5 model when the detection accuracy rate on the validation set is the highest during the training process, and name the weight file best.pt; Step 5. Load the weight file best.pt described in step 4 into the improved YOLOv5 target detection model, input the images in the test set into the improved YOLOv5 target detection model, and obtain the detection results of the model in the test set. The detected results include two parts, namely, the safety helmet worn in the construction scene and the normal head without wearing it.

2. The helmet wearing detection method based on the YOLOv5 target detection model according to claim 1, characterized in that: In step 1, the data set is a picture set, which is divided into two categories: head and helmet. Corresponding txt label files are created for pictures in different categories.

3. The helmet wearing detection method based on the YOLOv5 target detection model according to claim 1, characterized in that: In step 4, the size of the input images in the training set is set to 640×640, and the loss function is quantized using mini-batch stochastic gradient descent with a batch size of 8. For optimization, the initial learning rate is 0.01, the momentum value is set to 0.937, and the weight decay value is set to 5e-4; during the training process, the cosine annealing method is used to adjust the learning rate to accelerate the model convergence to the global optimum; The training set is used to train the model, and the validation set is used to provide feedback on the model training status. When the model has the highest detection accuracy on the validation set, the weight parameters of the model are saved to the weight file best.pt. The test set is mainly used for the final evaluation of the trained model.

4. The helmet wearing detection method based on the YOLOv5 target detection model according to claim 1, characterized in that: In step 5, the non-maximum suppression method is used to screen the detection results in the test set, where the detected helmet is framed by a green frame, and the detected head is framed by a red frame.

Images