Unmanned aerial vehicle aerial image small target detection method based on AFLW-YOLOv5s model

Abstract

The invention discloses an unmanned aerial vehicle aerial image small target detection method based on an AFLW-YOLOv5s model. The method comprises the following steps: firstly, preprocessing and dividing a pre-acquired unmanned aerial vehicle aerial image data set; secondly, constructing an unmanned aerial vehicle aerial image small target detection model based on AFLW-YOLOv5s; according to the model, a CSM channel space attention module is fused in a backbone network, so that the model can pay attention to channel information and space information, and the feature extraction capability of the model is enhanced; an LWinvotion module is introduced, so that the model space and channel self-adaptive capability is improved, and the model expression effect is enhanced; an MSA-BiFPN neck network structure is added, so that multi-scale feature fusion of the model is realized; then, a small target detection head ADH is introduced into a target detection layer, so that the detection capability of the model on a small target image is improved; according to the AFLW-YOLOv5s model provided by the invention, the detection precision of small target images aerial-photographed by the unmanned aerial vehicle is remarkably improved, and the method has a good application prospect.

Description

Technical Field

[0001] The present invention relates to a small target detection method for unmanned aerial vehicle (UAV) aerial images based on an AFLW-YOLOv5s model. Background Art

[0002] Computer vision technology has recently become increasingly important. Computer vision captures images or videos using devices like cameras and processes them using computers to extract the information humans need. Object detection, a core subtask of computer vision, aims to locate and identify specific objects in images or videos. Object detection in drone aerial imagery, a key branch of this field, has been widely applied in numerous real-world scenarios and has become a hot research topic.

[0003] Furthermore, drones, thanks to their low cost, lightweight design, high-resolution cameras, and flexible operation, are seeing significant applications in traffic monitoring, power inspections, military patrols, environmental regulation, agricultural management, and fire monitoring. Real-time drone monitoring not only effectively reduces labor risks and reduces the workload of both simple and complex tasks, but also significantly improves regulatory efficiency and promotes intelligent upgrades in related industries. The advancement of drone aerial imagery object detection technology provides strong technical support for the digital transformation of various sectors of society. However, due to the high altitude of drones, the target image occupies a relatively small proportion, resulting in insufficient feature information. This makes the detection of small targets more difficult and results in lower accuracy. Summary of the Invention

[0004] Purpose of the invention: The purpose of the present invention is to provide a small target detection method for UAV aerial images based on the AFLW-YOLOv5s model, which improves the difficulty and detection accuracy of small targets in UAV aerial images.

[0005] Technical solution: The method for detecting small targets in drone aerial images based on the AFLW-YOLOv5s model of the present invention comprises the following steps:

[0006] (1) Preprocess the labels in the pre-acquired Visdrone UAV aerial image dataset and convert the label files in the dataset into YOLO format;

[0007] (2) Perform data augmentation on the dataset images after label conversion, and then divide the dataset of images after data augmentation;

[0008] (3) Improve the YOLOv5s model and build a small target detection model for drone aerial images based on the AFLW-YOLOv5s model; implement the MSA-BiFPN neck network structure for multi-scale feature fusion; and capture the small target detection head ADH for small-scale objects;

[0009] (4) Input the training set images into the AFLW-YOLOv5s model for training to obtain the pre-trained AFLW-YOLOv5s model weights; use the pre-trained AFLW-YOLOv5s model weights to detect small targets in drone aerial images.

[0010] Furthermore, the implementation process of converting the label file in the dataset into the YOLO format in step (1) is as follows: first, obtain the label file of the VisDrone dataset in the annotations folder; secondly, traverse each annotation file and obtain the size of the image corresponding to the label; then, by opening the annotation file and obtaining the content, convert the label file into the xywh format of YOLO; finally, write the converted label into a new label folder.

[0011] Furthermore, the data enhancement processing in step (2) includes image HSV enhancement, image translation and scaling, image left-right flipping, Mosaic enhancement, MixUp enhancement, and Copy-Paste enhancement.

[0012] Furthermore, the adjustment range of the image HSV enhancement is ±0.015; the adjustment range of the image translation is ±0.1; the adjustment range of the image scaling is ±0.9; the probability of image left-right flipping is 0.5; the probability of Mosaic enhancement is 1.0, multiple images are stitched into one image; the probability of MixUp enhancement is 0.1, two images are mixed; the probability of Copy-Paste enhancement is 0.1, the target is copied from one image and pasted to another image.

[0013] Furthermore, in step (2), the image dataset after data enhancement is divided into a training set, a validation set and a test set.

[0014] Furthermore, the detection model in step (3) includes a CSM module for enhancing channel and spatial feature extraction capabilities, and an LWinvolution module for improving model spatial and channel adaptability.

[0015] Furthermore, the implementation process of the CSM module for enhancing the channel and spatial feature extraction capability is as follows: first, a CSM module for enhancing the channel and spatial feature extraction capability is introduced into the backbone network of the YOLOv5s model, and then the input feature map is sequentially fed into a 1×1 convolution and a 3×3 convolution for partial feature extraction; secondly, the channel feature information, spatial feature information and global feature information of the input feature map are extracted through two dual-branch structures; then the initial feature map is input into a 3×3 convolution for global feature extraction.

[0016] Furthermore, the two double-branch structures are:

[0017] The first dual-branch structure: First, a 1×1 convolution is used to extract partial features from the feature map. Then, the global average pooling layer and the global maximum pooling layer are used to extract channel information. Then, a 1×1 convolution and a ReLU activation function are used to output the channel attention weights. After weighted feature fusion, the sigmoid function is used to multiply the partial features originally extracted by the 1×1 convolution.

[0018] The second dual-branch structure: a 1×1 convolution is used to extract partial features, followed by extracting spatial information through the global average pooling layer and the global maximum pooling layer respectively, adding the extracted spatial information, and then outputting the spatial attention weight through a 7×7 convolution and a Sigmoid function, which is multiplied channel by channel with the original partial feature map extracted by a 1×1 convolution, realizing the model's ability to extract channel information and spatial information of the input feature map.

[0019] Furthermore, the implementation process of the LWinvolution module for improving the spatial and channel adaptability of the model is as follows: an LWinvolution module is integrated into the neck module of the model. The LWinvolution module first obtains the shape of the input feature map (B, C, H, W), where B represents the batch size, C represents the number of input feature channels, and H and W represent the height and width of the feature map; secondly, the input feature map is respectively passed through a 1×1 convolution, a ReLU activation function, and a 1×1 convolution to generate a dynamic convolution kernel with a shape of (B, G×K×K, H, W); then, the reshape operation is performed The generated dynamic convolution kernel shape is adjusted to (B, G, K × K, H, W), which contains G groups of K × K convolution kernels, and each (H, W) position has different kernel parameters; then a group convolution operation is performed to group the input features into groups, with a shape of (B, G, GC, H, W), which contains G groups of features, and each group has GC channels; then the dynamic convolution of each group is calculated, and the calculated feature map shape is grouped and multiplied with the feature map shape after the original grouping; the results of all group calculations are concat spliced; finally, the spliced ​​results are input into a 3 × 3 DWConv and a 1 × 1 PWConv.

[0020] Furthermore, the implementation process of the MSA-BiFPN neck network structure for realizing multi-scale feature fusion and introducing the small target detection head ADH in the target detection layer in step (3) is as follows: first, the BiFPN neck network structure is improved, and the network cross-layer connects the 20×20, 40×40, and 80×80 feature maps input by the Backbone backbone network with the feature maps of the same size in the BiFPN network; an upsampling operation is added after the 80×80 size feature map in the BiFPN network; then, the feature map size is converted to 160×160, fused with the feature map of the same scale of the Backbone, and a downsampling operation is added to match it; finally, the 160×160 size feature map information is passed to the 160×160 size detection head.

[0021] Beneficial effects: Compared with the existing technology, the present invention has the following significant features: the present invention addresses the problems of high flying altitude of UAVs, small proportion of targets in images, insufficient feature information, which leads to increased difficulty in detecting small targets and low detection accuracy, and thus constructs a small target detection model for UAV aerial images based on AFLW-YOLOv5s; the present invention incorporates the CSM module into the backbone network of the model, so that the model can focus on channel information and spatial information, and enhance the model's feature extraction capability; the present invention introduces the LWinvolution module into the neck network of the model, which improves the model's spatial and channel adaptability and enhances the model's expression effect; the present invention designs the MSA-BiFPN neck network structure and introduces the small target detection head ADH into the target detection layer, which realizes the multi-scale feature fusion of the model and improves the model's detection capability for small target images. Compared with the existing methods, the AFLW-YOLOv5s model proposed in the present invention significantly improves the accuracy of small target detection in UAV aerial images and has good application prospects. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] Figure 1 It is an operational flow chart of the present invention;

[0023] Figure 2 It is a schematic diagram of the AFLW-YOLOv5s model structure in the present invention;

[0024] Figure 3 It is the CSM module structure diagram of the present invention;

[0025] Figure 4 It is the structural diagram of the LWinvolution module in the present invention;

[0026] Figure 5 Schematic diagram of the MSA-BiFPN neck network structure in the present invention. DETAILED DESCRIPTION

[0027] The specific technical solutions of the present invention are further described in detail below with reference to specific examples.

[0028] As shown in the figure, the method for detecting small targets in drone aerial images based on the AFLW-YOLOv5s model described in the present invention includes the following steps:

[0029] (1) Preprocess the labels in the pre-acquired Visdrone UAV aerial image dataset and convert the label files in the dataset into YOLO format;

[0030] (2) Perform data augmentation on the label-converted dataset images, and then divide the data augmented image dataset into training set, validation set, and test set;

[0031] (3) Improve the YOLOv5s model and build a small target detection model for drone aerial images based on the AFLW-YOLOv5s model; implement the MSA-BiFPN neck network structure for multi-scale feature fusion; and capture the small target detection head ADH for small-scale objects;

[0032] (4) Input the training set images into the AFLW-YOLOv5s model for training to obtain the pre-trained AFLW-YOLOv5s model weights; use the pre-trained AFLW-YOLOv5s model weights to detect small targets in drone aerial images.

[0033] The implementation process of converting the label file in the dataset into the YOLO format in step (1) is as follows: first, obtain the label file (txt format) of the VisDrone dataset in the annotations folder; secondly, traverse each annotation file and obtain the size (width, height) of the image corresponding to the label; then, by opening the annotation file and obtaining the content, convert the label file into the YOLO xywh (x_center, y_center, width, height) format; finally, write the converted label into a new label folder (labels).

[0034] The data enhancement processing in step (2) includes image HSV enhancement with an adjustment range of ±0.015; image translation with an adjustment range of ±0.1; image scaling with an adjustment range of ±0.9; image left-right flipping with a probability of 0.5; Mosaic enhancement with a probability of 1.0, stitching multiple images into one image; MixUp enhancement with a probability of 0.1, mixing two images in a certain proportion; Copy-Paste enhancement with a probability of 0.1, copying and pasting the target from one image to another.

[0035] like Figure 2As shown, the detection model in step (3) includes a CSM module for introducing a feature extraction capability for enhancing the model's focus on channel information and spatial information; and an LWinvolution module for integrating and improving the model's spatial and channel adaptation capabilities.

[0036] The implementation process of the CSM module for enhancing the channel and spatial feature extraction capability is as follows: a CSM module for enhancing the channel and spatial feature extraction capability is introduced into the backbone network of the YOLOv5s model, such as Figure 3 As shown in the figure, the working principle of this module is as follows: first, the input feature map is sent to a 1×1 convolution and a 3×3 convolution in sequence for partial feature extraction; secondly, the channel feature information, spatial feature information and global feature information of the input feature map are extracted through two dual-branch structures respectively; then the initial feature map is input into a 3×3 convolution for global feature extraction, and the feature map output by this operation is added to the final output feature map to enhance the model's feature extraction capability.

[0037] The two dual-branch structures are: the first dual-branch structure is to first use a 1×1 convolution to extract partial features of the feature map, and then respectively extract channel information through the global average pooling layer and the global maximum pooling layer, and then output the channel attention weight through a 1×1 convolution and a ReLU activation function, and after weighted feature fusion, multiply it with the partial features originally extracted by a 1×1 convolution through a sigmoid function; the second dual-branch structure is to use a 1×1 convolution to extract partial features, and then respectively extract spatial information through the global average pooling layer and the global maximum pooling layer, add the extracted spatial information, and then output the spatial attention weight through a 7×7 convolution and a Sigmoid function, and multiply it channel by channel with the original partial feature map extracted by a 1×1 convolution, so as to realize the model's ability to extract channel information and spatial information of the input feature map.

[0038] like Figure 4As shown in FIG, the implementation process of the LWinvolution module for improving the spatial and channel adaptability of the model is as follows: an LWinvolution module is integrated into the neck module of the model, which mainly improves the spatial and channel adaptability of the model; the LWinvolution module first obtains the shape of the input feature map (B, C, H, W), where B represents the batch size, C represents the number of input feature channels, and H and W represent the height and width of the feature map; secondly, the input feature map is respectively passed through a 1×1 convolution, a ReLU activation function and a 1×1 convolution to generate a dynamic convolution kernel with a shape of (B, G×K×K, H, W); then the generated dynamic convolution kernel is reshaped by the reshape operation. The shape of the dynamic convolution kernel is adjusted to (B, G, K×K, H, W), which contains G groups of K×K convolution kernels, and each (H, W) position has different kernel parameters; then the group convolution operation is performed to group the input features into groups, and the shape is (B, G, GC, H, W), which contains G groups of features, and each group has GC channels; then the dynamic convolution of each group is calculated, and the shape of the feature map after calculation is multiplied by the shape of the feature map after the original grouping; the results of all group calculations are concat spliced; finally, the spliced ​​results are input into a 3×3 DWConv (DepthwiseConvolution) and a 1×1 PWConv (Pointwise Convolution) to further extract local spatial features and reduce the computational complexity of the model.

[0039] like Figure 5 As shown, the implementation process of the MSA-BiFPN neck network structure for realizing multi-scale feature fusion and introducing the small target detection head ADH in the target detection layer in step (3) is as follows: first, in order to make full use of the original feature map information of the backbone network and avoid partial loss of this information in the neck network, the BiFPN neck network structure is improved. The network cross-layer connects the 20×20, 40×40, and 80×80 feature maps input by the Backbone backbone network with the feature maps of the same size in the BiFPN network; an upsampling operation is added after the 80×80 size feature map in the BiFPN network to expand the feature map receptive field range; then, after converting the feature map size to 160×160, it is fused with the feature map of the same scale of the Backbone, and a downsampling operation is added to match it; finally, the 160×160 size feature map information is passed to the 160×160 size detection head; the improved MSA-BiFPN network uses a multi-scale adaptive detection head to effectively solve the problem of drastic changes in target scale in drone aerial images and significantly improve the accuracy of small target detection.

[0040] In the step (4), in order to fully utilize the performance of the experimental equipment, the model training is performed on the training set, and the batch size value is set to -1 to automatically adjust the batch size. The largest feasible batch size is automatically selected according to the GPU video memory, and no manual parameter adjustment is required, thereby improving the model training efficiency. The present invention sets the number of iterations to 200 rounds and the initial learning rate to 0.01; adopts the SGD optimizer, and sets the momentum parameter of SGD to 0.937; saves the weights of the last round of model training and the optimal weights after training; and uses the optimal weights obtained after 200 rounds of training to detect small targets in drone aerial images.

[0041] The present invention adopts the commonly used standards for evaluating the performance of target detection algorithms: precision, recall, and average precision (MAP).

[0042] (41) Precision: It indicates the proportion of correctly predicted samples in the forward prediction sample set. The formula is as follows:

[0043]

[0044] Where Precision represents the percentage of correctly predicted targets among all detected targets, TP represents the correct prediction of the model, and FP represents the incorrect prediction of the model;

[0045] (42) Recall: It indicates the proportion of correctly predicted targets among all targets. The formula is as follows:

[0046]

[0047] Where Recall represents the proportion of correctly predicted targets among all targets, TP represents the correct prediction of the model, and FN represents the target that needs to be predicted but is incorrectly detected by the model.

[0048] (43) Average Precision (MAP): represents the average accuracy of all samples; the formula is as follows:

[0049]

[0050] Where AP represents the area enclosed by the curve axes formed by precision and recall, MAP represents the average of the average accuracy of all samples, and c represents the number of samples.

[0051] The specific experimental comparison results are shown in Table 1.

[0052] Table 1 Comparison results before and after improvement:

[0053] Model Precision(%) Recall (%) MAP_0.5(%) MAP_0.5:0.95(%) YOLOv5s 46.88 35.43 35.01 19.26 AFLW-YOLOv5s 53.43 41.54 42.94 24.72

[0054] From the results in Table 1, we can see that compared with the YOLOv5s model, the precision, recall, MAP_0.5 and MAP_0.5:0.95 values ​​of the AFLW-YOLOv5s model are increased by 6.5, 6.1, 7.9 and 5.4 percentage points respectively, which significantly improves the accuracy of small target detection in UAV aerial images.

[0055] This method integrates the CSM module into the backbone network to optimize the feature extraction capability of the YOLOv5s model; introduces the LWinvolution module to improve the model's spatial and channel adaptability and enhance the model's expressiveness; designs the MSA-BiFPN neck network structure and introduces the small target detection head ADH into the target detection layer to achieve multi-scale feature fusion of the model, significantly improving the model's accuracy and robustness for small target detection.

Claims

1. A small target detection method for UAV aerial images based on the AFLW-YOLOv5s model, characterized by: The following steps are involved: (1) Preprocess the labels in the pre-acquired Visdrone UAV aerial image dataset and convert the label files in the dataset into YOLO format; (2) Perform data augmentation on the dataset images after label conversion, and then divide the dataset of images after data augmentation; (3) Improve the YOLOv5s model and build a small target detection model for drone aerial images based on the AFLW-YOLOv5s model; implement the MSA-BiFPN neck network structure for multi-scale feature fusion; and capture the small target detection head ADH for small-scale objects; (4) Input the training set images into the AFLW-YOLOv5s model for training to obtain the pre-trained AFLW-YOLOv5s model weights; use the pre-trained AFLW-YOLOv5s model weights to detect small targets in drone aerial images.

2. The method for detecting small targets in UAV aerial images based on the AFLW-YOLOv5s model according to claim 1 is characterized in that: The implementation process of converting the label file in the dataset into the YOLO format in step (1) is as follows: first, obtain the label file of the VisDrone dataset in the annotations folder; second, traverse each annotation file and obtain the size of the image corresponding to the label; then, by opening the annotation file and obtaining the content, convert the label file into the YOLO xywh format; finally, write the converted label into a new label folder.

3. The method for detecting small targets in UAV aerial images based on the AFLW-YOLOv5s model according to claim 1, wherein: The data enhancement processing in step (2) includes image HSV enhancement, image translation and scaling, image left-right flipping, Mosaic enhancement, MixUp enhancement, and Copy-Paste enhancement.

4. The method for detecting small targets in UAV aerial images based on the AFLW-YOLOv5s model according to claim 3 is characterized in that: The adjustment range of the image HSV enhancement is ±0.015; the adjustment range of the image translation is ±0.1; the adjustment range of the image scaling is ±0.9; the probability of image left-right flipping is 0.5; the probability of Mosaic enhancement is 1.0, which stitches multiple images into one image; the probability of MixUp enhancement is 0.1, which mixes two images; and the probability of Copy-Paste enhancement is 0.1, which copies and pastes the target from one image to another.

5. The method for detecting small targets in UAV aerial images based on the AFLW-YOLOv5s model according to claim 1 is characterized in that: In the step (2), the image dataset after data enhancement is divided into a training set, a validation set and a test set.

6. The method for detecting small targets in UAV aerial images based on the AFLW-YOLOv5s model according to claim 1, characterized in that: The detection model in step (3) includes a CSM module for enhancing channel and spatial feature extraction capabilities and an LWinvolution module for improving model space and channel adaptability.

7. The method for detecting small targets in UAV aerial images based on the AFLW-YOLOv5s model according to claim 6, characterized in that: The implementation process of the CSM module for enhancing the channel and spatial feature extraction capability is as follows: first, a CSM module for enhancing the channel and spatial feature extraction capability is introduced into the backbone network of the YOLOv5s model, and then the input feature map is sequentially fed into a 1×1 convolution and a 3×3 convolution for partial feature extraction; secondly, the channel feature information, spatial feature information and global feature information of the input feature map are extracted through two dual-branch structures; then the initial feature map is input into a 3×3 convolution for global feature extraction.

8. The method for detecting small targets in UAV aerial images based on the AFLW-YOLOv5s model according to claim 7, characterized in that: The two double-branch structures are: The implementation process of the first dual-branch structure is as follows: First, a 1×1 convolution is used to extract partial features from the feature map. Then, the global average pooling layer and the global maximum pooling layer are used to extract channel information. Then, a 1×1 convolution and a ReLU activation function are used to output the channel attention weights. After weighted feature fusion, the sigmoid function is used to multiply the partial features originally extracted by the 1×1 convolution. The implementation process of the second dual-branch structure is as follows: First, a 1×1 convolution is used to extract partial features. Secondly, the spatial information is extracted through the global average pooling layer and the global maximum pooling layer respectively. The extracted spatial information is added together. Then, a 7×7 convolution and a Sigmoid function are used to output the spatial attention weights, which are multiplied channel by channel with the original partial feature map extracted by a 1×1 convolution, realizing the model's ability to extract channel information and spatial information from the input feature map.

9. The method for detecting small targets in drone aerial images based on the AFLW-YOLOv5s model according to claim 6, characterized in that: The implementation process of the LWinvolution module for improving the spatial and channel adaptability of the model is as follows: an LWinvolution module is integrated into the neck module of the model. The LWinvolution module first obtains the shape of the input feature map, and then generates a dynamic convolution kernel with a shape of a 1×1 convolution, a ReLU activation function and a 1×1 convolution on the input feature map; then the shape of the generated dynamic convolution kernel is adjusted by a reshape operation; then a group convolution operation is performed to group the input features; then the dynamic convolution of each group is calculated, and the shape of the feature map after calculation is multiplied by the shape of the feature map after the original grouping; the results of all group calculations are concat spliced; finally, the spliced ​​results are input into a 3×3 DWConv and a 1×1 PWConv.

10. The method for detecting small targets in UAV aerial images based on the AFLW-YOLOv5s model according to claim 1, wherein: The implementation process of the MSA-BiFPN neck network structure for realizing multi-scale feature fusion and introducing the small target detection head ADH in the target detection layer in the step (3) is as follows: first, the BiFPN neck network structure is improved, and the network cross-layer connects the 20×20, 40×40, and 80×80 feature maps input by the Backbone backbone network with the feature maps of the same size in the BiFPN network; an upsampling operation is added after the 80×80 size feature map in the BiFPN network; then, the feature map size is converted to 160×160, fused with the feature map of the same scale of the Backbone, and a downsampling operation is added to match it; finally, the 160×160 size feature map information is passed to the 160×160 size detection head.

Images