A safety helmet detection method and device based on a YOLOv5 optimization model and a medium

Abstract

The application relates to a safety helmet detection method and device based on a YOLOv5 optimization model and a medium, wherein the safety helmet detection method based on the YOLOv5 optimization model comprises the following steps: obtaining a safety helmet detection image; inputting the safety helmet detection image into a YOLOv5 optimization model as a training sample for training to obtain a safety helmet detection model; wherein the YOLOv5 comprises a main structure, a neck structure and a head structure; the main structure comprises a plurality of CBS modules and a plurality of SwinT modules; each SwinT module is used for performing hierarchical feature extraction on a second feature map output by the safety helmet detection image after being processed by each CBS module; and an employee image collected in a work scene is input into the safety helmet detection model to obtain a safety helmet detection result. The application solves the problem that the existing YOLOv5 algorithm has low detection accuracy for small targets in the prior art, and improves the accuracy of detecting whether an employee wears a safety helmet in a power-related work scene.

Description

Technical Field

[0001] This application relates to the field of target detection, and in particular to a method, apparatus and medium for detecting safety helmets based on the YOLOv5 optimization model. Background Technology

[0002] Currently, with rapid social development, the demand for electricity is becoming increasingly urgent, and the number of jobs related to electricity is increasing. Safety issues should also become a more pressing focus. Safety helmets, as protective equipment for the head, are essential safety gear in electrical work. However, because workshop cameras are usually placed at a high position, a worker's safety helmet in the image is a relatively small target. Furthermore, environmental factors such as weather, lighting, and personnel density place higher demands on the target detection algorithms used.

[0003] Small object detection has always been a key research topic in the field of computer vision object detection. Currently, there are two main types of object detection methods: one is a regression-based single-stage algorithm, and the other is a two-stage algorithm based on candidate regions. The former offers higher real-time performance compared to the single-stage algorithm, but with slightly lower accuracy. YOLO (You Only Look Once) is a two-stage algorithm. YOLOv5, released by Ultralytics in 2020, is an excellent version in the YOLO series. Compared to other versions, it can be applied to a wider range of fields and is more flexible, but its ability to detect small objects is not outstanding, and its accuracy is not high.

[0004] There is currently no effective solution to the problem that the YOLOv5 algorithm has low accuracy in detecting small targets. Summary of the Invention

[0005] Therefore, it is necessary to provide a method, device, and medium for detecting safety helmets based on the YOLOv5 optimized model to address the aforementioned technical problems.

[0006] In a first aspect, embodiments of this application provide a helmet detection method based on a YOLOv5 optimized model, the method comprising:

[0007] Acquire images of the safety helmet;

[0008] The safety helmet detection image is used as a training sample and input into a YOLOv5 optimization model for training to obtain a safety helmet detection model. The YOLOv5 optimization model includes a trunk structure, a neck structure, and a head structure connected in sequence. The trunk structure is used to extract features from the safety helmet detection image to obtain a first feature map. The trunk structure includes multiple CBS modules and multiple SwinT modules. Each SwinT module is used to perform hierarchical feature extraction on the second feature map output from the CBS module. The neck structure is used to perform multi-scale feature fusion on the first feature map to obtain a third feature map. The head structure is used to make predictions based on the third feature map to obtain the detection result.

[0009] The images of employees collected in the work scene are input into the safety helmet detection model to obtain the safety helmet detection results.

[0010] In one embodiment, each of the SwinT modules includes a normalization layer, a window-based multi-head self-attention layer, a moving window-based multi-head self-attention layer, and a multilayer perceptron; the SwinT module is specifically used for:

[0011] The second feature map output after the safety helmet detection image is processed by the CBS module is passed through the normalization layer, the multi-head self-attention layer of the window, the normalization layer, and the multilayer perceptron to obtain the fourth feature map.

[0012] The fourth feature map is passed through the normalization layer, the multi-head self-attention layer of the moving window, the normalization layer, and the multilayer perceptron to obtain the fifth feature map.

[0013] In one embodiment, the plurality of SwinT modules include a first SwinT module, a second SwinT module, a third SwinT module, and a fourth SwinT module. The backbone structure further includes an SE module, which is connected to the output of the third SwinT module. The SE module is specifically used for:

[0014] The sixth feature map output after the safety helmet detection image is processed by the third SwinT module is subjected to global average pooling to output a one-dimensional vector.

[0015] Based on the one-dimensional vector, the weight value is calculated through an activation layer consisting of two fully connected layers;

[0016] The weight value is multiplied by the pixel value of the sixth feature map to obtain the output result.

[0017] In one embodiment, the backbone structure further includes a convolutional block attention module connected to the output of the fourth SwinT module. The convolutional block attention module includes a spatial attention module and a channel attention module. The channel attention module is used to adaptively correct the seventh feature map output by the safety helmet detection image after processing by the fourth SwinT module to generate an eighth feature map. The spatial attention module is used to correct the eighth feature map to output a ninth feature map.

[0018] In one embodiment, the channel attention module is used to adaptively correct the seventh feature map output by the fourth SwinT module after the safety helmet detection image is processed to generate an eighth feature map. The spatial attention module is used to correct the eighth feature map to output a ninth feature map in the following calculation method:

[0019]

[0020]

[0021]

[0022]

[0023]

[0024] in For performing convolution operations; F max C This is the channel max pooling feature map; MLP is a multilayer perceptron; MaxPool() is the max pooling function; FavgC is the channel average pooling feature map; AvePool() is the average pooling function. It is the Sigmoid activation function; f() is a standard convolutional layer; W C and W S These are channel attention weights and spatial attention weights, respectively.

[0025] In one embodiment, the step of inputting the safety helmet detection image as a training sample into the YOLOv5 optimization model for training to obtain the safety helmet detection model includes:

[0026] The safety helmet detection image is augmented using a mosaic method.

[0027] The data-augmented helmet detection images are input into the YOLOv5 optimized model for training, resulting in a helmet detection model.

[0028] In one embodiment, the step of inputting the safety helmet detection image as a training sample into the YOLOv5 optimization model for training to obtain the safety helmet detection model includes:

[0029] The data-augmented safety helmet detection image is input into the YOLOv5 optimized model for training, and the training results are obtained.

[0030] Based on the training results, the model parameters are adjusted for iterative optimization to obtain the safety helmet detection model.

[0031] Secondly, embodiments of this application also provide a helmet detection device based on optimized YOLOv5, the device comprising:

[0032] The acquisition module is used to acquire images of the safety helmet detection.

[0033] The training module is used to input the safety helmet detection image as training samples into the YOLOv5 optimization model for training, thereby obtaining a safety helmet detection model. The YOLOv5 optimization model includes a trunk structure, a neck structure, and a head structure connected in sequence. The trunk structure is used to extract features from the safety helmet detection image to obtain a first feature map. The trunk structure includes multiple CBS modules and multiple SwinT modules. Each SwinT module is used to perform hierarchical feature extraction on the second feature map output from the CBS module. The neck structure is used to perform multi-scale feature fusion on the first feature map to obtain a third feature map. The head structure is used to make predictions based on the third feature map to obtain the detection result.

[0034] The detection module is used to input employee images collected in the work scene into the safety helmet detection model to obtain safety helmet detection results.

[0035] Thirdly, embodiments of this application also provide a computer device, including a memory and a processor, wherein the memory stores a computer program, and the processor is configured to run the computer program to perform the method described in the first aspect above.

[0036] Fourthly, embodiments of this application also provide a storage medium storing a computer program, wherein the computer program, when executed by a processor, implements the method described in the first aspect above.

[0037] The aforementioned helmet detection method, apparatus, and medium based on a YOLOv5 optimized model acquire helmet detection images; these images are then used as training samples to train a YOLOv5 optimized model, resulting in a helmet detection model. The YOLOv5 optimized model includes a trunk structure, a neck structure, and a head structure connected sequentially. The trunk structure extracts features from the helmet detection images to obtain a first feature map. This trunk structure includes multiple CBS modules and multiple SwinT modules. Each SwinT module performs hierarchical feature extraction on the second feature map output from the CBS modules. The neck structure performs multi-scale feature fusion on the first feature map to obtain a third feature map. The head structure performs prediction based on the third feature map to obtain a detection result. Employee images collected in the work scene are input into the helmet detection model to obtain helmet detection results. This method solves the problem of low accuracy in detecting small targets using existing YOLOv5 algorithms and improves the accuracy of detecting whether employees are wearing helmets in power-related work scenarios.

[0038] Details of one or more embodiments of this application are set forth in the following drawings and description to make other features, objects and advantages of this application more readily apparent. Attached Figure Description

[0039] The accompanying drawings, which are included to provide a further understanding of this application and form part of this application, illustrate exemplary embodiments and are used to explain this application, but do not constitute an undue limitation of this application. In the drawings:

[0040] Figure 1 This is a schematic diagram of the application environment of a helmet detection method based on the YOLOv5 optimized model in one embodiment;

[0041] Figure 2 This is a flowchart illustrating a helmet detection method based on a YOLOv5 optimized model in one embodiment;

[0042] Figure 3 This is a diagram of an existing YOLOv5 architecture in one embodiment;

[0043] Figure 4 This is a computational structure diagram of self-attention in one embodiment;

[0044] Figure 5 This is a structural diagram of the SwinT module in one embodiment;

[0045] Figure 6 This is a structural diagram of the SE module in one embodiment;

[0046] Figure 7 This is a structural diagram of the convolutional block attention module in one embodiment;

[0047] Figure 8 This is a schematic diagram of an optimized YOLOv5 structure in a preferred embodiment;

[0048] Figure 9 This is a schematic diagram of the visualization results of helmet detection in one embodiment;

[0049] Figure 10 This is a schematic diagram of the PR curves of the YOLOv5 model before and after optimization in one embodiment;

[0050] Figure 11 This is a structural block diagram of a helmet detection device based on a YOLOv5 optimized model in one embodiment;

[0051] Figure 12 This is a schematic diagram of the computer device structure in one embodiment. Detailed Implementation

[0052] To make the objectives, technical solutions, and advantages of this application clearer, the application is described and illustrated below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the scope of this application. All other embodiments obtained by those skilled in the art based on the embodiments provided in this application without inventive effort are within the scope of protection of this application.

[0053] Obviously, the accompanying drawings described below are merely some examples or embodiments of this application. Those skilled in the art can apply this application to other similar scenarios based on these drawings without any inventive effort. Furthermore, it is understood that although the efforts made in this development process may be complex and lengthy, for those skilled in the art related to the content disclosed in this application, any changes to design, manufacturing, or production based on the technical content disclosed in this application are merely conventional technical means and should not be construed as insufficient disclosure of the content of this application.

[0054] In this application, the reference to "embodiment" means that a specific feature, structure, or characteristic described in connection with an embodiment may be included in at least one embodiment of this application. The appearance of this phrase in various places in the specification does not necessarily refer to the same embodiment, nor is it a separate or alternative embodiment that is mutually exclusive with other embodiments. It will be explicitly and implicitly understood by those skilled in the art that the embodiments described in this application may be combined with other embodiments without conflict.

[0055] Unless otherwise defined, the technical or scientific terms used in this application shall have the ordinary meaning understood by one of ordinary skill in the art to which this application pertains. The terms “a,” “an,” “an,” “the,” and similar words used in this application do not indicate quantity limitation and may indicate singular or plural. The terms “comprising,” “including,” “having,” and any variations thereof used in this application are intended to cover non-exclusive inclusion; for example, a process, method, system, product, or device that includes a series of steps or modules (units) is not limited to the listed steps or units, but may also include steps or units not listed, or may include other steps or units inherent to these processes, methods, products, or devices. The terms “connected,” “linked,” “coupled,” and similar words used in this application are not limited to physical or mechanical connections, but may include electrical connections, whether direct or indirect. “Multiple” used in this application refers to two or more. “And / or” describes the relationship between related objects, indicating that three relationships may exist; for example, “A and / or B” can represent: A alone, A and B simultaneously, and B alone. The character " / " generally indicates that the preceding and following objects are in an "or" relationship. The terms "first," "second," and "third" used in this application are merely to distinguish similar objects and do not represent a specific ordering of the objects.

[0056] The method embodiments provided in this example can be executed on a terminal, computer, or similar computing device. For example, it can run on a terminal. Figure 1 This is a hardware structure block diagram of the terminal for the helmet detection method based on the YOLOv5 optimized model in this embodiment. Figure 1 As shown, a terminal may include one or more ( Figure 1 Only one is shown in the diagram. A processor 102 and a memory 104 for storing data are also included. The processor 102 may be, but is not limited to, a microprocessor (MCU) or a programmable logic device (FPGA). The terminal may also include a transmission device 106 for communication functions and an input / output device 108. Those skilled in the art will understand that… Figure 1 The structure shown is for illustrative purposes only and does not limit the structure of the terminal described above. For example, the terminal may also include components that are larger than... Figure 1 The more or fewer components shown, or having the same Figure 1 The different configurations shown are illustrated.

[0057] The memory 104 can be used to store computer programs, such as application software programs and modules, like the computer program corresponding to the helmet detection method based on the YOLOv5 optimization model in this embodiment. The processor 102 executes various functional applications and data processing by running the computer programs stored in the memory 104, thereby implementing the above-described method. The memory 104 may include high-speed random access memory, and may also include non-volatile memory, such as one or more magnetic storage devices, flash memory, or other non-volatile solid-state memory. In some instances, the memory 104 may further include memory remotely located relative to the processor 102, and these remote memories can be connected to the terminal via a network. Examples of such networks include, but are not limited to, the Internet, corporate intranets, local area networks, mobile communication networks, and combinations thereof.

[0058] The transmission device 106 is used to receive or send data via a network. This network includes a wireless network provided by the terminal's communication provider. In one example, the transmission device 106 includes a Network Interface Controller (NIC), which can connect to other network devices via a base station to communicate with the Internet. In another example, the transmission device 106 can be a Radio Frequency (RF) module used for wireless communication with the Internet.

[0059] This application provides a helmet detection method based on a YOLOv5 optimized model, such as... Figure 2 As shown, the method includes the following steps:

[0060] Step S201: Obtain the safety helmet detection image;

[0061] Specifically, the helmet detection images obtained in this embodiment are from the open-source SHWD (Helmet Wearing Dataset). SHWD includes 7,581 images with different scenes, weather conditions, lighting conditions, number of people, and shooting distances. The images are divided into two categories: hat category (people wearing helmets) and people category (people not wearing helmets).

[0062] The software and hardware environment used in the experiment was as follows: Windows 11 Professional operating system, PyTorch version 1.10, CPU (Central Processing Unit) Intel Core(R) I510400F, RAM (Random Access Memory) 32GB, and GPU (Central Processing Unit) NVIDIA GeForce RTX 3060 12GB.

[0063] Step S202: The safety helmet detection image is used as a training sample and input into the YOLOv5 optimization model for training to obtain a safety helmet detection model. The YOLOv5 optimization model includes a trunk structure, a neck structure, and a head structure connected in sequence. The trunk structure is used to extract features from the safety helmet detection image to obtain a first feature map. The trunk structure includes multiple CBS modules and multiple SwinT modules. Each SwinT module is used to perform hierarchical feature extraction on the second feature map output after processing by the CBS modules. The neck structure is used to perform multi-scale feature fusion on the first feature map to obtain a third feature map. The head structure is used to make predictions based on the third feature map to obtain the detection result.

[0064] Specifically, Figure 3 The existing YOLOv5 architecture consists of a Backbone, a Neck, and a Head. The Backbone includes the CSP-Darknet 53 (C3) module, optimized using the CSP (Cross Stage Partial) method, significantly reducing computational cost. Following CSP-Darknet 53 is an SPPF (Spatial Pyramid Pooling-Fast) module. The Neck in YOLOv5 is a combination of FPN (Feature Pyramid Network) and PAN (Path Aggregation Network), capable of considering both large and small target features simultaneously. The Head receives the processed Neck features, generates the final detection result, and corrects the candidate box positions based on positional offsets for more accurate detection. Figure 3 In the diagram, Input represents the input image, and □×□×□ represents the pixel size of the input image or feature map; the CBS (Conv Batch normalization SiLU) module is a combination of convolution, batch normalization, and the SiLU image activation function; Bottleneck CSP (Cross Stage Partial) is the cross-local stage of the spine; UpSample is an upsampling layer, a neural network layer used to transform an input feature map into a larger feature map; Concat is a connection layer used to concatenate two or more tensors along a certain dimension; and Conv is a convolutional layer used for feature extraction.

[0065] This embodiment improves upon the existing YOLOv5 architecture. While the CSP Darknet53 (C3) module in the backbone of the existing YOLOv5 architecture has lightweight structure and computational complexity, its feature extraction capability for small targets is relatively weak. Therefore, this paper reduces the reuse of the C3 module by trimming and improving the existing C3 module in the YOLOv5 backbone architecture, replacing all C3 modules with SwinT (Swin Transformer) modules.

[0066] This embodiment proposes the SwinT module based on the application of Transformer in image processing. Since its inception, Transformer has been a dominant model in Natural Language Processing (NLP) tasks. It primarily establishes long-range correlations between pixels through a self-attention mechanism, exhibiting strong contextual feature extraction capabilities. Three matrices, Q, K, and V, can be obtained by learning the self-attention mechanism of each element in the input to determine the relative importance of each element. Figure 4 This is a computational structure diagram for self-attention, where Q (Query) represents the query, K (Key) represents the key, and V (Value) represents the value; DotProduct represents the dot product, which performs a dot product calculation on Q and K; Scale represents the weights, used to perform a feature dot product calculation on the normalized attention weights and the previous-level input; SoftMax is a normalized exponential function that transforms the output classification result into the corresponding probability; MatMul is a matrix multiplication function used to perform a tensor matrix multiplication on the previous-level input and the V value. Unlike Natural Language Processing (NLP), visual elements in computer vision are not fixed in scale; an image contains far more pixels than a piece of text. Therefore, directly applying Transformer to image processing is quite challenging.

[0067] The SwinT module introduces the hierarchical and local properties of convolution operations, performing attention computations across various window regions and also performing self-attention computations locally within non-overlapping windows. This limits the number of patches contained in each window and significantly reduces computational complexity. Furthermore, the SwinT module gradually merges adjacent blocks starting with smaller, shallower patches, ultimately obtaining a hierarchical feature map structure, making it more suitable for small object detection.

[0068] Step S203: Input the employee images collected in the work scene into the safety helmet detection model to obtain the safety helmet detection results.

[0069] Steps S201 to S203 above involve acquiring a safety helmet detection image; inputting the safety helmet detection image as a training sample into a YOLOv5 optimization model for training to obtain a safety helmet detection model; wherein, the YOLOv5 optimization model includes a trunk structure, a neck structure, and a head structure connected in sequence; the trunk structure is used to extract features from the safety helmet detection image to obtain a first feature map, wherein the trunk structure includes multiple CBS modules and multiple SwinT modules, each SwinT module is used to perform hierarchical feature extraction on the second feature map output by the safety helmet detection image after processing by each CBS module; the neck structure is used to perform multi-scale feature fusion on the first feature map to obtain a third feature map, and the head structure is used to make predictions based on the third feature map to obtain detection results; and employee images collected in the work scene are input into the safety helmet detection model to obtain safety helmet detection results. This solves the problem of low accuracy of existing YOLOv5 algorithms for detecting small targets in related technologies and improves the accuracy of detecting whether employees are wearing safety helmets in power-related work scenarios.

[0070] In one embodiment, each of the SwinT modules includes a normalization layer, a window-based multi-head self-attention layer, a moving window-based multi-head self-attention layer, and a multilayer perceptron.

[0071] Figure 5 The SwinT module structure diagram is shown. Figure 5 LayerNorm is a normalization layer that normalizes all features for each sample; W-MSA is a window-based multi-head self-attention layer, and SW-MSA is a shift-window-based multi-head self-attention layer; MLP is a multilayer perceptron; Z l-1 This represents the output feature of the MLP module in the (l-1)th block. and Z l These represent the output characteristics of the W-MSA module and the MLP module of the l-th block, respectively. and Z l+1 These represent the output characteristics of the SW-MSA module and the MLP module of the (l+1)th block, respectively.

[0072] The SwinT module in this embodiment is specifically used for:

[0073] The second feature map output after the safety helmet detection image is processed by the CBS module is passed through the normalization layer, the multi-head self-attention layer of the window, the normalization layer, and the multilayer perceptron to obtain the fourth feature map.

[0074] The fourth feature map is passed through the normalization layer, the multi-head self-attention layer of the moving window, the normalization layer, and the multilayer perceptron to obtain the fifth feature map.

[0075] In one embodiment, the plurality of SwinT modules include a first SwinT module, a second SwinT module, a third SwinT module, and a fourth SwinT module, and the backbone structure further includes an SE module, which is connected to the output terminal of the third SwinT module.

[0076] Figure 6 The diagram illustrates the structure of the SE (Squeeze and Excitation) module, a feature recalibration method primarily composed of compression and excitation components. The input image to the SE module first passes through a squeezing layer consisting of global average pooling, outputting a one-dimensional vector. Then, it passes through an excitation layer consisting of two fully connected layers to compute weights. These weights are multiplied by the pixel values ​​of the input feature map to obtain the output. The SE module can allocate larger weights to effective information in the input feature map while assigning smaller weights to ineffective information to enhance expressive power. Figure 6 In this context, Global Average Pool represents the global average pooling operation; FC (Full Connected layer) is a fully connected layer that acts as a classifier; ReLU and Sigmoid are both activation functions; and Scale represents the channel weight multiplication operation.

[0077] The SE module in this embodiment is specifically used for:

[0078] The sixth feature map output after the safety helmet detection image is processed by the third SwinT module is subjected to global average pooling to output a one-dimensional vector.

[0079] Based on the one-dimensional vector, the weight value is calculated through an activation layer consisting of two fully connected layers;

[0080] The weight value is multiplied by the pixel value of the sixth feature map to obtain the output result.

[0081] In one embodiment, the backbone structure further includes a convolutional block attention module connected to the output of the fourth SwinT module. The convolutional block attention module includes a spatial attention module and a channel attention module. The channel attention module is used to adaptively correct the seventh feature map output by the safety helmet detection image after processing by the fourth SwinT module to generate an eighth feature map. The spatial attention module is used to correct the eighth feature map to output a ninth feature map.

[0082] Figure 7The diagram shows the structure of the Convolutional Block Attention Module (CBAM). The CBAM module consists of two parts: (a) the Spatial Attention Module (SAM) and (b) the Channel Attention Module (CAM). These two parts are combined sequentially to serialize the attention feature map information in both the channel and spatial dimensions. When the feature map F is input into the CBAM module, it is first adaptively corrected by the channel attention module to generate the feature map F', and then corrected by the spatial attention module to output the feature map F''. Figure a shows: Feature F' is the feature map generated by the channel attention module; Channel-refined is the channel correction used to process the input feature map F'; [MaxPool, AvgPool] represents global max pooling and global average pooling, which are performed sequentially; Conv Layer represents a convolutional layer; Spatial Attention represents spatial attention; Ms is the final generated spatial attention feature. Figure b shows: Input feature F is the input feature map; MaxPool and AvgPool represent global max pooling and global average pooling, respectively, which are performed in parallel; Shared MLP is a two-layer multilayer perceptron; Channel Attention represents channel attention; Mc is the final generated channel attention feature.

[0083] In one embodiment, the channel attention module is used to adaptively correct the seventh feature map output by the fourth SwinT module after the safety helmet detection image is processed to generate an eighth feature map. The spatial attention module is used to correct the eighth feature map to output a ninth feature map in the following calculation method:

[0084]

[0085]

[0086]

[0087]

[0088]

[0089] in For performing convolution operations; F max C This is the channel max-pooling feature map; MLP is a multilayer perceptron; MaxPool() is the max-pooling function; FavgC is the channel average-pooling feature map; AvePool() is the average-pooling function. It is the Sigmoid activation function; f() is a standard convolutional layer; W Cand W S These are channel attention weights and spatial attention weights, respectively.

[0090] In one embodiment, the step of inputting the safety helmet detection image as a training sample into the YOLOv5 optimization model for training to obtain the safety helmet detection model includes:

[0091] The safety helmet detection image is augmented using a mosaic method.

[0092] The data-augmented helmet detection images are input into the YOLOv5 optimized model for training, resulting in a helmet detection model.

[0093] This embodiment uses a mosaic method to perform data augmentation processing on the safety helmet detection image to enrich sample information and enhance the robustness of the network.

[0094] In one embodiment, the step of inputting the safety helmet detection image as a training sample into the YOLOv5 optimization model for training to obtain the safety helmet detection model includes:

[0095] The data-augmented safety helmet detection image is input into the YOLOv5 optimized model for training, and the training results are obtained.

[0096] Based on the training results, the model parameters are adjusted for iterative optimization to obtain the safety helmet detection model.

[0097] The following describes and illustrates this embodiment through a preferred example. In this embodiment, the C3 module in the existing YOLOv5 backbone structure is pruned and improved, replacing all C3 modules with Swing modules. To further enhance its feature extraction capability, an SE module is added after the third Swing module. A CBAM module is added after the fourth Swing module to further improve the recognition accuracy of the entire network, generating the final optimized YOLOv5 structure. The final optimized YOLOv5 structure is as follows: Figure 8 As shown in Table 1, the optimized YOLOv5 main chain structure is illustrated.

[0098] Table 1

[0099]

[0100] In Table 1, From = -1 indicates that the current module is connected to the previous module; N indicates the number of times this module is repeated; Module indicates the name of the module.

[0101] The model parameters in this embodiment are set as follows: initial learning rate of 0.01, momentum of 0.937, learning rate reduction to 0.0005, epoch of 100, and batch size of 8. An evaluation metric is established to assess the performance of the YOLOv5 optimized model before and after optimization, using the mean accuracy (mAP) with an IoU (Intersection over Union) threshold of 0.5. IoU is a metric used in machine learning to evaluate model performance. A higher mAP value indicates higher model detection accuracy. The calculation process for mAP is as follows:

[0102]

[0103] Where C represents the number of categories; P rec It is the probability of correctly detecting a safety helmet, also known as accuracy; R rec Precision is the probability of correct identification, also known as recall. The formulas for precision and recall are as follows:

[0104]

[0105]

[0106] Where TP represents the number of correctly predicted samples; FP represents the number of samples predicted to be wearing a helmet but not wearing one; and FN represents the number of samples predicted to be wearing a helmet.

[0107] This invention is compared to the existing YOLOv5, trained under the same conditions, and the detection results of the two structures are compared. The visualization results of the helmet detection in this embodiment are as follows: Figure 9 As shown, hat[] represents the probability of identifying someone wearing a helmet; person[] represents the probability of identifying someone not wearing a helmet; the probability values ​​range from 0 to 1. Compared with the existing YOLOv5 model, the improved method of this invention has higher confidence and lower false detection rate in most cases. The PR curve is one of the important means of evaluating object detection models. The higher the accuracy, the lower the recall. Therefore, the closer the PR curve is to the right, the higher the accuracy and recall, and the better the detection effect of the model. The PR curves of the YOLOv5 model before and after optimization are shown in the figure. Figure 10 As shown in (a) and (b), it can be seen that the method proposed in this embodiment is closer to the upper right corner of the PR curve of the safety helmet sample, and its detection performance is better than the existing YOLOv5 model. Figure 10In the table, `all classed[]mAP@0.5` represents the probability of identifying all classes with average accuracy when using an IoU threshold of 0.5; the horizontal axis represents recall, and the vertical axis represents precision; the probability values ​​range from 0 to 1. Table 2 shows a comparison of the detection results between the existing YOLOv5 architecture and the optimized YOLOv5 architecture.

[0108] Table 2

[0109]

[0110] In Table 2, P represents the probability of correct detection; mAP50 represents the average accuracy (mAP) at an IoU threshold of 0.5; and mAP50-95 represents the average accuracy (mAP) at an IoU threshold of 0.5-0.95. As shown in Table 2, compared to the existing YOLOv5 model, the optimized YOLOv5 model in this embodiment improves the accuracy by 1.1%, the mAP50 value by 3.0%, and the mAP50-95 value by 4.3%. This demonstrates that the proposed method can achieve better helmet detection results.

[0111] To examine the impact of each module in the algorithm on the experimental results, three sets of comparative schemes (A, B, and C) were designed for ablation experiments. The results are shown in Table 3, where 1 indicates that the module was used in the experiment, and 0 indicates that the corresponding module was not used in the experiment. Three indicators, P, mAP 50, and mAP 50 -95, were also used for comparison in the ablation experiments.

[0112] Table 3

[0113]

[0114] As shown in Table 3, compared with the existing YOLOv5, the detection accuracy of safety helmets can be improved simply by adding the SwinT module or the SE and CBAM modules to the Backbone. The best results are achieved when both are performed simultaneously, demonstrating the effectiveness of the method proposed in this invention.

[0115] Secondly, embodiments of this application also provide a helmet detection device based on optimized YOLOv5, such as... Figure 11 As shown, the device includes:

[0116] The acquisition module 110 is used to acquire the safety helmet detection image;

[0117] Training module 120 is used to input the safety helmet detection image as a training sample into the YOLOv5 optimization model for training to obtain a safety helmet detection model. The YOLOv5 optimization model includes a trunk structure, a neck structure, and a head structure connected in sequence. The trunk structure is used to extract features from the safety helmet detection image to obtain a first feature map. The trunk structure includes multiple CBS modules and multiple SwinT modules. Each SwinT module is used to perform hierarchical feature extraction on the second feature map output from the safety helmet detection image after processing by the respective CBS module. The neck structure is used to perform multi-scale feature fusion on the first feature map to obtain a third feature map. The head structure is used to make predictions based on the third feature map to obtain a detection result.

[0118] The detection module 130 is used to input employee images collected in the work scene into the safety helmet detection model to obtain safety helmet detection results.

[0119] In one embodiment, each of the SwinT modules includes a normalization layer, a window-based multi-head self-attention layer, a moving window-based multi-head self-attention layer, and a multilayer perceptron; the SwinT module is specifically used for:

[0120] The second feature map output after the safety helmet detection image is processed by the CBS module is passed through the normalization layer, the multi-head self-attention layer of the window, the normalization layer, and the multilayer perceptron to obtain the fourth feature map.

[0121] The fourth feature map is passed through the normalization layer, the multi-head self-attention layer of the moving window, the normalization layer, and the multilayer perceptron to obtain the fifth feature map.

[0122] In one embodiment, the plurality of SwinT modules include a first SwinT module, a second SwinT module, a third SwinT module, and a fourth SwinT module. The backbone structure further includes an SE module, which is connected to the output of the third SwinT module. The SE module is specifically used for:

[0123] The sixth feature map output after the safety helmet detection image is processed by the third SwinT module is subjected to global average pooling to output a one-dimensional vector.

[0124] Based on the one-dimensional vector, the weight value is calculated through an activation layer consisting of two fully connected layers;

[0125] The weight value is multiplied by the pixel value of the sixth feature map to obtain the output result.

[0126] In one embodiment, the backbone structure further includes a convolutional block attention module connected to the output of the fourth SwinT module. The convolutional block attention module includes a spatial attention module and a channel attention module. The channel attention module is used to adaptively correct the seventh feature map output by the safety helmet detection image after processing by the fourth SwinT module to generate an eighth feature map. The spatial attention module is used to correct the eighth feature map to output a ninth feature map.

[0127] In one embodiment, the channel attention module is used to adaptively correct the seventh feature map output by the fourth SwinT module after the safety helmet detection image is processed to generate an eighth feature map. The spatial attention module is used to correct the eighth feature map to output a ninth feature map in the following calculation method:

[0128]

[0129]

[0130]

[0131]

[0132]

[0133]

[0134] in For performing convolution operations; F max C This is the channel max pooling feature map; MLP is a multilayer perceptron; MaxPool() is the max pooling function; FavgC is the channel average pooling feature map; AvePool() is the average pooling function. It is the Sigmoid activation function; f() is a standard convolutional layer; W C and W S These are channel attention weights and spatial attention weights, respectively.

[0135] In one embodiment, the training module 120 is further configured to:

[0136] The safety helmet detection image is augmented using a mosaic method.

[0137] The data-augmented helmet detection images are input into the YOLOv5 optimized model for training, resulting in a helmet detection model.

[0138] In one embodiment, the training module 120 is further configured to:

[0139] The data-augmented safety helmet detection image is input into the YOLOv5 optimized model for training, and the training results are obtained.

[0140] Based on the training results, the model parameters are adjusted for iterative optimization to obtain the safety helmet detection model.

[0141] It should be noted that the above modules can be functional modules or program modules, and can be implemented through software or hardware. For modules implemented through hardware, the above modules can reside in the same processor; or the above modules can be located in different processors in any combination.

[0142] In one embodiment, a computer device is provided, which may be a terminal, and its internal structure diagram may be as follows: Figure 12 As shown, the computer device includes a processor, memory, communication interface, display screen, and input devices connected via a system bus. The processor provides computing and control capabilities. The memory includes non-volatile storage media and internal memory. The non-volatile storage media stores the operating system and computer programs. The internal memory provides an environment for the operation of the operating system and computer programs stored in the non-volatile storage media. The communication interface is used for wired or wireless communication with external terminals; wireless communication can be achieved through Wi-Fi, mobile cellular networks, NFC (Near Field Communication), or other technologies. When executed by the processor, the computer program implements a helmet detection method based on the YOLOv5 optimized model. The display screen can be an LCD screen or an e-ink display screen. The input devices can be a touch layer covering the display screen, buttons, a trackball, or a touchpad on the computer device's casing, or an external keyboard, touchpad, or mouse.

[0143] Those skilled in the art will understand that Figure 12 The structure shown is merely a block diagram of a portion of the structure related to the present application and does not constitute a limitation on the computer device to which the present application is applied. Specific computer devices may include more or fewer components than those shown in the figure, or combine certain components, or have different component arrangements.

[0144] In one embodiment, a computer-readable storage medium is provided, on which a computer program is stored, which, when executed by a processor, implements the steps in any of the above embodiments of the helmet detection method based on the YOLOv5 optimization model.

[0145] Those skilled in the art will understand that all or part of the processes in the methods of the above embodiments can be implemented by a computer program instructing related hardware. The computer program can be stored in a non-volatile computer-readable storage medium, and when executed, it can include the processes of the embodiments of the methods described above. Any references to memory, storage, databases, or other media used in the embodiments provided in this application can include at least one of non-volatile and volatile memory. Non-volatile memory can include read-only memory (ROM), magnetic tape, floppy disk, flash memory, or optical storage, etc. Volatile memory can include random access memory (RAM) or external cache memory. By way of illustration and not limitation, RAM can be in various forms, such as static random access memory (SRAM) or dynamic random access memory (DRAM), etc.

[0146] The technical features of the above embodiments can be combined in any way. For the sake of brevity, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.

[0147] The embodiments described above are merely illustrative of several implementation methods of this application, and while the descriptions are relatively specific and detailed, they should not be construed as limiting the scope of the invention patent. It should be noted that those skilled in the art can make various modifications and improvements without departing from the concept of this application, and these all fall within the protection scope of this application. Therefore, the protection scope of this patent application should be determined by the appended claims.

Claims

1. A helmet detection method based on a YOLOv5 optimized model, characterized in that, The method includes: Acquire images of the safety helmet; The safety helmet detection image is used as a training sample and input into a YOLOv5 optimization model for training to obtain a safety helmet detection model. The YOLOv5 optimization model includes a trunk structure, a neck structure, and a head structure connected sequentially. The trunk structure is used to extract features from the safety helmet detection image to obtain a first feature map. The trunk structure includes multiple CBS modules and multiple SwinT modules. Each SwinT module is used to perform hierarchical feature extraction on the second feature map output from the CBS module. The neck structure is used to perform multi-scale feature fusion on the first feature map to obtain a third feature map. The head structure is used to make predictions based on the third feature map to obtain the detection result. The plurality of SwinT modules include a first SwinT module, a second SwinT module, a third SwinT module, and a fourth SwinT module. The backbone structure also includes an SE module, which is connected to the output of the third SwinT module. The SE module is specifically used for: performing global average pooling on the sixth feature map output from the safety helmet detection image after processing by the third SwinT module to output a one-dimensional vector; calculating weight values ​​based on the one-dimensional vector through an activation layer composed of two fully connected layers; and multiplying the weight values ​​by the pixel values ​​of the sixth feature map to obtain the output result. The backbone structure also includes a convolutional block attention module, which is connected to the output of the fourth SwinT module. The convolutional block attention module includes a spatial attention module and a channel attention module. The channel attention module is used to adaptively correct the seventh feature map output by the safety helmet detection image after processing by the fourth SwinT module to generate an eighth feature map. The spatial attention module is used to correct the eighth feature map to output a ninth feature map. The channel attention module is used to adaptively correct the seventh feature map output by the fourth SwinT module after the safety helmet detection image is processed to generate the eighth feature map. The spatial attention module is used to correct the eighth feature map to output the ninth feature map in the following calculation method: in For performing convolution operations; F represents the seventh feature map, F' represents the eighth feature map; F'' represents the ninth feature map; F max C It is the channel max-pooling feature map; MLP is a multilayer perceptron; MaxPool() is the max-pooling function; F avg C This is the channel average pooling feature map; AvePool() is the average pooling function; It is the Sigmoid activation function; f() is the standard convolutional layer; W C and W S These are channel attention weights and spatial attention weights, respectively. The images of employees collected in the work scene are input into the safety helmet detection model to obtain the safety helmet detection results.

2. The method according to claim 1, characterized in that, Each of the described SwinT modules includes a normalization layer, a window-based multi-head self-attention layer, a moving window-based multi-head self-attention layer, and a multilayer perceptron; the SwinT module is specifically used for: The second feature map output after the safety helmet detection image is processed by the CBS module is passed through the normalization layer, the multi-head self-attention layer of the window, the normalization layer, and the multilayer perceptron to obtain the fourth feature map. The fourth feature map is passed through the normalization layer, the multi-head self-attention layer of the moving window, the normalization layer, and the multilayer perceptron to obtain the fifth feature map.

3. The method according to claim 1, characterized in that, The step of using the detected helmet images as training samples to train the YOLOv5 optimization model includes: The safety helmet detection image is augmented using a mosaic method. The data-augmented helmet detection images are input into the YOLOv5 optimized model for training, resulting in a helmet detection model.

4. The method according to claim 3, characterized in that, The step of using the detected helmet images as training samples to train the YOLOv5 optimization model includes: The data-augmented safety helmet detection image is input into the YOLOv5 optimized model for training, and the training results are obtained. Based on the training results, the model parameters are adjusted for iterative optimization to obtain the safety helmet detection model.

5. A helmet detection device based on optimized YOLOv5, characterized in that, The device includes: The acquisition module is used to acquire images of the safety helmet detection. A training module is used to input the safety helmet detection image as training samples into a YOLOv5 optimization model for training, thereby obtaining a safety helmet detection model. The YOLOv5 optimization model includes a trunk structure, a neck structure, and a head structure connected sequentially. The trunk structure is used to extract features from the safety helmet detection image to obtain a first feature map. The trunk structure includes multiple CBS modules and multiple SwinT modules. Each SwinT module is used to perform hierarchical feature extraction on the second feature map output from the CBS module. The neck structure is used to perform multi-scale feature fusion on the first feature map to obtain a third feature map. The head structure is used to make predictions based on the third feature map to obtain a detection result. The plurality of SwinT modules include a first SwinT module, a second SwinT module, a third SwinT module, and a fourth SwinT module. The backbone structure also includes an SE module, which is connected to the output of the third SwinT module. The SE module is specifically used for: performing global average pooling on the sixth feature map output from the safety helmet detection image after processing by the third SwinT module to output a one-dimensional vector; calculating weight values ​​based on the one-dimensional vector through an activation layer composed of two fully connected layers; and multiplying the weight values ​​by the pixel values ​​of the sixth feature map to obtain the output result. The backbone structure also includes a convolutional block attention module, which is connected to the output of the fourth SwinT module. The convolutional block attention module includes a spatial attention module and a channel attention module. The channel attention module is used to adaptively correct the seventh feature map output by the safety helmet detection image after processing by the fourth SwinT module to generate an eighth feature map. The spatial attention module is used to correct the eighth feature map to output a ninth feature map. The channel attention module is used to adaptively correct the seventh feature map output by the fourth SwinT module after the safety helmet detection image is processed to generate the eighth feature map. The spatial attention module is used to correct the eighth feature map to output the ninth feature map in the following calculation method: in For performing convolution operations; F represents the seventh feature map, F' represents the eighth feature map; F'' represents the ninth feature map; F max C It is the channel max-pooling feature map; MLP is a multilayer perceptron; MaxPool() is the max-pooling function; F avg C This is the channel average pooling feature map; AvePool() is the average pooling function; It is the Sigmoid activation function; f() is the standard convolutional layer; W C and W S These are channel attention weights and spatial attention weights, respectively. The detection module is used to input employee images collected in the work scene into the safety helmet detection model to obtain safety helmet detection results.

6. A computer device comprising a memory and a processor, wherein the memory stores a computer program, characterized in that, When the processor executes the computer program, it implements the method of any one of claims 1 to 4.

7. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by a processor, it implements the method as described in any one of claims 1 to 4.

Images