Unmanned aerial vehicle image target detection algorithm based on double prediction head high-precision light weight

Abstract

The application relates to a high-precision lightweight unmanned aerial vehicle image target detection algorithm based on double prediction heads, which is optimized and improved in combination with the task characteristics of unmanned aerial vehicle image target detection, and a double target detection head module is proposed, which can improve small target detection precision while reducing the parameter quantity of a model. In addition, a small path feature fusion module (TP-Fusion) is further designed, the module can fuse more abundant detailed information, and the detector can more accurately match targets with different sizes and shapes. Finally, a coordinate attention module (CA) and an improved bottleneck module are further adopted, the interference of a complex background on foreground target information detection is reduced, and the unmanned aerial vehicle image small target detection is enhanced.

Description

TECHNICAL FIELD

[0001] The present application relates to a kind of unmanned aerial vehicle aerial image target detection using improved YOLOv8 algorithm, and for the deficiency of general YOLOv8 algorithm in unmanned aerial vehicle target detection, further optimization measures are taken, can be applied to the detection of specific target in unmanned aerial vehicle system. BACKGROUND

[0002] Unmanned aerial vehicle is widely used in military and civilian fields due to its small size, flexible movement and low cost. At the same time, with the continuous development of deep neural networks and hardware devices, great progress has been made in computer vision tasks such as image classification, target detection and semantic segmentation. The combination of unmanned aerial vehicle and target detection has become a research hotspot in recent years. Based on the difference between target detection in unmanned aerial vehicle scene and general scene, the background of unmanned aerial vehicle image is complex, with many clutter elements that can interfere with target detection, and there are also many small targets that are difficult to identify. While improving the detection accuracy of targets, the unmanned aerial vehicle target detection algorithm also needs to consider reducing the complexity of the model and reducing the number of model parameters as much as possible. Therefore, it is necessary to design and research the general target detection algorithm according to the characteristics of target detection in unmanned aerial vehicle scene.

[0003] To solve the problem of small size and complex background of unmanned aerial vehicle image target, some scholars add a detection head for small targets to the general target detection algorithm to detect small targets ignored by other detection heads, improving the detection accuracy of unmanned aerial vehicle image. However, this also increases the number of model parameters. Some scholars use image cropping and coarse-fine detection combination to detect unmanned aerial vehicle image, but image cropping may significantly increase the amount of data and prolong the training time, and using coarse-fine detection combination may slow down the inference speed. Therefore, the present application is designed and optimized based on the general YOLOv8 algorithm combined with the characteristics of the detection target of unmanned aerial vehicle image, and proposes a high-precision lightweight unmanned aerial vehicle image target detection algorithm based on double prediction head, to improve the detection accuracy and efficiency. The experimental results on two public unmanned aerial vehicle image datasets show that the present application can improve the detection accuracy of a large number of small targets in unmanned aerial vehicle image while reducing the number of model parameters. SUMMARY

[0004] In view of the deficiencies of the general YOLOv8 algorithm in unmanned aerial vehicle image detection, the general YOLOv8 algorithm is designed and optimized in combination with the characteristics of the detection target of the unmanned aerial vehicle aerial image. First, in view of the problem that the general YOLOv8 algorithm in the conventional scene adopts three detection heads and cannot effectively solve the problem of small target detection of aerial images in the unmanned aerial vehicle scene, a YOLOv8 algorithm based on double target detection heads is proposed, which improves the small target detection accuracy while reducing the parameter quantity of the model by 42.9 %; secondly, a small path feature fusion module (Tiny PathFusion, TP-Fusion) is designed, which can fuse more detailed information and enable the detector to more accurately match targets of different sizes and shapes; thirdly, the Coordinate Attention Block (CA) is introduced, so that the model can resist complex background information and focus on detecting foreground targets; finally, a bottleneck module is proposed, which is embedded in the target detection head of the general YOLOv8 algorithm, further enhancing the extraction of small target features.

[0005] The steps of the unmanned aerial vehicle image target detection algorithm based on double prediction heads and high-precision light weight are as follows:

[0006] Step one: design and improve the detection head of the YOLOv8 algorithm to make it more suitable for aerial target detection in the unmanned aerial vehicle scene

[0007] The target detection model of the YOLO series adopts three detection heads to detect targets of different sizes to solve the problem of size change of targets in the conventional scene. The three detection heads adopted are of different sizes, and relative to the input image size, the feature maps input into the three detection heads are down-sampled to 1 / 8, 1 / 16 and 1 / 32 of the original size, respectively, for detecting small, medium and large targets. Figure 1 is the target box size distribution result of the VisDrone2021 dataset, and Figure 1 It can be seen that the color of the lower left corner of the data set target box square is darker, which indicates from one aspect that the number of small targets accounts for a large proportion in the unmanned aerial vehicle aerial image dataset, further indicating the importance of small target detection in aerial images. In the general YOLOv8 target detection algorithm, the input size of the detection head for detecting small targets also adopts a 1 / 8 sampling method, but after the detection head is down-sampled, the spatial information of the detection target feature map is reduced, which also damages the feature representation of small targets, which is very unfavorable for the detection of small targets in the unmanned aerial vehicle aerial image.

[0008] In order to retain more detailed information of small targets for small target detection, the present application designs a larger resolution target detection head to detect tiny targets missed by other detection heads. For example,Figure 2 As shown, the input of the added micro target detection head of the present application comes from the underlying feature map with greater resolution, the feature map size of the 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 be more accurately positioned in the small target of the aerial image. However, after adding one target detection head, the parameter quantity and calculation quantity of the model also increase, and there is also a phenomenon of overlapping detection information between the four detection heads in the detection model. Further, according to the characteristics that the larger targets are relatively less in the unmanned aerial vehicle aerial image data set, the present application retains the proposed simplified 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 small targets in the model of the unmanned aerial vehicle image, but also reduces the model size.

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

[0010]

[0011]

[0012] Step two: design a TP-Fusion module to fuse the detail information of small targets

[0013] The feature extractor in the general target detection YOLOv8 model extracts features and filters part of the noise through downsampling, which also causes the loss of target information. Although the general YOLOv8 fuses the single-layer information close to the bottom layer in the form of a feature pyramid, enriching the final features, these feature information is still insufficient for the detection of small targets. As can be seen from Figure 3 The dark thick dashed line that keeps a downward trend in the figure shows 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 application proposes a TP-Fusion module to further fuse the features close to the input layer bottom layer.

[0014] The small path feature fusion module (Tiny Path Fusion, TP-Fusion) proposed by the present application, as shown in the figure, fuses the information of three feature maps, and the feature map of the bottom layer (Underlying Feature Map) is added to the feature map of the second layer (Middle Layer Feature Map) after a 3x3 deformable convolution, and then spliced in the channel with the feature map of the highest layer (Top Layer Feature Map). Figure 4

[0015] ​The deformable convolution is characterized in that:

[0016]

[0017] And the traditional convolution is defined as:

[0018]

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

[0020] Among the three feature maps, the feature map with the largest size and the least number of channels contains the most detailed information for small target detection; the feature map with the smallest size and the most number of channels has the most abundant semantic information for small target classification. The present application fuses the three feature maps, which not only retains the detailed information of the small target to a large extent, but also obtains relatively abundant semantic information, which provides stronger support for the subsequent decision of the target detection head. Further, the deformable convolution technology introduced by the present application not only changes the size and the 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 the deformable parameters, so that the input features at different positions have different influences on the output, so that the detector can better capture the detailed information in the image and adapt to targets of different sizes and shapes.

[0021] Step three: introduce CA attention module to resist the interference of background noise

[0022] Since the military unmanned aerial vehicle cruises at a high altitude, the field of view of the aerial image obtained is wide, and there are many chaotic background information in the aerial image. Therefore, the target detection of the unmanned aerial vehicle aerial image is actually a target detection in a complex background. To this end, the present application introduces a coordinate attention module (Coordinate Attention Block, CA), as shown in Figure 5 It helps the network model to only focus on the foreground target to be detected and ignore the interference of complex background noise by assigning different weights to different parts of the feature map, further improving the processing capacity of information. CA is a plug-and-play lightweight attention module, which embeds the position information of the feature into the channel attention, so that the network can pay attention to a larger range of context information, while avoiding a large amount of calculation overhead.

[0023] Firstly, the coordinate attention module CA introduced in the general YOLOv8 algorithm can obtain the attention in the width and height of the image by performing global average pooling on the input feature map in the width and height directions, and encode the accurate position information of the target; then, the two direction feature maps obtained in the previous step are spliced together, and then subjected to convolution, batch normalization and nonlinear activation function operations to obtain an intermediate feature map; then, the intermediate feature map is split into high and wide direction tensors, and after convolution and Sigmoid activation function, the height direction and width direction coordinate attention are obtained; 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 in the application, the model can consciously filter the feature information after multi-layer fusion, highlight the small target region feature information, and eliminate the background noise interference.

[0024] Step four: design a bottleneck module to enhance the semantic information of the target

[0025] The two tasks of target detection are positioning and classification. For positioning, only when the detection box is positioned within a certain range and the target in the detection box is correctly classified, the target is considered to be successfully detected. For classification, whether the target can be correctly classified depends largely on whether the semantic information of the target is sufficient. When the semantic information of the target is sufficient, the probability of correct classification of the target is greater. Compared with the normal size target in the conventional video image, the small target in the unmanned aerial vehicle aerial image has less coverage and low resolution. If a general YOLOv8 model is used for target classification, it is difficult to learn sufficient semantic information from it. In order to compensate for the insufficient feature extraction caused by the deletion of the medium target detection head and the large target detection head, a method for enhancing the semantic information of the target is used to solve it. For this purpose, a bottleneck module composed of two 3x3 convolution modules is designed to be embedded in front of the detection layer of the model, as shown in Figure 6 After the input feature map passes through the first 3x3 convolution module, the channel number of the model is reduced to half of the original; after passing through the second 3x3 convolution module, the channel number of the model is restored to the size of the input channel number; on the basis of improving the accuracy of the features and target classification of the model, the calculation amount of the model is further reduced and the time delay caused by the newly added module is shortened.

[0026] In addition, the application 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] In the formula, x represents a tensor input to the activation function.

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

[0031] Step five: integrate and fuse the modules proposed in steps one to four into the general YOLOv8 target detection algorithm to obtain a high-precision lightweight UAV image target detection algorithm based on double prediction heads

[0032] This step takes the general YOLOv8 algorithm as the baseline model, integrates and fuses the modules of steps one to four, and completes the high-precision lightweight UAV image target detection algorithm based on double prediction heads, improving the detection effect of small targets in UAV aerial images.

[0033] Step six: training and adjusting training parameters to obtain the optimal weight of the high-precision lightweight UAV image target detection algorithm based on double prediction heads

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

[0035] (1) Set the input image size

[0036] The resolution size of the UAV aerial image will affect the hardware platform GPU computation and memory consumption during 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 the performance of the model. On the other hand, if the input resolution of the aerial image is too large during training, the computation and memory consumption may increase, resulting in a decrease in the training and prediction speed of the model, and even exceeding the hardware limit. In summary, according to the characteristics of the UAV aerial image and the hardware platform, the input image resolution of the present application is set to 960x960, 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 input into the model at the same time during the training process. The size of Batch Size will affect the effect and speed of training the model on the hardware platform, neither too large nor too small. If Batch Size is too large, the gradient direction of each update is too smooth, which may exceed the GPU memory capacity, leading to overfitting or reduced generalization ability of the training; if Batch Size is too small, the gradient direction of each update is unstable, which cannot fully utilize the parallel computing capability of the GPU, leading to non-convergence or slow convergence of the training. Considering the GPU memory capacity of the trained model, the Batch Size of the present application is set to 2.

[0039] (3) Set the size of the impulse

[0040] Impulse (Momentum) is a technique for optimizing gradient descent, which can accelerate the convergence of the model, avoid falling into local optimum or saddle point, and improve the generalization ability of the model. The size of the impulse determines the influence of the last update direction on the current update direction, which is generally between 0 and 1, and the commonly used value is 0.9. The impulse of the present application is set to 0.937.

[0041] (4) Set the learning rate

[0042] Learning rate (Learning rate) is a hyperparameter that adjusts the step size of parameter update in 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, the convergence may be too slow or fall into local optimum. The initial learning rate of the present application is set to 0.01, and the final learning rate is set to 0.001.

[0043] (5) Set the training rounds

[0044] Training rounds (epochs) refer to single training iterations of all batches in forward and backward propagation. An epoch means that each sample in the training dataset has the opportunity to update the internal model parameters. The size and quality of the training dataset determine how much information the model needs to learn, and the difficulty and effect of learning. According to the learning difficulty and size of the used dataset, the present application sets the training epochs of VisDrone2021 dataset to 200 and sets the training epochs of UAVDT dataset to 15.

[0045] Step seven: load the designed model and evaluate the model detection performance using aerial data set

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

[0047] Advantages

[0048] In view of the problems such as small target to be detected in the unmanned aerial vehicle image and complex image background, the application is designed and optimized in combination with the task characteristics of unmanned aerial vehicle image target detection on the basis of the general YOLOv8 algorithm, and a YOLOv8 detection algorithm with higher precision and lighter weight is proposed. First, the general three target detection heads in the model are replaced by double target detection heads more suitable for the unmanned aerial vehicle scene, which effectively improves the detection precision of small targets of the model and greatly reduces the model parameter quantity. Then, a small path feature fusion module is designed, a bottleneck module is designed, and a CA attention module is introduced, which further enhances the detection performance of the model in the unmanned aerial vehicle scene. A large number of experiments are carried out on two benchmark data sets, and the results show that, compared with the baseline model, the model of the application improves the AP of the VsiDrone2021 data set by 4.5%, improves the AP by 5.2%, and reduces the parameter quantity by 33.0%, which has considerable competitiveness in the actual application of unmanned aerial vehicle target detection. s Advantages BRIEF DESCRIPTION OF DRAWINGS

[0049] Figure 1 : Pixel distribution diagram of VisDrone data set

[0050] Figure 2 : Network structure diagram of unmanned aerial vehicle image target detection algorithm based on double prediction head high precision and light weight

[0051] Figure 3 : Diagram of the trend of detail information and semantic information of targets of different sizes and model layers

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

[0053] Figure 5 : CA attention module structure diagram

[0054] Figure 6 : Diagram of bottleneck structure module

[0055] Figure 7 : Overall step flowchart

[0056] Figure 8 : Comparison diagram of results of baseline algorithm and algorithm of the application on VisDrone2021 data set DETAILED DESCRIPTION

[0057] The application will be further described in combination with embodiments and drawings:

[0058] The specific implementation is as follows:

[0059] Step one: design the detection head of YOLOv8 algorithm to better match the aerial target detection in the unmanned aerial vehicle scene

[0060] To solve the scale variation problem of targets in conventional real scenes, the target detection model of the YOLO series adopts three detection heads to detect targets of different sizes. The three detection heads used have different sizes, and relative to the size of the input image, the feature maps input to the three detection heads are down-sampled to 1 / 8, 1 / 16 and 1 / 32 of the original size, respectively, for detecting small, medium and large targets. Figure 1 is the target box size distribution result of the VisDrone2021 dataset, from Figure 1 It can be seen that the color of the lower left corner of the data set target box grid is darker, which indicates from one aspect that the number of small targets occupies a large proportion in the unmanned aerial vehicle aerial image dataset, further illustrating the importance of small target detection in aerial images. In the general YOLOv8 target detection algorithm, the input size of the detection head for detecting small targets also adopts a 1 / 8 sampling method, but after the detection head is down-sampled, the spatial information of the target feature map is reduced, which also damages the feature representation of small targets, which is very unfavorable for the detection of small targets which account for the majority in unmanned aerial vehicle aerial images.

[0061] In order to retain more detailed information of small targets for small target detection, the present application designs a larger resolution target detection head to detect tiny targets missed by other detection heads. As Figure 2 shown, the input of the tiny target detection head added by the present application comes from the bottom layer feature map with larger resolution, and the feature map size of this detection head is twice that of the small target detection head, which can retain more detailed information such as edges and contours of the image, so that the detector needle can more accurately locate the small targets in the aerial image. However, after adding a target detection head, the parameter quantity and computational quantity of the model also increase, and there is also a phenomenon of overlapping detection information between the four detection heads in the detection model. Further, according to the characteristics of the unmanned aerial vehicle aerial image dataset that larger targets are relatively few, the present application retains the proposed simplified 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 small targets in the unmanned aerial vehicle image by the model, but also reduces the model size.

[0062] Step two: design a TP-Fusion module to fuse the detailed information of small targets

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

[0064] The Tiny Path Fusion (TP-Fusion) module proposed by the present application, as shown in Figure 4 The module fuses the information of three layers of feature maps. The feature map of the bottom layer (Underlying Feature Map) is added to the feature map of the second layer (Middle Layer Feature Map) after a 3x3 deformable convolution, and then spliced with the feature map of the highest layer (Top Layer Feature Map) in the channel.

[0065] The deformable convolution is characterized in that:

[0066]

[0067] The traditional convolution is defined as:

[0068]

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

[0070] Among the three feature maps, the one with the largest feature map and the least number of channels contains the most detailed information for small target detection; the one with the smaller feature map and the most number of channels has the most abundant semantic information for small target recognition. By fusing the three feature maps, the present application not only retains the detailed information of small targets to a large extent, but also obtains relatively abundant semantic information, which provides stronger support for the subsequent decision of the target detection head. Further, the deformable convolution technology introduced by the present application not only changes the size and the 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 the deformable parameters, so that the input features at different positions have different effects on the output, thereby enabling the detector to better capture the detailed information in the image and adapt to targets of different sizes and shapes.

[0071] Step three: introducing a CA attention module to resist the interference of background noise

[0072] Since the military unmanned aerial vehicle cruises at a high altitude, the field of view of the aerial image obtained is wide, and there is much chaotic background information in the aerial image. Therefore, the target detection of the unmanned aerial vehicle aerial image is actually a target detection in a complex background. In this regard, the present application introduces a coordinate attention module (Coordinate Attention Block, CA), as shown in Figure 5 which helps the network model to focus only on the foreground target to be detected and ignore the interference of complex background noise, thereby further improving the processing capacity of information. The CA is a plug-and-play lightweight attention module, which embeds the position information of the feature into the channel attention, so that the network can focus on a larger range of context information while avoiding a large amount of computational overhead.

[0073] Firstly, the coordinate attention module CA introduced by the present application in the general YOLOv8 algorithm can perform global average pooling on the input feature map in the width and height directions, thereby obtaining attention in the width and height directions of the image and encoding the accurate position information of the target; secondly, the two direction feature maps obtained in the previous step are spliced together, and then subjected to convolution, batch normalization and nonlinear activation function operations, thereby obtaining an intermediate feature map; then, the intermediate feature map is split into high and wide direction tensors, which are respectively subjected to convolution and Sigmoid activation function, thereby obtaining the coordinate attention in the height direction and the width direction; finally, the input feature map is multiplied by the obtained coordinate attention weight, thereby obtaining a feature map with attention weight in the height and width directions. By introducing the CA into the feature fusion module proposed by the present application, the feature information after multi-layer fusion of the model can be consciously screened, thereby highlighting the small target region feature information while eliminating the background noise interference.

[0074] Step four: design a bottleneck module to enhance the semantic information of the target

[0075] The two 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 sufficient. When the semantic information of the target is sufficient, the probability of correct classification of the target is greater. Compared with the normal-sized target in the conventional video image, the small target in the unmanned aerial vehicle image has fewer pixels and lower resolution. If a general YOLOv8 model is used for target classification, it is difficult to learn sufficient semantic information from it. In order to compensate for the insufficient feature extraction caused by the deletion of the medium target detection head and the large target detection head, the method of enhancing the semantic information of the target is used to solve it. For this purpose, the present application designs a bottleneck module composed of two 3x3 convolution modules embedded in front of the detection layer of the model, as shown in Figure 6 After the input feature map passes through the first 3x3 convolution module, the number of channels of the model is reduced to half of the original; after passing through the second 3x3 convolution module, the number of channels of the model is restored to the size of the input channel number; on the basis of improving the features of the model and the accuracy of target classification, the computational complexity of the model is further reduced and the time delay caused by the newly added module is shortened.

[0076] In addition, the present application 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] In the formula, x represents a tensor input into the activation function.

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

[0081] Step five: integrate and fuse the modules proposed in steps one to four into the general target detection YOLOv8 algorithm to obtain a DPH-YOLOv8 design model

[0082] The step takes the general YOLOv8 algorithm as a baseline model, integrates and fuses the modules of steps 1 to 4, and completes the high-precision lightweight UAV image target detection algorithm based on the double prediction head, thereby improving the small target detection effect of the UAV aerial image.

[0083] Step six: training and adjusting training parameters to obtain the optimal weight of DPH-YOLOv8

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

[0085] (1) Set the input image size

[0086] The resolution size of the UAV aerial image will affect the hardware platform GPU calculation amount and memory consumption during 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 the performance of the model. On the other hand, if the input resolution of the aerial image during training is too large, the calculation amount and memory consumption may increase, resulting in a decrease in the training and prediction speed of the model, and even exceeding the limit of the hardware. In summary, according to the characteristics of the UAV aerial image and the hardware platform, the input image resolution of the present application is set to 960x960, which can achieve a balance between efficiency and accuracy.

[0087] (2) Set the training batch size

[0088] The batch size refers to the number of data samples input into the model at the same time during each training process. The size of the batch size will affect the effect and speed of the hardware platform in training the model, and it cannot be too large or too small. If the batch size is too large, the gradient direction updated each time is too smooth, which may exceed the GPU memory capacity, resulting in overfitting or a decrease in the generalization ability during training. If the batch size is too small, the gradient direction updated each time is unstable, and the parallel computing capability of the GPU cannot be fully utilized, resulting in non-convergence or slow convergence speed during training. Considering the GPU memory capacity of the trained model, the batch size of the present application is set to 2.

[0089] (3) Set the momentum size

[0090] Momentum is a technique for optimizing gradient descent, which can accelerate the convergence of the model, avoid falling into local optimum or saddle point, and improve the generalization ability of the model. The size of the momentum determines the influence of the update direction of the last time on the current update direction, and is generally between 0 and 1. The commonly used value is 0.9, and the momentum of the present application is set to 0.937.

[0091] (4) Set the learning rate

[0092] The learning rate is a hyperparameter that adjusts the step size of parameter updates in each iteration of the optimization algorithm. The size of the learning rate affects the convergence speed and effectiveness of the model. If the learning rate is too large, it may cause the loss function to oscillate or diverge; if the learning rate is too small, it may cause slow convergence or fall into 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 the training epochs

[0094] Training epochs refer to a single training iteration of all batches in forward and backward propagation. An epoch means that each sample in the training dataset has the opportunity to update the internal model parameters. The size and quality of the training dataset determine how much information the model needs to learn, and the difficulty and effectiveness 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 seven: Load the designed model and evaluate the model's detection performance using aerial data sets

[0096] Through training, the weight best.pt is obtained, and the performance of the model is evaluated using the VisDrone2021 and UAVDT two unmanned aerial vehicle aerial data sets.

[0097] The overall steps of the algorithm are shown in the attached Figure 7 .

[0098] First, according to the characteristics of the small target proportion in the unmanned aerial vehicle aerial image and the tight memory resources of the unmanned aerial vehicle device, the detection head of YOLOv8 algorithm is designed;

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

[0100] Third, in view of the problem of complex background of unmanned aerial vehicle image, CA attention module is introduced;

[0101] Fourth, the 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 the design model of the unmanned aerial vehicle image target detection algorithm based on double prediction head high-precision lightweight;

[0103] Sixth, start training and gradually adjust the training parameters to obtain the optimal weight of the unmanned aerial vehicle image target detection algorithm based on the double prediction head high-precision lightweight;

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

[0105] 1. Experimental conditions

[0106] To test the performance of the unmanned aerial vehicle image target detection algorithm based on the double prediction head high-precision lightweight, the present application uses the VisDrone2021 dataset to perform ablation experiments and comparative experiments of different algorithms, and uses the UAVDT dataset to perform comparative experiments. VisDrone2021 and UAVDT are both public unmanned aerial vehicle image datasets with high usage frequency, suitable for studying target detection tasks based on unmanned aerial vehicle images.

[0107] Experimental test environment: the system is Ubuntu20.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 application first optimizes the YOLOv8 model from four different angles, 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 algorithm structure after design is as shown in Figure 2 .

[0110] To verify the superiority of the unmanned aerial vehicle image target detection algorithm based on the double prediction head high-precision lightweight in detection accuracy and model size, ablation experiments and comparative experiments of different algorithms are performed on the VisDrone2021 dataset, and comparative experiments are performed on the UAVDT dataset.

[0111] 3. Classification evaluation index

[0112] The detection quality evaluation index used in the present application is the mean average precision (Mean Average Precision, mAP). The basic concepts involved in the calculation of mAP will be introduced below.

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

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

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

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

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

[0118] (6) Precision: Accuracy, also known as precision. It represents the proportion of samples that are truly positive in the output results of the detection that are positive samples, and the calculation formula is:

[0119]

[0120] (7) Recall: Recall, also known as recall. It represents the proportion of truly identified positive examples among all positive examples, and the calculation formula is:

[0121]

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

[0123]

[0124] When detecting multiple classes of targets, the average of the AP of all classes can be obtained to measure the average average precision (mAP) of the multi-class target detection quality, and the calculation formula is as follows:

[0125]

[0126] In the formula, 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 requirement, can it be applied to industry, which is of great significance to realize real-time detection of the unmanned aerial vehicle. A commonly used measurement index is frame per second (FPS), that is, the number of images that can be processed per second, which is defined as follows:

[0128]

[0129] In the formula, Tot represents the total time consumed when the image or video is detected; and FC refers to the number of image frames processed. Generally, the FPS obtained by the target detection algorithm using different hardware configurations is quite different.

[0130] 4. Simulation test

[0131] In order to test and verify the detection performance of the algorithm proposed in the application on the image of the unmanned aerial vehicle, experiments are performed on the VisDrone2021 dataset and the UAVDT dataset, and actual detection result visualization analysis is performed.

[0132] 1. Ablation experiment on the VisDrone dataset

[0133] (1) Analysis of the detection result of the optimized detection head

[0134] The application performs experiments on two modified detection head modes, including adding a micro target detection head in the baseline model and removing the two detection heads with small general feature maps of the model on the basis of the above modification. As can be seen from the data in Table 2, both of the two modified detection head modes improve the detection accuracy of the model, and the detection accuracy of small targets is improved most. Among them, the first modification mode improves the AP and AP 50 by 1.9% and 2.3% respectively, and the AP values for small targets and medium targets are improved by 2.4% and 0.9% respectively, and the second modification mode improves the AP and AP 50The accuracy was improved by 1.9% and 2.6% respectively, with AP values ​​for small and medium targets improving by 2.5% and 1.2% respectively. This indicates that the micro-target detection head designed in this invention is helpful for detecting a large number of small targets in UAV images. The first modification method has slightly lower detection accuracy than the second. The reason is that the combined effect of the newly added micro-target detection head and the general small target detection head in the model can detect small targets in UAV images relatively comprehensively. The two detection heads with smaller feature maps in the first method misjudged small targets, leading to incorrect decisions about them and ultimately affecting the model's final decision. Furthermore, small targets constitute a large portion of the total targets in the VisDrone2021 dataset, thus slightly reducing the model's detection accuracy. In addition, both modifications to the detection head improved the AP accuracy for large targets. l This had a negative impact. The first modification method, in particular, affected AP. l The accuracy was reduced by 1.6%. The reason for this is that the introduced small target detection head has high resolution, a large feature map, and a small receptive field. This results in inaccurate localization of large targets during bounding box generation, thus affecting the model's final decision on large targets. The second modification method improved AP... l The 1.5% reduction is attributed to the increased false negative rate for large targets after removing two detectors with smaller feature maps. In terms of detection efficiency, both modifications reduced the number of model parameters: the first reduced it by 5.4%, and the second by 42.9%. Furthermore, the second modification had a significantly smaller impact on the model's detection speed compared to the first.

[0135] Table 2 shows the ablation experiments on the VisDrone dataset.

[0136]

[0137]

[0138] In summary, this invention improves the detection accuracy of UAV images by adding a small target detection head and removing two detection heads with smaller feature maps, thereby significantly reducing the model size and having almost no impact on the model's detection speed.

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

[0140] The rich detail contained in the low-level feature maps helps the model locate the target. According to row 4 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 metrics improved, including AP. sThe detection precision of small targets is improved by 1.0%, which indicates that the TP-Fusion module proposed in the application helps to improve the detection precision of small targets.

[0141] (3) Analysis of the designed bottleneck module

[0142] Adding the bottleneck module in front of the detection head helps the model to further extract features and make up for the loss of semantic information caused by discarding two detection heads. As can be seen from Table 2, the increase of each precision index after adding the bottleneck module is between 1.0% and 2.2%, among which the detection precision of large targets is improved by 2.2%, which is more than the 1.5% reduction of AP l caused by removing the two smaller feature maps of the detection head, and has little effect on the parameter quantity and detection speed of the model. This indicates that the bottleneck module introduced in the application solves the problem of reduced detection precision of large targets caused by discarding two detection heads without affecting the detection efficiency of the model.

[0143] (4) Result analysis 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 targets that need to be focused on. The effect of adding the CA attention module is shown in the last row of Table 2. Except for AP l , each detection precision index is improved by about 1.0%, while AP l is slightly reduced.

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

[0146] The network proposed in the application and different algorithms are compared on the VisDrone2021 dataset and the UAVDT dataset, including the YOLO series target detection algorithm and other target detection algorithms for unmanned aerial vehicle 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. The model of the application achieves the best results on all indicators except AP l . In AP lThe index lags behind ClusDet and DMNet, but the two algorithms need to perform coarse detection and fine detection on the target and fuse the target frame, which delays the inference speed, while the algorithm of the application is an end-to-end one-stage target detection algorithm, which has an advantage in detection real-time and is more suitable for deployment on unmanned aerial vehicles with high time delay requirements.

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

[0149] Method AP (%) AP 50 (%)]] AP 75 (%)]] AP s (%)]] AP m (%)]] 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, and the AP value of the algorithm of the application exceeds all the compared algorithms, proving that the algorithm of the application has good generalization.

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

[0152] Method AP (%) AP 50 (%)]] 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, Actual detection result visual analysis

[0154] In order to more intuitively compare the performance gap between the algorithm of the application and the baseline model, the prediction results of the two models on the VisDrone2021 test set are visualized, as shown in Figure 8 The upper half is the detection result of the baseline model, and the lower half is the detection result of the model of the application. At the same time, in order to make the visual effect better, the target class and the confidence score in the detection result are removed, and the area that can highlight the performance gap of the two models is enlarged and displayed. From the figure, it can be seen that compared with the baseline model, the first scene, the application method detects more small vehicle targets at a long distance, the second scene, the application method detects more small pedestrian targets at a long distance, the third scene, the application method detects more small vehicle targets at a long distance, and the fourth scene, the application method detects more small vehicle targets in the case of dim light. Therefore, the algorithm of the application can improve the detection rate of the model for targets with a longer line-of-sight distance and a smaller volume, and has a good detection effect on unmanned aerial vehicle images with a complex background and small targets.

[0155] In summary, the algorithm of the application can improve the target detection accuracy in the unmanned aerial vehicle scene and reduce the size of the model, and has considerable competitiveness in practical application.

Claims

1. A high-precision, lightweight UAV image target detection algorithm based on dual prediction heads, characterized in that, Includes the following steps: Step 1: Construct a high-precision, lightweight YOLOv8 target detection algorithm; Step 1-1: Simplify the prediction head of the YOLOv8 algorithm; add a small target prediction head with a larger feature map resolution to the general YOLOv8 algorithm, and remove the two prediction heads with smaller feature map resolution in the general YOLOv8 algorithm to obtain a dual-prediction head YOLOv8 target detection algorithm. Steps 1-2: Add a Tiny Path Fusion (TP-Fusion) module to the neck network of the general YOLOv8. This module fuses information from three layers of feature maps. The feature map of the bottom layer (Underlying Feature Map) is added to the feature map of the second layer (Middle Layer Feature Map) after passing through a 3×3 deformable convolution, and then concatenated with the feature map of the top layer (Top Layer Feature Map) channel by channel. Steps 1-3: Continue by adding a 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 the width and height directions to obtain attention in the image width and height, and encodes precise location information. Then, the feature maps obtained in the previous step are concatenated together, and then processed through convolution, batch normalization, and non-linear activation functions to obtain an intermediate feature map. The intermediate feature map is then split into tensors in the height and width directions, and processed through convolution and sigmoid activation functions respectively to obtain coordinate attention in the height and width directions. Finally, the input feature map is multiplied by the obtained attention weights to obtain the final feature map with attention weights in the height and width directions. Steps 1-4: Continue to add a bottleneck module to the head network of the general YOLOv8; this module consists of two 3×3 convolutional modules; after the input feature map passes through the first convolutional module, the number of channels is reduced to half of the original number, and after passing through the second convolutional module, the number of channels is restored to the input size, and then residual connections are made with the input feature map, and finally the final output is obtained through the HSwish activation function; Steps 1-5: Using the general YOLOv8 algorithm as the baseline model, integrate and fuse the modules from Steps 1-1 to 1-4 respectively to complete the target detection model for UAV aerial images—a dual-predictor head high-precision lightweight UAV image target detection algorithm, which improves the small target detection effect of UAV aerial images. Step 2: Train a high-precision, lightweight UAV image target detection algorithm with dual prediction heads; set training parameters: input image size, batch size, impulse size, learning rate, and maximum number of iterations; train the network using the VisDrone2021 and UAVDT datasets; The training process yields the final dual-predictor head high-precision lightweight UAV image target detection algorithm. Step 3: Load the dual-predictor high-precision and lightweight UAV image target detection algorithm and evaluate the algorithm.

2. The UAV image target detection algorithm based on dual prediction heads with high precision and lightweight design as described in claim 1, characterized in that, The high-precision and lightweight design means that, in the scenario of target detection in UAV aerial images, the accuracy value is improved compared to the baseline network, while the number of parameters is reduced compared to the baseline network.

3. The high-precision, lightweight UAV image target detection algorithm based on dual prediction heads according to claim 1, characterized in that, The deformable convolution is defined as: The traditional definition of convolution is: Where 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 p is the weight of the convolution kernel at position k. k It is the offset of p within the convolution kernel range, Δp k It is the offset introduced by deformable convolution.

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

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

Images