Target detection model improvement method based on YOLOv5 in coal gangue detection

Abstract

The invention belongs to the technical field of target detection, and aims to solve the problem that the detection performance of YOLOv5 may be limited in a coal dressing scene or a small target. The invention provides an improved method based on a YOLOv5 target detection model in coal gangue detection. The method comprises the following steps: inserting an ECA module into key nodes of feature extraction and feature fusion in a backbone network and a feature fusion network of an original YOLOv5 model; defining a loss function of the new YOLOv5 model; carrying out new YOLOv5 model training by adopting a data set which is the same as the original YOLOv5 model; and collecting real-time data, inputting the real-time data into the optimized YOLOv5 model, evaluating the output of the model, and verifying the screening performance of the new YOLOv5. The ECA mechanism is introduced, so that the feature details of the small target can be better highlighted, and the method has remarkable advantages for application scenes needing to accurately detect the small target.

Description

Technical Field

[0001] The present invention belongs to the technical field of target detection, and in particular relates to an improved method based on a YOLOv5 target detection model in coal gangue detection. Background Art

[0002] In the coal preparation process, the mixing of coal gangue, domestic waste and coal is a common problem. The traditional manual sorting method is inefficient and easily affected by human factors.

[0003] YOLOv5 is a convolutional neural network model widely used in target detection tasks. YOLOv5 performs well in real-time target detection due to its fast detection speed, multi-target detection, and high accuracy. Its basic idea is to regard the target detection task as a single-stage regression problem and directly predict the bounding box and category of the target. In the coal preparation scenario, the model needs to have high real-time performance and be able to perform multi-target detection, so that coal can be quickly screened out from more objects and used efficiently. YOLOv5 can meet these requirements.

[0004] However, when faced with complex scenes or small targets, the detection performance of YOLOv5 may be limited. This is because in these cases, the features of the target may be submerged by background noise, or the feature expression of the target is not sufficient to distinguish different categories. For example, in the coal selection scene, coal will reflect light, while coal gangue and domestic garbage basically do not reflect light. The selection of coal requires special attention to the grayscale information of the object.

[0005] In the traditional structure of YOLOv5, the lack of effective use of inter-channel correlation in the feature extraction process may cause some important feature information to be masked or lost. Summary of the invention

[0006] In order to solve at least one of the above technical problems existing in the prior art, the present invention provides an improved method based on the YOLOv5 target detection model in coal gangue detection.

[0007] The present invention is implemented by the following technical scheme: an improved method based on the YOLOv5 target detection model in coal gangue detection, comprising the following steps: S1: inserting an ECA module into key nodes of feature extraction and feature fusion in a backbone network and a feature fusion network of an original YOLOv5 model, and keeping the overall structure of the original network unchanged; S2: defining a loss function of a new YOLOv5 model, including definitions of a classification loss function, a positioning loss function and a confidence loss function and weight coefficients of the three in the loss function of the new YOLOv5 model; S3: using the same data set as the original YOLOv5 model to train the new YOLOv5 model, and optimizing the parameters of the new YOLOv5 model; S4: collecting real-time data to input into the optimized YOLOv5 model, and evaluating the output of the model to verify the screening performance of the new YOLOv5.

[0008] Preferably, the locations where the ECA module is inserted include:

[0009] Inserted after the Bottleneck module of the C3 module, it is used to improve the expressiveness of inter-channel features while keeping the network lightweight;

[0010] After being inserted into the PANet feature fusion module of the feature fusion network, it is used to enhance the semantic information fusion effect between features of different scales;

[0011] After the feature concatenation is inserted into the spatial pyramid pooling fusion module, the important features are strengthened and irrelevant information is suppressed by weighted processing of each channel.

[0012] Preferably, the loss function L is expressed as:

[0013] L=α class ×L class +β loc ×L GIoU +γ con ×L con

[0014] Where, L class is the classification loss function, L GIoU is the positioning loss function, L con is the confidence loss function; α class , β loc , γ con They are the weight coefficients of the classification loss function, positioning loss function and confidence loss function respectively.

[0015] Preferably, the classification loss function L class The expression is:

[0016]

[0017] y i =sigmoid(x i )

[0018] In the formula, N is the total number of categories, x i is the predicted value of this category, y i is the value of the current category after the activation function, is the true value of the current category (either 0 or 1).

[0019] Preferably, the positioning loss function L GIoU的 The expression is:

[0020] L GIoU (B,B gt )=1-GLoU(B,B gt )

[0021] In the formula, Where C is the smallest box that can contain both the real box and the candidate box. B is the candidate box, B gt The real label box.

[0022] Preferably,

[0023]

[0024] In the formula, x1 is the minimum horizontal coordinate of the box, x2 is the maximum horizontal coordinate of the box; y1 is the minimum vertical coordinate of the box, y2 is the maximum vertical coordinate of the box, and the superscripts C, B, and B gt All are boxes where they are located; C is the area of ​​box C.

[0025] Preferably, the expression of the confidence loss function is:

[0026] L con =(1.0-Self gr )+Self gr ×Score GIoU

[0027] In the formula, Self gr is the GIoU loss ratio, which is a hyperparameter in the YOLOv5 model used to balance the classification loss and GIoU loss in the confidence score; Score GIoU Represents the GIoU score, with a value between 0 and 1, where 1 indicates complete overlap.

[0028] Preferably, the EAC module inserted after the Bottleneck module of the C3 module performs EAC operation after each bottleneck is completed.

[0029] Compared with the prior art, the present invention has the following beneficial effects:

[0030] The present invention can enhance the feature expression capability by introducing the ECA module into the YOLOv5 model. Based on the ECA mechanism, the discrimination of feature expression can be effectively improved, making the network more targeted and accurate. On the other hand, semantic information fusion can be improved. The ECA mechanism can optimize the information transmission efficiency of each channel in the feature map, so that the semantic information of features of different scales can be more comprehensively and effectively integrated. In the coal mining scenario, this technology can more effectively capture the reflective characteristics of coal, so that the coal can be quickly and accurately screened out. The introduction of the ECA mechanism can better highlight the feature details of small targets, which has significant advantages for application scenarios that require accurate detection of small targets. It is applied in coal preparation scenarios to realize the automatic identification and separation of coal gangue, domestic waste and coal. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative work.

[0032] Figure 1 It is a structural diagram of introducing the ECA mechanism into the C3 module;

[0033] Figure 2 It is a structural diagram that introduces the ECA mechanism into feature fusion;

[0034] Figure 3 It is a structural diagram of introducing the ECA mechanism into the spatial pyramid pooling fusion module;

[0035] Figure 4 It is the positioning loss function for the candidate box B and the real annotation box B gt And a schematic diagram of the smallest box C that can frame both the real box and the candidate box. DETAILED DESCRIPTION

[0036] In conjunction with the drawings in the embodiments of the present invention, the technical solutions in the embodiments of the present invention are clearly and completely described. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other implementation methods obtained by ordinary technicians in this field without making creative work are within the scope of protection of the present invention.

[0037] It should be noted that the structures, proportions, sizes, etc. illustrated in the drawings of this specification are only used to match the contents disclosed in the specification so that people familiar with this technology can understand and read them. They are not used to limit the conditions under which the present invention can be implemented, and therefore have no substantive technical significance. Any structural modification, change in proportional relationship or adjustment of size should fall within the scope of the technical contents disclosed in the present invention without affecting the effects and purposes that can be achieved by the present invention. It should be noted that in this specification, relational terms such as first and second are only used to distinguish one entity from several other entities, and do not necessarily require or imply any actual relationship or order between these entities.

[0038] The present invention provides an embodiment:

[0039] The invention discloses an improved method for a target detection model based on YOLOv5 in gangue detection, comprising the following steps: S1: inserting an ECA module into key nodes of feature extraction and feature fusion in a backbone network and a feature fusion network of an original YOLOv5 model, and keeping the overall structure of the original network unchanged; S2: defining a loss function of a new YOLOv5 model, including definitions of a classification loss function, a positioning loss function and a confidence loss function, and weight coefficients of the three in the loss function of the new YOLOv5 model; S3: using the same data set as the original YOLOv5 model to train the new YOLOv5 model, and optimizing the parameters of the new YOLOv5 model; S4: collecting real-time data to input into the optimized YOLOv5 model, and evaluating the output of the model to verify the screening performance of the new YOLOv5.

[0040] 1. Select the location to insert the ECA module: The original YOLOv5 model mainly consists of three parts: Backbone (backbone network): used to extract image features; Neck (feature fusion network): used to fuse features at different levels; Head (detection head): used to predict the category and position of the target. According to the architectural characteristics of the YOLOv5 model, choose to insert the ECA module in the important modules in Backbone (backbone network) and Neck (feature fusion network). These locations are the key nodes for model feature extraction and feature fusion. Inserting the ECA module can effectively improve the feature representation capability.

[0041] After the Bottleneck module is inserted into the C3 (three convolutional layers) module, the C3 module is the key part of feature extraction. The ECA operation is performed immediately after each Bottleneck module is completed to improve the expressiveness of inter-channel features while keeping the network lightweight; Figure 1As shown, ECA is performed after each bottleneck ends, rather than after n bottlenecks end, so the bottleneck+ECA writing method is adopted.

[0042] After being inserted into the PANet feature fusion module of the feature fusion network, PANet is responsible for the fusion of features of different scales. In the YOLOv5 network, upsampling is concat (feature fusion) with features of the corresponding level. Introducing ECA after each feature fusion can enhance the semantic information fusion effect between features of different scales; Figure 2 shown.

[0043] After the feature concatenation of the spatial pyramid pooling fusion module SPPF, the ECA module is introduced after the Concat (feature concatenation) of SPPF to improve the expression ability of the concatenated feature map. The ECA module can enhance important features and suppress irrelevant information by weighting each channel; Figure 3 As shown in the figure, except for inserting the ECA module at the corresponding position described above in the original YOLOv5 network, the overall structure of the original network remains unchanged.

[0044] right Figures 1 to 3 The abbreviations in the following are explained:

[0045] BCHW:

[0046] B: Batch size, indicating the batch size; C: Channels, indicating the number of channels of the input image; H: Height, indicating the height of the input image; W: Width, indicating the width of the input image. BCHW means that the size of the input image is batch size × number of channels × height × width.

[0047] ConvBNSiLU(k1,s1,p0,c(C / 2)):

[0048] Conv: convolution layer, representing convolution operation; BN: Batch Normalization, batch normalization layer, used to speed up training and stabilize the training process; SiLU: Sigmoid Linear Unit, an activation function; k1: the size of the convolution kernel is 1x1; s1: the stride is 1, that is, the step length of the convolution kernel movement is 1; p0: padding is 0, that is, no additional edge padding is used; c (C / 2): the number of convolution kernels is C / 2, indicating that the number of output channels is half of the number of input channels; the convolution operation of other parameters is similar.

[0049] MaxPool2d:

[0050] MaxPool: Max Pooling, maximum pooling operation; 2d: two-dimensional, indicating that the maximum pooling operation is performed in two-dimensional space.

[0051] upsample: Upsample.

[0052] 2. Design loss function L:

[0053] 2.1 Classification loss function L class The expression is:

[0054]

[0055] y i =sigmoid(x i )

[0056] In the formula, N is the total number of categories, x i is the predicted value of this category, y i is the value of the current category after the activation function, is the true value of the current category (either 0 or 1).

[0057] 2.2. Positioning loss is used to measure the difference between the predicted bounding box and the true bounding box. This patent uses GIoU (Generalized Intersection over Union) as its positioning loss function. GIoU improves the traditional IoU and can better handle the situation where the bounding boxes do not overlap.

[0058] Positioning loss function L GIoU The expression is:

[0059] L GIoU (B,B gt )=1-GLoU(B,B gt )

[0060] In the formula, Where C is the smallest box that can contain both the real box and the candidate box. B is the candidate box, B gt The real label box.

[0061]

[0062]

[0063] In the formula, x1 is the minimum horizontal coordinate of the box, x2 is the maximum horizontal coordinate of the box; y1 is the minimum vertical coordinate of the box, y2 is the maximum vertical coordinate of the box, and the superscripts C, B, and B gt All are boxes where they are located; C is the area of ​​box C. Figure 4 shown.

[0064] 2.3. The expression of confidence loss function is:

[0065] L con =(1.0-Self gr )+Self gr ×Score GIoU

[0066] In the formula, Self gr is the GIoU loss ratio, which is a hyperparameter in the YOLOv5 model used to balance the classification loss and GIoU loss in the confidence score; Score GIoU Represents the GIoU score, with a value between 0 and 1, where 1 indicates complete overlap.

[0067] 2.4. In summary, the expression of the loss function L is:

[0068] L=α class ×L class +β loc ×L GIoU +γ con ×L con

[0069] Where, L class is the classification loss function, L GIoU is the positioning loss function, L con is the confidence loss function;

[0070] α class , β loc , γ con They are the weight coefficients of the classification loss function, the positioning loss function, and the confidence loss function. The weight coefficients can be adjusted dynamically during the training process. Consider that the confidence coefficient can be slightly larger in the early stage of training to make the model more focused on target detection, and the confidence coefficient can be lowered in the later stage to improve the positioning accuracy of the model. The weights of the classification loss and the positioning loss can be equal at the beginning and gradually adjusted according to the training results (accuracy, etc.).

[0071] This application proposes a method to introduce the Efficient Channel Attention (ECA) mechanism into the YOLOv5 model. ECA is an efficient channel attention mechanism. This module can use a 1x1 convolution layer directly after the global average pooling layer, removing the fully connected layer. This module avoids dimensionality reduction and effectively captures cross-channel interactions. ECA only involves a few parameters to achieve good results, which not only reduces computational overhead but also improves model performance. Integrating the ECA mechanism into an important position in the YOLOv5 network can significantly improve its detection accuracy while maintaining the model's efficiency.

[0072] The above is only a preferred specific embodiment of the present invention, but the protection scope of the present invention is not limited thereto. Any changes or substitutions that can be easily thought of by a person skilled in the art within the technical scope disclosed by the present invention should be included in the protection scope of the present invention. Therefore, the protection scope of the present invention shall be based on the protection scope of the claims.

Claims

1. A method for improving the YOLOv5 target detection model in coal gangue detection, characterized in that: The following steps are involved: S1: Insert the ECA module into the key nodes of feature extraction and feature fusion in the backbone network and feature fusion network of the original YOLOv5 model, and keep the overall structure of the original network unchanged; S2: Define the loss function of the new YOLOv5 model, including the definitions of the classification loss function, positioning loss function, and confidence loss function, as well as the weight coefficients of the three in the loss function of the new YOLOv5 model; S3: Use the same dataset as the original YOLOv5 model to train a new YOLOv5 model and optimize the parameters of the new YOLOv5 model; S4: Collect real-time data to input the optimized YOLOv5 model and evaluate the output of the model to verify the screening performance of the new YOLOv5.

2. The improved method based on the YOLOv5 target detection model in gangue detection according to claim 1, characterized in that: The locations where the ECA modules are inserted include: Inserted after the Bottleneck module of the C3 module, it is used to improve the expressiveness of inter-channel features while keeping the network lightweight; After being inserted into the PANet feature fusion module of the feature fusion network, it is used to enhance the semantic information fusion effect between features of different scales; After the feature concatenation is inserted into the spatial pyramid pooling fusion module, the important features are strengthened and irrelevant information is suppressed by weighted processing of each channel.

3. The improved method based on the YOLOv5 target detection model in coal gangue detection according to claim 1, characterized in that: The expression of the loss function L is: L=a class ×L class +b loc ×L GIoU +g con ×L con Where, L class is the classification loss function, L GIoU is the positioning loss function, L con is the confidence loss function; α class , β loc , γ con They are the weight coefficients of the classification loss function, positioning loss function and confidence loss function respectively.

4. The improved method based on the YOLOv5 target detection model in gangue detection according to claim 3 is characterized in that: Classification loss function L class The expression is: y i =sigmoid(x i ) In the formula, N is the total number of categories, x i is the predicted value of this category, y i is the value of the current category after the activation function, is the true value of the current category (either 0 or 1).

5. The improved method based on the YOLOv5 target detection model in gangue detection according to claim 3 is characterized in that: Positioning loss function L GIoU The expression is: L GIoU (B,B gt )=1-GLoU(B,B gt ) In the formula, Where C is the smallest box that can contain both the real box and the candidate box. B is the candidate box, B gt The real label box.

6. The improved method based on the YOLOv5 target detection model in gangue detection according to claim 5, characterized in that: In the formula, x1 is the minimum horizontal coordinate of the box, x2 is the maximum horizontal coordinate of the box; y1 is the minimum vertical coordinate of the box, y2 is the maximum vertical coordinate of the box, and the superscripts C, B, and B gt All are boxes where they are located; C is the area of ​​box C.

7. The improved method based on the YOLOv5 target detection model in gangue detection according to claim 3 is characterized in that: The expression of the confidence loss function is: L con =(1.0-Self gr )+Self gr ×Score GIoU In the formula, Self gr is the GIoU loss ratio, which is a hyperparameter in the YOLOv5 model used to balance the classification loss and GIoU loss in the confidence score; Score GIoU Represents the GIoU score, with a value between 0 and 1, where 1 indicates complete overlap.

8. The improved method based on the YOLOv5 target detection model in gangue detection according to claim 2, characterized in that: The EAC module inserted after the Bottleneck module of the C3 module performs EAC operation after each bottleneck.

Images