High-precision lightweight unmanned aerial vehicle image target detection algorithm based on double prediction heads

Abstract

The invention relates to a high-precision lightweight unmanned aerial vehicle image target detection algorithm based on double prediction heads, optimization and improvement are performed in combination with task characteristics of unmanned aerial vehicle image target detection, a double-target detection head module is provided, and the module can improve small target detection precision and can also reduce the parameter quantity of a model. In addition, a small path feature fusion module (Tiny Path Fusion, TP-Fusion) is further designed, the module can fuse richer detail information, and the detector can be more accurately matched with targets of different sizes and shapes. And finally, a coordinate attention block (CA) and an improved bottleneck module are further adopted, so that interference of a complex background on foreground target information detection is reduced, and unmanned aerial vehicle image small target detection is enhanced.

Description

Technical Field

[0001] The present invention relates to a method of using an improved YOLOv8 algorithm to detect targets in unmanned aerial vehicle (UAV) aerial images, and further adopts optimization measures to address the deficiencies of the general YOLOv8 algorithm in UAV target detection, so that the algorithm can be applied to detect specific targets in UAV systems. Background Art

[0002] Drones are widely used in military and civilian fields due to their small size, flexible movement and low cost. At the same time, with the continuous development of deep neural networks and hardware equipment, computer vision tasks such as image classification, target detection and semantic segmentation have made great progress. Combining drones with target detection has also become a research hotspot in recent years. Based on the difference between target detection in drone scenes and general scenes, the background in drone images is complex and there are many messy elements, which can easily interfere with target detection. In addition, there are a large number of small targets that are difficult to identify in drone images. While improving the detection accuracy of targets, drone target detection algorithms also need to consider reducing the complexity of the model as much as possible and reducing the number of model parameters. Therefore, it is necessary to design and study general target detection algorithms in a targeted manner based on the characteristics of target detection in drone scenes.

[0003] In view of the small size and complex background of drone image targets, some scholars have added a detection head for tiny targets on the basis of the general target detection algorithm to detect small targets ignored by other detection heads, which improves the detection accuracy in drone images, but this also increases the number of model parameters. Other scholars use image cropping and a combination of coarse and fine detection to detect drone images, but image cropping may greatly increase the amount of data and significantly prolong the training time, while the combination of coarse and fine detection will slow down the reasoning speed. Therefore, the present invention is designed and optimized based on the general YOLOv8 algorithm and combined with the detection target characteristics of drone images, and proposes a high-precision and lightweight drone image target detection algorithm based on dual prediction heads, thereby achieving improved detection accuracy and efficiency. The experimental results in two public drone image data sets show that the present invention can improve the detection accuracy of a large number of small targets in drone images while reducing the number of model parameters. Summary of the invention

[0004] Aiming at the shortcomings of the general YOLOv8 algorithm in drone image detection, the present invention designs and optimizes the general YOLOv8 algorithm in combination with the characteristics of the detection targets of drone aerial images. Firstly, in view of the problem that the general YOLOv8 algorithm in conventional scenarios uses three detection heads and cannot effectively solve the problem of small target detection in drone aerial images, a YOLOv8 algorithm based on dual target detection heads is proposed, which improves the accuracy of small target detection and reduces the number of model parameters by 42.9%; secondly, a Tiny Path Fusion (TP-Fusion) module is designed, which can fuse richer detail information and enable the detector to more accurately match targets of different sizes and shapes; thirdly, the present invention introduces a Coordinate Attention Block (CA) module, which enables the model to resist complex background information and focus on detecting foreground targets; finally, the present invention proposes a bottleneck module, which is embedded in front of the target detection head of the general YOLOv8 algorithm to further enhance the extraction of small target features.

[0005] The steps of the high-precision and lightweight UAV image target detection algorithm based on dual prediction heads are as follows:

[0006] Step 1: Design and improve the detection head of the YOLOv8 algorithm to make it more suitable for aerial target detection in drone scenarios

[0007] The YOLO series of target detection models use three detection heads to detect targets of different sizes in order to solve the problem of scale changes in conventional real-world scenarios. The three detection heads used have different sizes, so the feature maps input to the three detection heads are downsampled to 1 / 8, 1 / 16, and 1 / 32 of the original size, respectively, to detect smaller targets, medium targets, and larger targets, respectively. Figure 1 is the target box size distribution result of the VisDrone2021 dataset, Figure 1 It can be seen that the color of the target box of the dataset becomes darker as it moves toward the lower left corner. This shows that small targets account for a large proportion of the drone aerial image dataset, which further illustrates the importance of small target detection in aerial images. In the general YOLOv8 target detection algorithm, the input size of the detection head used to detect small targets also uses a 1 / 8 sampling method. However, after the detection head is downsampled, the spatial information of the detection target feature map is reduced, and the feature representation of small targets is also damaged, which is very unfavorable for the detection of small targets in the majority of drone aerial images.

[0008] In order to retain more detailed information of small targets for easy detection of small targets, the present invention designs a target detection head with higher resolution to detect tiny targets missed by other detection heads. Figure 2 As shown, the input of the tiny target detection head added by the present invention comes from the underlying feature map with higher resolution. The feature map size of this detection head is twice that of the small target detection head, which can retain the edge, contour and other detail information of the image to a greater extent, so that the detector needle can more accurately locate small targets in aerial images. However, after adding a target detection head, the number of parameters and the amount of calculation of the model also increase, and there is also a phenomenon of detection information overlap between the four detection heads in the detection model. Furthermore, based on the fact that there are relatively few large targets in the drone aerial image data set, the present invention retains the proposed streamlined detection head, and removes the two prediction heads with smaller feature map resolution in the original algorithm. As can be seen from Table 1, this method not only improves the detection accuracy of the model for small targets in drone images, but also reduces the scale of the model.

[0009] Table 1 Detection results of the model before and after the simplified detection head on the VisDrone validation set

[0010]

[0011]

[0012] Step 2: Design the TP-Fusion module to fuse the details of small objects

[0013] The feature extractor in the general target detection YOLOv8 model extracts features and filters out some noise through downsampling, which also results in the loss of target information. Although the general YOLOv8 uses a feature pyramid to fuse the single-layer information close to the bottom layer and enrich the final features, these feature information is still insufficient for the detection of small targets. Figure 3 It can be seen from the dark thick dashed line that keeps a downward trend that the loss of small target detail information is mainly concentrated in the downsampling of the first few layers. Therefore, in order to retain the small target detail information in the extracted final features to a greater extent, the present invention proposes a TP-Fusion module to further fuse the bottom features close to the input layer.

[0014] The Tiny Path Fusion (TP-Fusion) module proposed in the present invention is as follows: Figure 4 As shown in the figure, this module integrates the information of three layers of feature maps. The feature map of the bottom layer (Underlying Feature Map) is added with the feature map of the second layer (Middle Layer Feature Map) through a 3×3 deformable convolution, and then spliced ​​with the feature map of the top layer (Top Layer Feature Map) on the channel.

[0015] The deformable convolution is characterized by:

[0016]

[0017] The traditional convolution is defined as:

[0018]

[0019] Among them, x(p) and y(p) are the feature representations of the input feature map and the output feature map at position p, respectively, and w k is the weight of the convolution kernel at position k, p k is the offset of p within the convolution kernel, Δp k is the offset introduced by the deformable convolution.

[0020] Among these three layers of feature maps, the one with the largest feature map and the least number of channels contains the most detail information for small target detection; the one with the smaller feature map and the most channels has the richest semantic information for small target classification. By fusing these three feature maps, the present invention not only retains the detail information of small targets to a large extent, but also obtains relatively rich semantic information, which provides stronger support for the decision-making of subsequent target detection heads. Furthermore, the deformable convolution technology introduced in the present invention not only changes the size and number of channels of the feature map to match the feature map of the second layer. More importantly, the deformable convolution can adjust the sampling position and weight of the receptive field by learning deformable parameters, so that the input features at different positions have different effects on the output, so that the detector can better capture the detail information in the image and adapt to targets of different sizes and shapes.

[0021] Step 3: Introduce CA attention module to resist the interference of background noise

[0022] Since military drones cruise at a high altitude, the aerial images captured in the field of view have a wide coverage area, and there is a lot of confusing background information in the aerial images. Therefore, drone aerial image target detection is actually a target detection in a complex background. To this end, the present invention introduces a coordinate attention block (CA), such as Figure 5 As shown in the figure, it helps the network model focus only on the foreground targets that need to be detected and ignores the interference of complex background noise by assigning different weights to different parts of the feature map, further improving the information processing capability. CA is a plug-and-play lightweight attention module that embeds the location information of the feature into the channel attention, enabling the network to focus on a wider range of contextual information while avoiding a large amount of computational overhead.

[0023] First, the coordinate attention module CA introduced in the general YOLOv8 algorithm of the present invention can perform global average pooling of the input feature map in both width and height directions, so as to obtain the attention on the image width and height, and encode the precise position information of the target; then, the two directional feature maps obtained in the previous step are spliced ​​together, and then the intermediate feature map is obtained after operations such as convolution, batch normalization and nonlinear activation function; then, the intermediate feature map is split into tensors in the height and width directions, and the coordinate attention in the height direction and width direction is obtained after convolution and Sigmoid activation function respectively; finally, the input feature map is multiplied by the obtained coordinate attention weight, and finally the feature map with attention weight in the height and width directions is obtained. By introducing CA into the feature fusion module proposed by the present invention, the feature information after the multi-layer fusion of the model can be consciously screened, while highlighting the feature information of the small target area, eliminating the purpose of background noise interference.

[0024] Step 4: Design bottleneck modules to enhance the semantic information of the target

[0025] The two major tasks of target detection are positioning and classification. For positioning, the target is successfully detected only when the detection box is positioned within a certain range and the target in the detection box is correctly classified. For classification, whether the target can be correctly classified depends largely on whether the semantic information of the target is rich enough. When the semantic information of the target is rich enough, the probability of the target being correctly classified is greater. Compared with normal-sized targets in conventional video images, small targets in drone aerial images cover fewer pixels and have lower resolution. If a general YOLOv8 model is used for target classification, it is difficult to learn enough semantic information from it. In order to compensate for the insufficient feature extraction caused by the deletion of medium target detection heads and large target detection heads, a method of enhancing the semantic information of the target is adopted. To this end, the present invention designs a bottleneck module composed of two 3×3 convolution modules and is embedded in front of the detection layer of the model, such as Figure 6 As shown in the figure. After the input feature map passes through the first 3×3 convolution module, the number of channels of the model is reduced to half of the original number; after passing through the second 3×3 convolution module, the number of channels of the model is restored to the number of input channels; on the basis of improving the accuracy of model features and target classification, the computational complexity of the model is further reduced and the delay caused by the new module is shortened.

[0026] In addition, the present invention selects HSwish as the activation function of the structure, and the HSwish activation function is defined as:

[0027]

[0028] ReLU[x]=max(0,x)(4)

[0029] Where x represents the tensor input to the activation function.

[0030] The HSwish activation function is a new type of activation function that combines the advantages of the ReLU and Swish activation functions. Compared with ReLU and Swish, the HSwish curve is smoother, and the gradient and output are relatively stable, which can better ensure the stability of the model with low computational overhead.

[0031] Step 5: Integrate and fuse the modules proposed in steps 1 to 4 into the general YOLOv8 target detection algorithm to obtain a high-precision and lightweight drone image target detection algorithm based on dual prediction heads.

[0032] This step uses the general YOLOv8 algorithm as the baseline model, integrates and fuses the modules of steps one to four respectively, completes a high-precision and lightweight drone image target detection algorithm based on dual prediction heads, and improves the small target detection effect of drone aerial images.

[0033] Step 6: Train and adjust training parameters to obtain the optimal weights of the high-precision and lightweight UAV image target detection algorithm based on dual prediction heads

[0034] In order to obtain stable detection performance and training speed of the model, the present invention needs to optimize and adjust the relevant parameters according to the characteristics of each data set to find the optimal weight.

[0035] (1) Set the input image size

[0036] The resolution of drone aerial images will affect the GPU computing power and memory consumption of the hardware platform used for model training. On the one hand, if the input resolution of the image is too small, some important information may be lost during training, resulting in a decrease in model performance. On the other hand, if the input resolution of the aerial image is too large during training, the computing power and memory consumption may increase, resulting in slower model training and prediction speeds, or even exceeding hardware limitations. In summary, according to the characteristics of drone aerial images and the hardware platform, the input image resolution of the present invention is set to 960×960, which can achieve a balance between efficiency and accuracy.

[0037] (2) Set the training batch size

[0038] Batch size refers to the number of data samples simultaneously input into the model during a training process. The size of Batch Size will affect the effect and speed of the hardware platform training model, and it should be neither too large nor too small. If the Batch Size is too large, the gradient direction of each update will be too smooth, which may exceed the GPU's video memory capacity, resulting in overfitting of the training or decreased generalization ability; if the Batch Size is too small, the gradient direction of each update will be unstable, and the parallel computing capability of the GPU cannot be fully utilized, resulting in non-convergence of training or very slow convergence. Considering the GPU video memory capacity of the training model, the Batch Size of the present invention is set to 2.

[0039] (3) Set the impulse size

[0040] Momentum is a technique for optimizing the gradient descent method. It can accelerate the convergence of the model, avoid falling into local optimality or saddle points, and improve the generalization ability of the model. The size of the momentum determines the degree of influence of the previous update direction on the current update direction. It is generally between 0 and 1, and the commonly used value is 0.9. The momentum of the present invention is set to 0.937.

[0041] (4) Setting the learning rate

[0042] The learning rate is a hyperparameter that adjusts the parameter update step size at each iteration of the optimization algorithm. The size of the learning rate will affect the convergence speed and effect of the model. If the learning rate is too large, the loss function may oscillate or diverge; if the learning rate is too small, it may lead to slow convergence or fall into a local optimum. The initial learning rate of the present invention is set to 0.01, and the final learning rate is set to 0.001.

[0043] (5) Set training rounds

[0044] A training round (epochs) refers to a single training iteration of all batches in forward and backward propagation. An epoch means that each sample in the training dataset has a chance to update the internal model parameters. The size and quality of the training dataset determine how much information the model needs to learn, as well as the difficulty and effect of learning. According to the learning difficulty and size of the dataset used, the present invention sets the training epochs of the VisDrone2021 dataset to 200 and the training epochs of the UAVDT dataset to 15.

[0045] Step 7: Load the designed model and evaluate the model detection performance using the aerial photography dataset

[0046] The weight best.pt is obtained through training, and the performance of the model is evaluated using the VisDrone2021 and UAVDT datasets.

[0047] Beneficial Effects

[0048] In order to solve the problems of small targets to be detected and complex image backgrounds in drone images, the present invention designs and optimizes the general YOLOv8 algorithm based on the task characteristics of drone image target detection, and proposes a YOLOv8 detection algorithm with higher accuracy and lighter weight. First, the present invention replaces the three common target detection heads in the model with a dual target detection head that is more suitable for drone scenarios, which effectively improves the model's detection accuracy for small targets and greatly reduces the number of model parameters. Next, the present invention designs a small path feature fusion module, designs a bottleneck module, and introduces a CA attention module to further enhance the detection performance of the model in drone scenarios. A large number of experiments were carried out on two benchmark datasets. The results show that compared with the baseline model, the AP of the model of the present invention on the VsiDrone2021 dataset is improved by 4.5%, and the AP s It is improved by 5.2% and the number of parameters is reduced by 33.0%, which has considerable competitiveness in the practical application of UAV target detection. BRIEF DESCRIPTION OF THE DRAWINGS

[0049] Figure 1 : VisDrone dataset pixel distribution map

[0050] Figure 2 :Network structure diagram of the high-precision and lightweight UAV image target detection algorithm based on dual prediction heads

[0051] Figure 3 : The detail information and semantic information of objects of different sizes and the changing trend of the model layers

[0052] Figure 4 :TP-Fusion module structure diagram

[0053] Figure 5 :CA attention module structure diagram

[0054] Figure 6 :Bottleneck structure module diagram

[0055] Figure 7 : Overall steps flow chart

[0056] Figure 8 :Comparison of the results of the baseline algorithm and the proposed algorithm on the VisDrone2021 dataset DETAILED DESCRIPTION

[0057] The present invention will now be further described with reference to the embodiments and the accompanying drawings:

[0058] The specific implementation is as follows:

[0059] Step 1: Design the detection head of the YOLOv8 algorithm to make it more suitable for aerial target detection in drone scenarios

[0060] The YOLO series of target detection models use three detection heads to detect targets of different sizes in order to solve the problem of scale changes in conventional real-world scenarios. The three detection heads used have different sizes, so the feature maps input to the three detection heads are downsampled to 1 / 8, 1 / 16, and 1 / 32 of the original size, respectively, to detect smaller targets, medium targets, and larger targets, respectively. Figure 1 is the target box size distribution result of the VisDrone2021 dataset, Figure 1 It can be seen that the color of the target box of the dataset becomes darker as it moves toward the lower left corner. This shows that small targets account for a large proportion of the drone aerial image dataset, which further illustrates the importance of small target detection in aerial images. In the general YOLOv8 target detection algorithm, the input size of the detection head used to detect small targets also uses a 1 / 8 sampling method. However, after the detection head is downsampled, the spatial information of the detection target feature map is reduced, and the feature representation of small targets is also damaged, which is very unfavorable for the detection of small targets in the majority of drone aerial images.

[0061] In order to retain more detailed information of small targets for easy detection of small targets, the present invention designs a target detection head with higher resolution to detect tiny targets missed by other detection heads. Figure 2 As shown, the input of the tiny target detection head added by the present invention comes from the underlying feature map with higher resolution. The feature map size of this detection head is twice that of the small target detection head, which can retain the edge, contour and other detail information of the image to a greater extent, so that the detector needle can more accurately locate small targets in aerial images. However, after adding a target detection head, the number of parameters and the amount of calculation of the model also increase, and there is also a phenomenon of detection information overlap between the four detection heads in the detection model. Furthermore, based on the fact that there are relatively few large targets in the drone aerial image data set, the present invention retains the proposed streamlined detection head, and removes the two prediction heads with smaller feature map resolution in the original algorithm. As can be seen from Table 1, this method not only improves the detection accuracy of the model for small targets in drone images, but also reduces the scale of the model.

[0062] Step 2: Design the TP-Fusion module to fuse the details of small objects

[0063] The feature extractor in the general target detection YOLOv8 model extracts features and filters out some noise through downsampling, which also results in the loss of target information. Although the general YOLOv8 uses a feature pyramid to fuse the single-layer information close to the bottom layer and enrich the final features, these feature information is still insufficient for the detection of small targets. Figure 3 It can be seen from the dark thick dashed line that keeps a downward trend that the loss of small target detail information is mainly concentrated in the downsampling of the first few layers. Therefore, in order to retain the small target detail information in the extracted final features to a greater extent, the present invention proposes a TP-Fusion module to further fuse the bottom features close to the input layer.

[0064] The Tiny Path Fusion (TP-Fusion) module proposed in the present invention is as follows: Figure 4 As shown in the figure, this module integrates the information of three layers of feature maps. The feature map of the bottom layer (Underlying Feature Map) is added with the feature map of the second layer (Middle Layer Feature Map) through a 3×3 deformable convolution, and then spliced ​​with the feature map of the top layer (Top Layer Feature Map) on the channel.

[0065] The deformable convolution is characterized by:

[0066]

[0067] The traditional convolution is defined as:

[0068]

[0069] Among them, x(p) and y(p) are the feature representations of the input feature map and the output feature map at position p, respectively, and w k is the weight of the convolution kernel at position k, p k is the offset of p within the convolution kernel, Δp k is the offset introduced by the deformable convolution.

[0070] Among these three layers of feature maps, the one with the largest feature map and the least number of channels contains the most detail information for small target detection; the one with the smaller feature map and the most channels has the richest semantic information for small target recognition. By fusing these three feature maps, the present invention not only retains the detail information of small targets to a large extent, but also obtains relatively rich semantic information, which provides stronger support for the decision-making of subsequent target detection heads. Furthermore, the deformable convolution technology introduced in the present invention not only changes the size and number of channels of the feature map to match the second layer feature map. More importantly, the deformable convolution can adjust the sampling position and weight of the receptive field by learning deformable parameters, so that the input features at different positions have different effects on the output, so that the detector can better capture the detail information in the image and adapt to targets of different sizes and shapes.

[0071] Step 3: Introduce CA attention module to resist the interference of background noise

[0072] Since military drones cruise at a high altitude, the aerial images captured in the field of view have a wide coverage area, and there is a lot of confusing background information in the aerial images. Therefore, drone aerial image target detection is actually a target detection in a complex background. To this end, the present invention introduces a coordinate attention block (CA), such as Figure 5 As shown in the figure, it helps the network model focus only on the foreground targets that need to be detected and ignores the interference of complex background noise by assigning different weights to different parts of the feature map, further improving the information processing capability. CA is a plug-and-play lightweight attention module that embeds the location information of the feature into the channel attention, enabling the network to focus on a wider range of contextual information while avoiding a large amount of computational overhead.

[0073] First, the coordinate attention module CA introduced in the general YOLOv8 algorithm of the present invention can perform global average pooling of the input feature map in both width and height directions, so as to obtain the attention on the image width and height, and encode the precise position information of the target; then, the two directional feature maps obtained in the previous step are spliced ​​together, and then the intermediate feature map is obtained after operations such as convolution, batch normalization and nonlinear activation function; then, the intermediate feature map is split into tensors in the height and width directions, and the coordinate attention in the height direction and width direction is obtained after convolution and Sigmoid activation function respectively; finally, the input feature map is multiplied by the obtained coordinate attention weight, and finally the feature map with attention weight in the height and width directions is obtained. By introducing CA into the feature fusion module proposed by the present invention, the feature information after the multi-layer fusion of the model can be consciously screened, while highlighting the feature information of the small target area, eliminating the purpose of background noise interference.

[0074] Step 4: Design bottleneck modules to enhance the semantic information of the target

[0075] The two major tasks of target detection are positioning and classification. For positioning, the target is successfully detected only when the detection box is positioned within a certain range and the target in the detection box is correctly classified. For classification, whether the target can be correctly classified depends largely on whether the semantic information of the target is rich enough. When the semantic information of the target is rich enough, the probability of the target being correctly classified is greater. Compared with normal-sized targets in conventional video images, small targets in drone aerial images cover fewer pixels and have lower resolution. If a general YOLOv8 model is used for target classification, it is difficult to learn enough semantic information from it. In order to compensate for the insufficient feature extraction caused by the deletion of medium target detection heads and large target detection heads, a method of enhancing the semantic information of the target is adopted. To this end, the present invention designs a bottleneck module composed of two 3×3 convolution modules and is embedded in front of the detection layer of the model, such as Figure 6 As shown in the figure. After the input feature map passes through the first 3×3 convolution module, the number of channels of the model is reduced to half of the original number; after passing through the second 3×3 convolution module, the number of channels of the model is restored to the number of input channels; on the basis of improving the accuracy of model features and target classification, the computational complexity of the model is further reduced and the delay caused by the new module is shortened.

[0076] In addition, the present invention selects HSwish as the activation function of the structure, and the HSwish activation function is defined as:

[0077]

[0078] ReLU[x]=max(0,x)(4)

[0079] Where x represents the tensor input to the activation function.

[0080] The HSwish activation function is a new type of activation function that combines the advantages of the ReLU and Swish activation functions. Compared with ReLU and Swish, the HSwish curve is smoother, and the gradient and output are relatively stable, which can better ensure the stability of the model with low computational overhead.

[0081] Step 5: Integrate and fuse the modules proposed in steps 1 to 4 into the general target detection YOLOv8 algorithm to obtain the DPH-YOLOv8 design model

[0082] This step uses the general YOLOv8 algorithm as the baseline model, integrates and fuses the modules of steps one to four respectively, completes a high-precision and lightweight drone image target detection algorithm based on dual prediction heads, and improves the small target detection effect of drone aerial images.

[0083] Step 6: Train and adjust training parameters to obtain the optimal weights of DPH-YOLOv8

[0084] In order to obtain stable detection performance and training speed of the model, the present invention needs to optimize and adjust the relevant parameters according to the characteristics of each drone aerial photography data set to find the optimal weight.

[0085] (1) Set the input image size

[0086] The resolution of drone aerial images will affect the GPU computing power and memory consumption of the hardware platform used for model training. On the one hand, if the input resolution of the image is too small, some important information may be lost during training, resulting in a decrease in model performance. On the other hand, if the input resolution of the aerial image is too large during training, the computing power and memory consumption may increase, resulting in slower model training and prediction speeds, or even exceeding hardware limitations. In summary, according to the characteristics of drone aerial images and the hardware platform, the input image resolution of the present invention is set to 960×960, which can achieve a balance between efficiency and accuracy.

[0087] (2) Set the training batch size

[0088] Batch size refers to the number of data samples simultaneously input into the model during a training process. The size of Batch Size will affect the effect and speed of the hardware platform training model, and it should be neither too large nor too small. If the Batch Size is too large, the gradient direction of each update will be too smooth, which may exceed the GPU's video memory capacity, resulting in overfitting of the training or decreased generalization ability; if the Batch Size is too small, the gradient direction of each update will be unstable, and the parallel computing capability of the GPU cannot be fully utilized, resulting in non-convergence of training or very slow convergence. Considering the GPU video memory capacity of the training model, the Batch Size of the present invention is set to 2.

[0089] (3) Set the impulse size

[0090] Momentum is a technique for optimizing the gradient descent method. It can accelerate the convergence of the model, avoid falling into local optimality or saddle points, and improve the generalization ability of the model. The size of the momentum determines the degree of influence of the previous update direction on the current update direction. It is generally between 0 and 1, and the commonly used value is 0.9. The momentum of the present invention is set to 0.937.

[0091] (4) Setting the learning rate

[0092] The learning rate is a hyperparameter that adjusts the parameter update step size at each iteration of the optimization algorithm. The size of the learning rate will affect the convergence speed and effect of the model. If the learning rate is too large, the loss function may oscillate or diverge; if the learning rate is too small, it may lead to slow convergence or fall into a local optimum. The initial learning rate of the present invention is set to 0.01, and the final learning rate is set to 0.001.

[0093] (5) Set training rounds

[0094] A training round (epochs) refers to a single training iteration of all batches in forward and backward propagation. An epoch means that each sample in the training dataset has a chance to update the internal model parameters. The size and quality of the training dataset determine how much information the model needs to learn, as well as the difficulty and effect of learning. According to the learning difficulty and size of the dataset used, the present invention sets the training epochs of the VisDrone2021 dataset to 200 and the training epochs of the UAVDT dataset to 15.

[0095] Step 7: Load the designed model and evaluate the model detection performance using the aerial photography dataset

[0096] The weight best.pt is obtained through training, and the performance of the model is evaluated using two drone aerial photography datasets, VisDrone2021 and UAVDT.

[0097] The overall steps of the algorithm are as follows Figure 7 shown.

[0098] First, considering the large proportion of small objects in drone aerial images and the limited memory resources of drone devices, we designed a detection head for the YOLOv8 algorithm.

[0099] Second, in order to strengthen the information interaction between feature maps of different layers and improve the detailed information of the target, the TP-Fusion module is designed;

[0100] Third, to address the problem of complex backgrounds in drone images, the CA attention module is introduced;

[0101] Fourth, a bottleneck module is designed to enhance the semantic information of the target;

[0102] Fifth, the modules proposed in the above four steps are integrated into the YOLOv8 algorithm to obtain a design model of a high-precision and lightweight drone image target detection algorithm based on dual prediction heads;

[0103] Sixth, start training and gradually adjust the training parameters to obtain the optimal weights of the high-precision and lightweight UAV image target detection algorithm based on dual prediction heads;

[0104] Seventh, load the design model and perform model evaluation.

[0105] 1. Experimental conditions

[0106] In order to test the performance of the high-precision and lightweight drone image target detection algorithm based on dual prediction heads, this paper uses the VisDrone2021 dataset for ablation experiments and comparative experiments of different algorithms, and uses the UAVDT dataset for comparative experiments. Both VisDrone2021 and UAVDT are publicly available drone image datasets with high frequency of use, which are suitable for studying target detection tasks based on drone images.

[0107] Experimental test environment: The system is Ubuntu 20.04, the CPU is Intel (R) Core (TM) i9-10940X CPU @ 3.30GHz, the memory is 48GB, the GPU processor is NVIDIA GeForce GTX 2080Ti, and the video memory is 11GB.

[0108] 2. Experimental content

[0109] The present invention first designs and optimizes the YOLOv8 model from four different perspectives, namely, optimizing the model detection head, designing the TP-Fusion module, introducing the coordinate attention module CA, and designing the bottleneck module and embedding it in front of the detection head. The designed algorithm structure is as follows Figure 2 shown.

[0110] In order to verify the superiority of the high-precision and lightweight drone image target detection algorithm based on dual prediction heads proposed in the present invention in detection accuracy and model size, ablation experiments and comparative experiments of different algorithms were carried out on the VisDrone2021 dataset, and comparative experiments were carried out on the UAVDT dataset.

[0111] 3. Classification evaluation indicators

[0112] The detection quality evaluation index used in the present invention is mean average precision (mAP). The basic concepts involved in the process of calculating mAP are introduced below.

[0113] (1) Intersection over Union (IoU): Intersection over Union ratio. In object detection, the detection box and the true value box are both rectangles. The intersection of the two divided by their union is the IoU value. The IoU value represents the degree of overlap between the detection box and the true value box, and is used to measure the accuracy of positioning. When the overlap rate is greater than the set threshold, the positioning is considered accurate, otherwise it is considered incorrect. When the two boxes completely overlap, the IoU reaches the maximum value of 1.

[0114] (2) True Positive (TP): The number of samples that are correctly identified as positive samples, indicating those samples that are successfully detected.

[0115] (3) True Negatives (TN): The number of negative samples that are correctly identified as negative samples, indicating that the background is not misidentified as a target.

[0116] (4) False Positives (FP): The number of negative samples that are mistakenly identified as positive samples, indicating that the background is identified as the target, i.e., false detection.

[0117] (5) False Negatives (FN): The number of positive samples that are mistakenly identified as negative samples, indicating that the target is mistakenly classified as the background and not detected, that is, missed detection.

[0118] (6) Precision: Accuracy, also known as precision. It indicates the proportion of samples that are actually positive examples among those samples that are positive examples in the detection output results. The calculation formula is:

[0119]

[0120] (7) Recall: Recall rate, also known as recall rate. It indicates the proportion of positive examples that are actually identified to all positive examples. The calculation formula is:

[0121]

[0122] In VOC2010 and later, the average precision (AP) is the area enclosed by the PR (Precision-Recall) curve and the coordinate axis, which is used to measure the detection quality of a class of targets. The calculation formula is as follows:

[0123]

[0124] When detecting multiple categories of targets, averaging the AP of all categories can obtain the mean average precision (mAP) used to measure the quality of multi-category target detection. The calculation formula is as follows:

[0125]

[0126] Where: M represents the number of categories, i∈(1,M).

[0127] Another important evaluation index of the target detection algorithm is speed. Only when the speed meets the real-time requirements can it be applied to industry, which is of great significance for the real-time detection of drones. The commonly used measurement index is the number of frames processed per second (Frame Per Second, FPS), that is, the number of images that can be processed per second, which is defined as follows:

[0128]

[0129] Where: Tot represents the total time spent on detecting an image or video; FC refers to the number of image frames processed. Usually, the FPS obtained by the target detection algorithm varies greatly depending on the hardware configuration used.

[0130] 4. Simulation test

[0131] In order to test and verify the detection performance of the algorithm proposed in this invention on drone images, experiments were carried out on the VisDrone2021 dataset and the UAVDT dataset, and the actual detection results were visualized and analyzed.

[0132] 1. Ablation experiments on the VisDrone dataset

[0133] (1) Optimize the detection result analysis of the detection head

[0134] The present invention experiments with two methods of modifying the detection head, including adding a small target detection head to the baseline model and removing the two detection heads with smaller feature maps of the model based on the above modification. From the data in Table 2, it can be seen that both methods of modifying the detection head improve the detection accuracy of the model, and the detection accuracy of small targets is improved the most. Among them, the first modification method makes AP and AP 50 The AP values ​​for small and medium targets increased by 2.4% and 0.9% respectively. The second modification increased AP and AP 50The AP values ​​for small and medium targets were increased by 1.9% and 2.6% respectively, and the AP values ​​for small and medium targets were increased by 2.5% and 1.2% respectively. This shows that the small target detection head designed by the present invention is helpful for the detection of a large number of small targets in drone images. The detection accuracy of the first modification method is slightly lower than that of the second modification method. The reason is that the combined effect of the newly added small target detection head and the common small target detection head in the model can more comprehensively detect small targets in drone images, while the two detection heads with smaller feature maps in the first method make incorrect judgments on small targets, which mislead the model's final decision on small targets. In addition, small targets account for a large part of the total targets in the VisDrone2021 dataset, so the detection accuracy of the model is slightly reduced. In addition, the modifications of the two detection heads have an impact on the detection accuracy AP of large targets. l The first modification method makes AP l The reason is that the introduced small target detection head has high resolution, large feature map, small receptive field, and in the process of generating the detection frame, the positioning of the large target is not accurate enough, which in turn affects the model's final decision on the large target. l It was reduced by 1.5%. The reason was that after removing the two detection heads with smaller feature maps, the model's missed detection rate for large targets increased. From the perspective of the model's detection efficiency, both modifications reduced the number of model parameters. The first modification reduced it by 5.4%, and the second modification reduced it by 42.9%. The second modification had a much smaller impact on the model's detection speed than the first modification.

[0135] Table 2 Ablation experiments on the VisDrone dataset

[0136]

[0137]

[0138] In summary, the present invention adds a tiny target detection head and removes two detection heads with smaller feature maps. This greatly reduces the model size and hardly affects the detection speed of the model while improving the detection accuracy of the model for drone images.

[0139] (2) TP-Fusion module results analysis

[0140] The rich details contained in the underlying feature map help the model locate the target. According to the fourth row of Table 2, after modifying the feature fusion module, the number of model parameters increased by 0.3M, and the detection speed was slightly affected, but all accuracy indicators were improved, including AP sThe improvement is 1.0%, which indicates that the TP-Fusion module proposed in the present invention is helpful to improve the detection accuracy of small targets.

[0141] (3) Analysis of bottleneck modules after design

[0142] Adding the bottleneck module before the detection head helps the model to further extract features and make up for the semantic information lost by discarding the two detection heads. As can be seen from Table 2, the increase in various accuracy indicators after adding the bottleneck module is between 1.0% and 2.2%, among which the detection accuracy of large objects is improved the most, by 2.2%, which exceeds the AP increase brought by removing the two detection heads with smaller feature maps. l The loss of the detection accuracy of large targets caused by discarding two detection heads is only 1.5%, and the impact on the number of model parameters and detection speed is very small. This shows that the bottleneck module introduced by the present invention solves the problem of reduced large target detection accuracy caused by discarding two detection heads without substantially affecting the detection efficiency of the model.

[0143] (4) Analysis of the results of the coordinate attention module CA

[0144] The main purpose of introducing the CA attention module is to reduce the interference of background information on the network and focus on detecting the foreground target that needs attention. The effect of adding the CA attention module is shown in the last row of Table 2. l In addition, each detection accuracy index has been improved by about 1.0%, and AP l Slightly reduced.

[0145] 2. Comparative experiments on VisDrone2021 and UAVDT datasets

[0146] The present invention compares the proposed network with different algorithms on the VisDrone2021 dataset and the UAVDT dataset, including the YOLO series of target detection algorithms and other target detection algorithms for drone images: YOLOv7-tiny, YOLOv8-s, ClusDet (ResNet50), GLSAN (ResNet50), DMNet (ResNet50), DSHNet (ResNet50), and CEASC (ResNet18).

[0147] The results obtained on the VisDrone2021 validation set are shown in Table 3. l The best results were achieved in other indicators. lThe indicators lag behind ClusDet and DMNet, but these two algorithms need to perform coarse and fine detection of the target and fuse the target frame during the target prediction process, which slows down the reasoning speed. The algorithm of the present invention is an end-to-end one-stage target detection algorithm, which has more advantages in real-time detection and is more suitable for deployment on drones with high latency requirements.

[0148] Table 3 Detection results of the proposed algorithm and other algorithms on the VisDrone2021 validation set

[0149] Method AP(%) <![CDATA[AP 50 (%)]]> <![CDATA[AP 75 (%)]]> <![CDATA[AP s (%)]]> <![CDATA[AP m (%)]]> <![CDATA[AP l (%)]]> YOLOv7 24.0 42.2 23.6 15.2 34.3 45.2 YOLOv8 27.5 46.0 27.7 16.3 37.2 42.0 ClusDet 26.7 50.6 24.7 17.6 38.9 51.4 DMNet 28.2 47.6 28.9 19.9 39.6 55.8 DSHNet 30.3 51.8 30.9 - - - CEASC 28.7 50.7 28.4 - - - Ours 32.0 52.0 32.3 21.5 40.9 42.5

[0150] The results obtained on the UAVDT dataset are shown in Table 4. The AP value of the algorithm of the present invention exceeds that of all compared algorithms, proving that the algorithm of the present invention has good generalization performance.

[0151] Table 4 Detection results of the algorithm of the present invention and other algorithms on UAVDT

[0152] Method AP(%) <![CDATA[AP 50 (%)]]> <![CDATA[AP 75 (%)]]> YOLOv7 18.6 32.7 18.8 ClusDet 13.7 26.5 12.5 GLSAN 19.0 30.5 21.7 DMNet 14.7 24.6 16.3 DSHNet 17.8 30.4 19.7 CEASC 17.1 30.9 17.8 Ours 20.3 33.4 22.7

[0153] 3. Visual analysis of actual test results

[0154] In order to more intuitively compare the performance gap between the proposed algorithm and the baseline model, the prediction results of the two models on the VisDrone2021 test set are visualized, as shown in the figure. Figure 8 As shown, the upper part is the detection result of the baseline model, and the lower part is the detection result of the model of the present invention. At the same time, in order to make the visual effect better, the target category and confidence score in the detection result are removed, and the area that can highlight the performance gap between the two models is enlarged and displayed. As can be seen from the figure, compared with the baseline model, the method of the present invention detects more long-distance tiny vehicle targets in the first scene, the method of the present invention detects more long-distance tiny pedestrian targets in the second scene, the method of the present invention detects more long-distance tiny vehicle targets in the third scene, and the method of the present invention detects more tiny vehicle targets in the fourth scene in dim light. Therefore, the algorithm of the present invention can improve the detection rate of the model for targets with a farther line of sight and a smaller size, and has a better detection effect on drone images with complex backgrounds and small target sizes.

[0155] In summary, the algorithm of the present invention can improve the target detection accuracy in UAV scenarios and reduce the size of the model, and has considerable competitiveness in practical applications.

Claims

1. A high-precision and lightweight UAV image target detection algorithm based on dual prediction heads, characterized in that: The following steps are involved: Step 1: Construct a high-precision and lightweight YOLOv8 target detection algorithm; Step 1-1: Streamlining the prediction head of the YOLOv8 algorithm; adding a small target prediction head with a larger feature map resolution to the general YOLOv8 algorithm, removing the two prediction heads with smaller feature map resolutions in the general YOLOv8 algorithm, and obtaining a dual-prediction-head YOLOv8 target detection algorithm; Step 1-2: Add a Tiny Path Fusion (TP-Fusion) module to the neck network of the general YOLOv8. This module fuses the information of three layers of feature maps. The feature map of the bottom layer (Underlying Feature Map) is added with the feature map of the second layer (Middle Layer Feature Map) through a 3×3 deformable convolution, and then spliced ​​with the feature map of the top layer (Top Layer Feature Map) on the channel. Step 1-3: Continue to add the coordinate attention block (CA) to the neck network of the general YOLOv8; first, CA performs global average pooling on the input feature map in both width and height directions to obtain the attention on the image width and height and encode the precise position information; then the feature maps in the two directions obtained in the previous step are concatenated together, and then the intermediate feature map is obtained through operations such as convolution, batch normalization, and nonlinear activation functions; then the intermediate feature map is split into tensors in the height and width directions, and the coordinate attention in the height and width directions is obtained after convolution and Sigmoid activation function respectively; finally, the input feature map is multiplied by the obtained attention weight to obtain the final feature map with attention weights in the height and width directions; Step 1-4: Continue to add a bottleneck module to the head network of the general YOLOv8; this module consists of two 3×3 convolution modules; the number of channels of the input feature map is reduced to half after passing through the first convolution module, and the number of channels is restored to the input size after passing through the second convolution module, and then a residual connection is made with the input feature map, and finally the final output is obtained through the HSwish activation function; Step 1-5: Using the general YOLOv8 algorithm as the baseline model, integrate and fuse the modules of steps 1-1 to 1-4 respectively to complete the target detection model of drone aerial images - a high-precision and lightweight drone image target detection algorithm with dual prediction heads, which improves the small target detection effect of drone aerial images. Step 2: Train a high-precision and lightweight dual-prediction head drone image target detection algorithm; set training parameters: input image size, batch size, impulse size, learning rate, and maximum number of iterations; use VisDrone2021 and UAVDT datasets to train the network; After training, the final dual-prediction head high-precision and lightweight UAV image target detection algorithm is obtained; Step 3: Load the dual-prediction head high-precision and lightweight drone image target detection algorithm and perform algorithm evaluation.

2. According to claim 1, a high-precision and lightweight UAV image target detection algorithm based on dual prediction heads is characterized in that: The high-precision and lightweight method means that in the target detection scenario of drone aerial images, the accuracy value is improved compared with the baseline network, and the parameter amount is reduced compared with the baseline network.

3. The high-precision and lightweight UAV image target detection algorithm based on dual prediction heads according to claim 1 is characterized in that: The deformable convolution is defined as: The traditional convolution is defined as: Among them, x(p) and y(p) are the feature representations of the input feature map and the output feature map at position p, respectively, and w k is the weight of the convolution kernel at position k, p k is the offset of p within the convolution kernel, Δp k is the offset introduced by the deformable convolution.

4. The high-precision and lightweight UAV image target detection algorithm based on dual prediction heads according to claim 1 is characterized in that: The HSwish activation function is defined as: ReLU[x]=max(0,x) Where x represents the tensor input to the activation function.

5. The high-precision and lightweight UAV image target detection algorithm based on dual prediction heads according to claim 1 is characterized in that: The training parameters are set as follows: the input image size is 960×960, the batch size is 2, the impulse size is 0.937, the initial learning rate is 0.01, the final learning rate is 0.001, the maximum number of iterations of the VisDrone2021 dataset is 200, and the maximum number of iterations of the UAVDT dataset is 15.

Images