A small target recognition method based on improved YOLOv5

Abstract

The application relates to the technical field of target detection, and discloses a small target recognition method based on an improved YOLOv5, which comprises the following steps: S1, constructing an improved YOLOv5 model, wherein the step S1 comprises the following steps: S11, constructing a channel-space parallel attention mechanism module; S12, adding the attention module into a corresponding network layer; S13, performing a concat operation on a 160x160 size feature map output by a 21st layer of a trunk network CSPDarknet and a 2nd layer feature map of the network to serve as an input module of a 4th detection head; and S2, replacing an original CIOU loss function of the YOLOv5 to train the model; the small target recognition method based on the improved YOLOv5 can ensure the calculation speed, increase the detection head, enhance the multi-scale learning ability of the model, increase the context information around the small target through the sub attention mechanism, finally optimize the loss function of the target, improve the regression ability of the model for the target frame, and comprehensively improve the detection precision of the model for the small target, thereby guaranteeing the practicality.

Description

A Small Target Recognition Method Based on Improved YOLOv5 Technical Field

[0001] This invention relates to the field of target detection technology, specifically to a small target recognition method based on an improved YOLOv5. Background Technology

[0002] Currently, object detection methods based on deep neural networks are mainly divided into two categories: one-stage and two-stage. Two-stage object detection algorithms first distinguish foreground and background using a region extraction network, then classify and regress the foreground region of interest to obtain the final detection result. Representative models include Faster R-CNN and DetectorRs. One-stage object detection algorithms directly classify and regress from the feature maps extracted by the network, then use non-maximum suppression methods to obtain the final detection result. Representative models include the YOLO series. Two-stage models have higher detection accuracy, but they are complex and computationally intensive, making them unsuitable for deployment on the edge segments of inspection robots. Therefore, this patent uses the YOLOv5 model as the benchmark model for detecting server indicator lights.

[0003] However, small targets occupy relatively few pixels in an image. On a standard definition image with a resolution of 1920x1080, an indicator light occupies approximately 8x7 pixels. After the image undergoes feature extraction via the backbone network and multiple downsampling processes, the target's feature information is gradually lost, making it difficult for the network to detect and prone to missed detection. Therefore, it is necessary to improve and optimize the YOLOv5 model to enhance its accuracy in detecting small targets. Summary of the Invention

[0004] To address the shortcomings of existing technologies, this invention provides a small target recognition method based on an improved YOLOv5. This method enhances the model's multi-scale learning capability by adding a detection head while maintaining computational speed. It also increases the contextual information surrounding the small target through a sub-attention mechanism. Finally, by optimizing the target's loss function, it improves the model's regression ability for detecting bounding boxes, thereby comprehensively improving the model's detection accuracy for small targets and solving the problems mentioned in the background technology.

[0005] To achieve the goal of enhancing the model's multi-scale learning ability by adding detection heads, increasing the contextual information around small targets through sub-attention mechanisms, and finally improving the model's regression ability for detecting target boxes by optimizing the target loss function, thereby comprehensively improving the model's detection accuracy for small targets, this invention provides the following technical solution:

[0006] A small target recognition method based on improved YOLOv5 includes the following steps:

[0007] S1. Construct an improved YOLOv5 model, wherein S1 includes:

[0008] S11. Construct a channel-space parallel attention mechanism module;

[0009] S12. Add the attention module to the corresponding network layer;

[0010] S13. Perform a concat operation between the 160x160 feature map output from the 21st layer of the backbone network CSPDarknet and the feature map from the 2nd layer of the network, and use it as the input module for the 4th detection head.

[0011] S2. Replace the original CIOU loss function in YOLOv5 and train the model;

[0012] S3. After training is complete, deploy the model.

[0013] Preferably, step S11 includes a first step:

[0014] S111, the attention mechanism module performs feature enhancement on the two-dimensional spatial features in the feature map. S111 includes the following steps:

[0015] S1111: Two 1×1 convolutions are used to generate the query weight wsq and key weight wsk in the spatial attention mechanism module;

[0016] S1112, Perform global max pooling on the query weights;

[0017] S1113. Perform reshape operations on the query weight and key value weight respectively to reshape the matrix form of the weight size;

[0018] S1114. Use the softmax function to activate the query matrix and map it within (0, 1);

[0019] S1115. Perform matrix multiplication, reshape operation and sigmoid function activation on the query matrix and key matrix to obtain the weighted feature tensor of the feature map in space.

[0020] S1116. Perform a dot product between the spatial weight feature tensor and the input feature map to obtain the output of the enhanced features of the feature map in two-dimensional space.

[0021] Preferably, step S11 further includes a second step:

[0022] S112, The attention module performs feature enhancement on the one-dimensional channel features in the spliced ​​feature map. S112 includes the following steps:

[0023] S1121. Two 1×1 convolutions are used to generate the query weight wcq and key weight wck in the channel attention mechanism;

[0024] S1122. Perform a reshape operation on the query weight and key value weight to generate the corresponding matrix;

[0025] S1123. Map the query matrix to the range (0, 1) using softmax;

[0026] S1124. Multiply the query matrix and the key matrix to obtain the channel weight matrix;

[0027] S1125. The channel weight matrix is ​​subjected to 1×1 convolution, layer normalization operation and Sigmoid feature mapping to obtain the channel feature weight tensor.

[0028] S1126. Perform a dot product operation between the channel feature weight tensor and the input feature map to achieve feature enhancement of the one-dimensional channel features in the feature map.

[0029] Preferably, the backbone network is provided with a plurality of attention mechanism modules, which are located at layers 21, 24, 27 and 30 of the backbone network, and the channel and spatial position are treated as parallel attention mechanism modules.

[0030] Preferably, the improved YOLOv5 uses IOU loss as the loss function for calculating the target box regression.

[0031] Compared with existing technologies, this invention provides a small target recognition method based on improved YOLOv5, which has the following beneficial effects:

[0032] This small target recognition method based on an improved YOLOv5 improves the detection head by adding a higher-resolution detection head. The 160x160 feature map output from layer 21 of the backbone network CSPDarknet is concatenated with the feature map from layer 2 of the network to obtain a higher-resolution feature map, which is then used as input to the detection head. The output of this detection head is a 10x10 target, roughly the same pixel size as an indicator light, thus significantly improving the target detection accuracy. A sub-attention mechanism module is added to the feature extraction part of the backbone network to enhance the contextual information around the small indicator light while suppressing the weight of background regions without indicator lights, thereby enhancing the target weight in the feature map. Parallel attention mechanism extraction of channels and spatial positions improves the overall representational ability of the network. This method enhances the model's multi-scale learning ability by adding a detection head while maintaining computational speed. The sub-attention mechanism also enhances the contextual information around the small target. Finally, by optimizing the target loss function, the model's regression ability for detected target boxes is improved, thus comprehensively improving the model's detection accuracy for small targets and ensuring practicality. Attached Figure Description

[0033] Figure 1 is a schematic diagram of the improved YOLOv5 overall architecture in the small target recognition method based on improved YOLOv5 proposed in this invention;

[0034] Figure 2 is a schematic diagram of the parallel attention mechanism module for channel and spatial position in a small target recognition method based on improved YOLOv5 proposed in this invention.

[0035] Figure 3 is a schematic diagram of the loss function in the small target recognition method based on the improved YOLOv5 proposed in this invention;

[0036] Figure 4 is a schematic diagram of the existing YOLOv5 basic architecture. Detailed Implementation

[0037] The technical solutions of the present invention will be clearly and completely described below with reference to the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0038] Please refer to Figure 1-3, which illustrates a small target recognition method based on an improved YOLOv5, including the following steps:

[0039] S1. Construct an improved YOLOv5 model, wherein S1 includes:

[0040] S11. Construct a channel-space parallel attention mechanism module;

[0041] S12. Add the attention module to the corresponding network layer;

[0042] S13. Perform a concat operation between the 160x160 feature map output from the 21st layer of the backbone network CSPDarknet and the feature map from the 2nd layer of the network, and use it as the input module for the 4th detection head.

[0043] S2. Replace the original CIOU loss function in YOLOv5 and train the model;

[0044] S3. After training is complete, deploy the model.

[0045] Further, S11 includes a first step:

[0046] S111, the attention mechanism module performs feature enhancement on the two-dimensional spatial features in the feature map. S111 includes the following steps:

[0047] S1111: Two 1×1 convolutions are used to generate the query weight wsq and key weight wsk in the spatial attention mechanism module;

[0048] S1112, Perform global max pooling on the query weights;

[0049] S1113. Perform reshape operations on the query weight and key value weight respectively to reshape the matrix form of the weight size;

[0050] S1114. Use the softmax function to activate the query matrix and map it within (0, 1);

[0051] S1115. Perform matrix multiplication, reshape operation and sigmoid function activation on the query matrix and key matrix to obtain the weighted feature tensor of the feature map in space.

[0052] S1116. Perform a dot product between the spatial weight feature tensor and the input feature map to obtain the output of the enhanced features of the feature map in two-dimensional space.

[0053] Furthermore, S11 also includes a second step:

[0054] S112, The attention module performs feature enhancement on the one-dimensional channel features in the spliced ​​feature map. S112 includes the following steps:

[0055] S1121. Two 1×1 convolutions are used to generate the query weight wcq and key weight wck in the channel attention mechanism;

[0056] S1122. Perform a reshape operation on the query weight and key value weight to generate the corresponding matrix;

[0057] S1123. Map the query matrix to the range (0, 1) using softmax;

[0058] S1124. Multiply the query matrix and the key matrix to obtain the channel weight matrix;

[0059] S1125. The channel weight matrix is ​​subjected to 1×1 convolution, layer normalization operation and Sigmoid feature mapping to obtain the channel feature weight tensor.

[0060] S1126. Perform a dot product operation between the channel feature weight tensor and the input feature map to achieve feature enhancement of the one-dimensional channel features in the feature map.

[0061] Furthermore, the backbone network is equipped with several attention mechanism modules, which are located at layers 21, 24, 27, and 30 of the backbone network, i.e., the layers where the model extracts feature maps as input to the detection head. These modules are used to enhance the contextual interaction information of the output feature maps of the backbone network, thereby improving the model's ability to detect all targets. Moreover, the parallel attention mechanism modules for channels and spatial positions can improve the network's representation ability as a whole.

[0062] Furthermore, the improved YOLOv5 uses IOU loss as the loss function for calculating bounding box regression, while traditional YOLOv5 uses CIOU loss. Traditional object detection loss functions rely on the aggregation of bounding box regression metrics, such as the distance, overlap, and aspect ratio between the predicted and ground truth boxes (i.e., GIoU, CIoU, ICIoU, etc.). Although CIOU Loss considers the overlap area, center distance, and aspect ratio of bounding box regression, the difference in aspect ratio reflected by 'v' in its formula, rather than the true difference between width and height and their confidence scores, can sometimes hinder the model's effective optimization of similarity. Therefore, this invention uses IOU loss as the loss function for calculating bounding box regression.

[0063] The specific implementation plan is as follows:

[0064] (1) The standard YOLOv5 model is shown in Figure 4. It includes a backbone network for extracting image features, a pyramid layer for multi-scale feature fusion, and three prediction heads that make predictions at different scales;

[0065] (2) When the size of the input image is 640x640, the YOLOv5 model will get three different scale outputs: 80x80 (640 / 8), 40x40 (640 / 16), and 20x20 (640 / 32). The indicator lights that the inspection robot needs to detect are generally about 10 pixels in length and width, which is difficult to detect by the standard YOLOv5 detection model. Therefore, this invention first improves the detection head by adding a higher resolution detection head. The 160x160 feature map output from the 21st layer of the backbone network CSPDarknet is concatted with the feature map of the second layer of the network to obtain a higher resolution feature map as the input of the detection head. Then the output of the detection head is a 10x10 target, which is basically the same as the pixel size of the indicator light, thereby greatly improving the detection accuracy of the target.

[0066] (3) To further improve the recognition rate of occluded indicator lights, this invention adds a sub-attention mechanism module to the feature extraction part of the backbone network to enhance the contextual information around the indicator lights and suppress the weights of background areas without indicator lights, thereby enhancing the target weights of the feature map. The specific sub-attention mechanism module is shown in Figure 2. This invention proposes to extract attention mechanisms for channels and spatial positions in parallel, thereby improving the overall representation ability of the network. As shown in Figure 2, the input of the feature map is X, with a height of H, a width of W, and a number of channels of C. After passing through a 1×1 convolution operation, the query weights wcq and key weights wck of the channel attention mechanism, as well as the query weights wsq and key weights wsk of the position, are generated. The query weights wcq of the channel attention mechanism are reshaped into a C / 2×HW matrix. The key weights (wck) are resized to obtain a WH×1 matrix. Then, a softmax activation function is applied to map the key matrix to the range (0, 1). A matrix multiplication operation is then performed between the lookup matrix and the key matrix to obtain a heavy matrix of size C / 2×1. This matrix undergoes a 1×1 convolution operation, followed by layer normalization and a sigmoid activation function to obtain the channel weight features, each of size C×1×1 with values ​​between 0 and 1. The channel weight features are then multiplied by the input original feature map (dot multiplication is the multiplication of corresponding values ​​in two feature maps, a common operation in deep neural networks) to obtain the feature information of each pixel in the channel direction. The position lookup weights (wsq) undergo global pooling to obtain the position information representation features of each channel, C / 2×1×1. This is then resized to obtain a 1×C / 2 matrix, and a softmax activation function is applied to map the channel feature weights of the position information to the range (0, 1). Similar to channel key-value pairs, positional key-value pairs are converted into matrices through 1×1 convolution and resizing operations. These matrices are then multiplied with the channel feature weights of the positional information to obtain the weight matrix for each position. After resizing and applying the sigmoid activation function, the final weight features for each position in the feature map are obtained. These weights range from 0 to 1; a larger value indicates a higher probability of a target being present at that position. The positional weights are then multiplied by the input feature map to obtain the final attention features at the positional information. Finally, the positional attention features and channel attention features are added to obtain the output of the attention mechanism module.

[0067] (4) Since the target detection of the robot for the inspection room is not only for small indicator lights, but also for equipment, large indicator lights, display panels, instruments, etc., the attention mechanism module is added to the 21st, 24th, 27th and 30th layers of the backbone network, that is, the layer where the model extracts feature maps as the input of the detection head, to enhance the contextual interaction information of the output feature maps of the backbone network, thereby improving the model's ability to detect all targets. The overall network structure is shown in Figure 1.

[0068] (5) When training the model, the loss function of YOLOv5 was optimized. This invention uses the following defined IOU loss as the loss function for calculating the target box regression:

[0069]

[0070] In the above formula, α is a hyperparameter we designed, which gives the detector greater flexibility in achieving regression accuracy for target boxes of different sizes; the function ρ(n, m) represents the calculation of the Euclidean distance between points n and m; in the above formula, b represents the center point of the target box predicted by the model; b gt The center point of the actual object bounding box is represented by ; w represents the width of the predicted bounding box; h represents the height of the predicted bounding box; and h is the corresponding value. gt and w gt c represents the height and width of the actual object bounding box; c represents the distance from the diagonal of the smallest rectangle containing both the predicted and actual object bounding boxes. h This represents the height of the smallest rectangle that encloses both the predicted and ground truth boxes; c w This represents the width of the smallest rectangle that encloses the predicted bounding box and the ground truth bounding box.

[0071] Wherein, IOU is the intersection-union ratio of bounding boxes, a commonly used method in object detection, which is well-known and will not be elaborated further.

[0072]

[0073] In the above formula, d represents the distance between the center points of the predicted bounding box and the ground truth bounding box; c represents the distance along the diagonal of the smallest rectangle containing both the predicted bounding box and the ground truth bounding box; the representations of c and d are shown in Figure 3, where:

[0074]

[0075]

[0076] Where y1, y2, x1, and x2 are the coordinates of the smallest rectangle containing the predicted bounding box and the ground truth bounding box, as shown in Figure 3; w preThis represents the width of the prediction box; h pre This represents the height of the prediction box. Based on the data we collected from the server room, setting α=1.65 resulted in the fastest convergence speed and the highest accuracy of the model.

[0077] Furthermore, the TPH-YOLOv5 model, proposed in the "Vision Meets Drones: A Challenge" at the ICCV 2021 workshop, is primarily used for small target recognition in UAVs. Its attention mechanism employs the CBAM model, which has large parameters, resulting in slow training and inference speeds, failing to meet real-time requirements. The attention mechanism model of this invention has relatively fewer parameters, meets real-time requirements, and exhibits superior accuracy compared to the TPH-YOLOv5 model.

[0078] The beneficial effects of this invention are:

[0079] 1. This small target recognition method based on improved YOLOv5 improves the detection head by adding a higher-resolution detection head. The 160x160 feature map output from layer 21 of the backbone network CSPDarknet is concatenated with the feature map from layer 2 of the network to obtain a higher-resolution feature map, which is then used as input to the detection head. The output of this detection head is a 10x10 target, which is roughly the same pixel size as the indicator light, thus significantly improving the target detection accuracy. A sub-attention mechanism module is added to the feature extraction part of the backbone network to enhance the contextual information around the small indicator light while suppressing the weight of background regions without indicator lights, thereby enhancing the target weight in the feature map. Parallel attention mechanism extraction of channels and spatial positions improves the overall representation ability of the network. While maintaining computational speed, the method enhances the model's multi-scale learning ability by adding a detection head and increases the contextual information around the small target through the sub-attention mechanism. Finally, by optimizing the target loss function, the model's regression ability for detecting target boxes is improved, thus comprehensively improving the model's detection accuracy for small targets and ensuring practicality.

[0080] Although embodiments of the invention have been shown and described, it will be understood by those skilled in the art that various changes, modifications, substitutions and alterations can be made to these embodiments without departing from the principles and spirit of the invention, the scope of which is defined by the appended claims and their equivalents.

Claims

1. A small target recognition method based on improved YOLOv5, characterized in that: The steps include: S1, constructing an improved YOLOv5 model, S1 including: S11, constructing a channel-space parallel attention mechanism module, S11 including the first step: S111, the attention mechanism module performs feature enhancement on the two-dimensional spatial features in the feature map, S111 including the following steps: S1111, using two 1×1 convolutions to generate query weights wsq and key weights wsk in the spatial attention mechanism module; S1112, performing global max pooling on the query weights to obtain the expression features of position information in each channel C / 2×1×1, where C is the number of channels; S1113, performing reshape operations on the query weights and key weights respectively to reshape the matrix form of the weight size to obtain a 1×C / 2 matrix; S1114, using the softmax function to activate the query matrix, mapping the channel feature weights of position information within (0,1); S1115, performing matrix multiplication between the query matrix and the key matrix to obtain the weight matrix for each position, and then reshaping... The APE operation and Sigmoid activation function are used to obtain the spatial weight feature tensor of the feature map, whose value is between (0, 1). The larger the value, the greater the probability that the target exists at the corresponding position; S1116, the spatial weight feature tensor is multiplied by the input feature map to obtain the output of the enhanced feature of the feature map in two-dimensional space; S11 also includes a second step: S112, the attention module enhances the one-dimensional channel features in the spliced ​​feature map, S112 includes the following steps: S1121, two 1×1 convolutions are used to generate the query weight wcq and key weight wck in the channel attention mechanism; S1122, the query weight and key weight are reshaped to generate corresponding matrices, a C / 2×HW matrix and a WH×1 matrix, where H is the height dimension and W is the width dimension; S1123, the query matrix is ​​mapped to (0, 1) through softmax, and then the key matrix is ​​mapped to (0, 1) through the softmax activation function; S1124, Multiply the query matrix and the key matrix to obtain the channel weight matrix, which has a size of C / 2×1; S1125, the channel weight matrix is ​​subjected to 1×1 convolution, layer normalization, and Sigmoid feature mapping to obtain the channel feature weight tensor, which has a size of C×1×1; S1126, the channel feature weight tensor is multiplied with the input feature map to achieve feature enhancement of the one-dimensional channel features in the feature map; S12, the attention module is added to the corresponding network layer; S13, the 160x160 feature map output from the 21st layer of the backbone network CSPDarknet is concatted with the feature map of the second layer of the network as the input module of the fourth detection head; S2, replace the original CIOU loss function of YOLOv5 and train the model; S3, after training is completed, deploy the model.

2. The small target recognition method based on improved YOLOv5 according to claim 1, characterized in that: The backbone network is equipped with several attention mechanism modules, which are located at layers 21, 24, 27, and 30 of the backbone network, respectively, and the channel and spatial position are treated as parallel attention mechanism modules.

3. The small target recognition method based on improved YOLOv5 according to claim 1, characterized in that: The improved YOLOv5 uses IOU loss as the loss function for calculating the target box regression.

Images