Model training method and aerial image small target detection method

Abstract

The invention discloses a model training method and an aerial image small target detection method, and belongs to the technical field of target detection. According to the method, a YOLOv8 improved model comprising a triple attention mechanism module Triplet Attention, a small-scale target detection layer P2 and a full-dimensional dynamic convolution layer ODConv is trained based on a sample aerial image and a sample aerial marking image, and a target detection model is obtained. According to the method, the attention degree on the small target can be improved, the detection robustness and stability are enhanced, and the small target detection performance of the model is improved.

Description

Technical Field

[0001] The present application relates to the technical field of target detection, and in particular to a model training method and a method for detecting small targets in aerial images. Background Art

[0002] With the rapid development of the Internet of Things and advanced sensing technology, the types and number of monitoring equipment have increased dramatically, bringing about the demand for the collection and processing of massive image data. Especially in crowded public places such as squares and stations, real-time and accurate image monitoring is crucial for crime prevention, traffic control and key target tracking. The use of Internet of Things technologies such as surveillance cameras and drones can obtain massive surveillance images in real time and achieve no-blind-angle monitoring. Based on the massive image data collected in real time, researchers have developed various real-time monitoring systems, such as real-time pedestrian monitoring systems, through advanced algorithms to achieve precise positioning and tracking. In the field of drones, the detection of complex multi-target and small targets in surveillance images is a key core technology in intelligent monitoring systems.

[0003] The target detection methods in related technologies often detect small targets through models such as YOLO and SSD. However, drone aerial images often involve complex backgrounds and small-sized targets, which limits the performance of target detection methods in related technologies in small target detection. There are the following problems: resolution limitations make small targets difficult to identify, complex and changing backgrounds increase the false detection rate, the dynamic changes of drones require the model to be highly adaptable, and there is a balance between real-time requirements and accuracy. Summary of the invention

[0004] In view of this, the present application provides a model training method and a small target detection method for aerial images, which can increase the attention to small targets, enhance the robustness and stability of detection, and improve the small target detection performance of the model.

[0005] Specifically, the following technical solutions are included:

[0006] In a first aspect, the present application provides a model training method, the method comprising:

[0007] Obtain sample aerial images, sample aerial annotated images and a YOLOv8 improved model, wherein the YOLOv8 improved model includes a triple attention mechanism module TripletAttention, a small-scale target detection layer P2, and a full-dimensional dynamic convolution layer ODConv;

[0008] Based on the sample aerial images and the sample aerial annotated images, a YOLOv8 improved model is trained to obtain a target detection model.

[0009] In some embodiments, the training of the YOLOv8 improved model based on the sample aerial image and the sample aerial annotated image to obtain the target detection model includes:

[0010] Based on the sample aerial images and the sample aerial annotated images, training a YOLOv8 improved model to obtain an intermediate YOLOv8 improved model;

[0011] In response to the performance of the intermediate YOLOv8 improved model reaching a preset standard, the intermediate YOLOv8 improved model reaching the preset standard is used as the target detection model; in response to the performance of the intermediate YOLOv8 improved model not reaching the preset standard, the training parameters of the intermediate YOLOv8 improved model are adjusted until the performance of the intermediate YOLOv8 improved model reaches the preset standard, the training parameters comprising one or more parameters of an initial learning rate, a number of training iterations and a hyperparameter.

[0012] In some embodiments, the YOLOv8 improved model is a model improved on the basis of the YOLOv8 model. Both the YOLOv8 model and the YOLOv8 improved model include a feature extraction part and a feature fusion part. The improvement method includes:

[0013] A triplet attention mechanism module is added after each upsampling layer Upsamp in the feature fusion part of the YOLOv8 model to obtain the first model;

[0014] On the basis of the three detection layers in the feature fusion part of the first model, a small-scale target detection layer P2 is added, and the large-scale target detection layer P5 and the feature extraction part corresponding to the large-scale target detection layer P5 in the feature fusion part of the first model are deleted to obtain a second model;

[0015] A full-dimensional dynamic convolution layer ODConv is set between the feature extraction part and the feature fusion part of the second model, and the number of output channels of the feature fusion part is set to 256 to obtain the YOLOv8 improved model.

[0016] In some embodiments, after the YOLOv8 improved model is trained based on the sample aerial images and the sample aerial annotated images to obtain the intermediate YOLOv8 improved model, the method further includes calculating a loss function.

[0017] In some embodiments, the numbers of output channels of the three detection layers in the YOLOv8 improved model are 128, 256, and 512, respectively.

[0018] In some embodiments, the triple attention mechanism module TripletAttention has three parallel connected branches, the first branch includes a first replacement layer Premute, a first maximum pooling layer Z-Pool, a first dimensional convolution layer, a first batch normalization layer BatchNorm, a first activation layer, and a second replacement layer Premute connected in sequence, the output of the first replacement layer Premute is introduced between the first batch normalization layer Batch Norm and the second replacement layer Premute, the second branch includes a third replacement layer Premute, a second maximum pooling layer Z-Pool, a second dimensional convolution layer, a second batch normalization layer Batch Norm, a second activation layer and a fourth replacement layer Premute connected in sequence, the output of the third replacement layer Premute is introduced between the second batch normalization layer BatchNorm and the fourth replacement layer Premute, the third branch includes a channel pooling layer, a third dimensional convolution layer, a third batch normalization layer BatchNorm and a third activation layer connected in sequence, and the first activation layer, the second activation layer and the third activation layer all generate attention weights through the Sigmoid function.

[0019] In some embodiments, the pixel resolution of the small-scale target detection layer P2 is 160 pixels×160 pixels.

[0020] In some embodiments, before training the YOLOv8 improved model, the initial learning rate is set to 0.01, data enhancement is turned off in the last 10 rounds of training, and the waiting rounds of early stopping training are set to 50.

[0021] In some embodiments, the preset standard is met when the training results remain unchanged for 50 consecutive rounds, and the training results include precision, recall rate, mAP and parameter quantity.

[0022] In a second aspect, the present application provides a method for detecting small targets in aerial images, the method comprising:

[0023] Obtaining the target detection model in the model training method as described in the first aspect;

[0024] The aerial image to be tested is input into the target detection model to obtain a target detection result, which includes a target category, a target bounding box position, a confidence level, and a category probability.

[0025] The beneficial effects of the technical solution provided by the embodiments of the present application include at least:

[0026] The embodiment of the present application provides a model training method and a method for detecting small targets in aerial images. The method is based on sample aerial images and sample aerial annotated images, and trains a YOLOv8 improved model including a triple attention mechanism module TripletAttention, a small-scale target detection layer P2, and a full-dimensional dynamic convolution layer ODConv to obtain a target detection model. The triple attention mechanism module TripletAttention can simultaneously focus on spatial, channel, and pixel-level features, and enhance the network's feature representation capabilities and the detection capabilities of small targets by focusing on important feature areas to help the model distinguish multiple closely arranged objects. The small-scale target detection layer P2 can solve the problem of insufficient capture of small target information features in aerial images, and the full-dimensional dynamic convolution layer ODConv can better balance the detail information in the high-resolution feature map with the semantic information in the low-resolution feature map, ensuring that the feature map can be effectively fused, enhancing the robustness and stability of the detection, and improving the small target detection performance of the model. BRIEF DESCRIPTION OF THE DRAWINGS

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

[0028] Figure 1 A method flow chart of the model training method provided in an embodiment of the present application;

[0029] Figure 2 A schematic diagram of the structure of the YOLOv8 improved model of the model training method provided in the embodiment of the present application;

[0030] Figure 3 A schematic diagram of the structure of the triple attention mechanism module Triplet Attention of the YOLOv8 improved model of the model training method provided in an embodiment of the present application;

[0031] Figure 4 A schematic diagram of the improvement of the small-scale target detection layer P2 by the YOLOv8 improved model of the model training method provided in an embodiment of the present application;

[0032] Figure 5 Schematic diagram of the effect of the full-dimensional dynamic convolution layer ODConv after unifying channels in the model training method provided in an embodiment of the present application. DETAILED DESCRIPTION

[0033] The following will be combined with the drawings in the embodiments of the present application to clearly and completely describe the technical solutions in the embodiments of the present application. Obviously, the described embodiments are part of the embodiments of the present application, not all of the embodiments. Based on the embodiments in the present application, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of this application.

[0034] In order to make the technical solutions and advantages of the present application more clear, the implementation methods of the present application will be further described in detail below with reference to the accompanying drawings.

[0035] In the first aspect, the present application provides a model training method, such as Figure 1 As shown, the method includes:

[0036] Step 101, obtain sample aerial images, sample aerial annotated images and a YOLOv8 improved model, where the YOLOv8 improved model includes a triple attention mechanism module TripletAttention, a small-scale target detection layer P2, and a full-dimensional dynamic convolution layer ODConv.

[0037] The sample aerial images and sample aerial annotated images can be used to train the YOLOv8 improved model. Since the YOLOv8 improved model includes the triple attention mechanism module TripletAttention ( Figure 2 Triplet Attention in the small-scale target detection layer P2 ( Figure 2 P2 in ) and the full-dimensional dynamic convolution layer ODConv ( Figure 2 The triple attention mechanism module TripletAttention can focus on the features at the spatial, channel and pixel levels at the same time. It focuses on important feature areas to help the model distinguish multiple closely arranged objects to enhance the feature representation ability of the network and the detection ability of small targets. The small-scale target detection layer P2 can solve the problem of insufficient feature capture of small target information in aerial images. The full-dimensional dynamic convolution layer ODConv can better balance the detail information in the high-resolution feature map with the semantic information in the low-resolution feature map, ensuring that the feature maps can be effectively fused and improving the performance of the model.

[0038] In some embodiments, the sample aerial images and the sample aerial annotated images may come from the VisDrone-2019 dataset.

[0039] It should be noted that the VisDrone-2019 dataset was collected and produced by the AISKYEYE team of Tianjin University and the Data Mining Laboratory. This dataset is a large dataset collected by multiple drone platforms under various conditions such as different weather and lighting conditions. The images are captured by cameras mounted on drones, including location (14 cities in China thousands of kilometers apart), environment (urban or rural), objects (people or various vehicles), and density (sparse or crowded scenes). The images used in the target detection task in this dataset are divided into 6471 training images, 1610 test set images, and 548 validation set images, including 10 categories: pedestrians, people, bicycles, cars, vans, trucks, tricycles, awning tricycles, buses, and motorcycles. The dataset contains many complex scenes, with a large impact on the lighting environment, a wide range of target sizes, and a variety of shooting angles. It also contains a large number of small targets with occlusion, which makes the recognition of small targets challenging.

[0040] In some embodiments, the YOLOv8 improved model is a model improved on the basis of the YOLOv8 model. Both the YOLOv8 model and the YOLOv8 improved model include a feature extraction part and a feature fusion part. Figure 2 As shown, the improvements include:

[0041] (1) A triplet attention mechanism module is added before each upsample layer Upsample in the feature fusion part of the YOLOv8 model to obtain the first model.

[0042] The triple attention mechanism module Triplet Attention enhances the network's feature representation ability by focusing on important feature areas to help the model distinguish multiple closely arranged objects, thereby enhancing the detection ability of small targets. Before upsampling, the feature map is processed by the triple attention mechanism module TripletAttention, which can enhance the quality of the feature map and make the upsampled feature map more representative and discriminative, so that important feature information can be better retained during the upsampling process, avoiding information loss and noise introduction during the upsampling process.

[0043] See also Figure 3In some embodiments, the triple attention mechanism module Triplet Attention has three parallel connected branches, the first branch includes a first permutation layer Premute, a first maximum pooling layer Z-Pool, a first dimensional convolution layer, a first batch normalization layer BatchNorm, a first activation layer, and a second permutation layer Premute connected in sequence, the output of the first permutation layer Premute is introduced between the first batch normalization layer BatchNorm and the second permutation layer Premute, the second branch includes a third permutation layer Premute, a second maximum pooling layer Z-Pool, a second dimensional convolution layer, a second batch normalization layer Batch Norm, a second activation layer, and a fourth permutation layer Premute connected in sequence, the output of the third permutation layer Premute is introduced between the second batch normalization layer BatchNorm and the fourth permutation layer Premute, the third branch includes a channel pooling layer, a third dimensional convolution layer, a third batch normalization layer BatchNorm, and a third activation layer connected in sequence, and the first activation layer, the second activation layer, and the third activation layer all generate attention weights through a Sigmoid function.

[0044] In some embodiments, the dimensions of the first dimension convolution layer, the second dimension convolution layer, and the third dimension convolution layer are 7×7.

[0045] In some embodiments, in the YOLOv8 improved model, the input image needs to be converted into a tensor for forward propagation and reasoning of the model. The YOLOv8 improved model requires that the channel order of the input tensor is the number of channels, height, width (ie, C×H×W), such as Figure 3 The input tensor X(C×H×W) is passed in.

[0046] See also Figure 3 , TripletAttention mechanism module, the first branch establishes interaction between the H dimension and the C dimension of the image, the first branch sees Figure 3 First, through the first replacement layer Premute (i.e. Figure 3 The first Premute in the leftmost branch rotates the input tensor X 90° counterclockwise along the H axis to obtain the rotated tensor The rotation tensor The shape of is W×H×C, and then passes through the first maximum pooling layer Z-Pool (i.e. Figure 3 After processing by Z-Pool in the leftmost branch, we get the tensor Tensor The shape of is 2×H×C. Through the standard convolution layer with a kernel size of 7×7 (i.e. the first dimension convolution layer, i.e. Figure 37×7Conv in the leftmost branch), and then through the first batch normalization layer Batch Norm (i.e. Figure 3 The Batch Norm in the leftmost branch) obtains an intermediate output tensor of dimension 1×H×C. Then, through the first activation layer (i.e. Figure 3 Sigmoid in the leftmost branch) converts the intermediate output tensor and the rotation tensor Generate attention weights and obtain the first branch result to be output. Finally, the first branch result to be output is passed through the second replacement layer Premute (i.e. Figure 3 The second Premute from the top in the leftmost branch is rotated 90° clockwise along the H axis to obtain the first branch result, thereby maintaining a shape consistent with the input.

[0047] In the second branch, the image is interactively connected between the C dimension and the W dimension. Figure 3 The middle branch. First, through the third replacement layer Premute (i.e. Figure 3 Premute in the middle branch rotates the input tensor X 90° counterclockwise along the W axis to obtain the rotated tensor The rotation tensor The shape of is H×C×W, and then passes through the second maximum pooling layer Z-Pool (i.e. Figure 3 After processing by Z-Pool in the middle branch, a tensor with a shape of 2×C×W is obtained. Through the standard convolution layer with a kernel size of 7×7 (i.e., the second dimension convolution layer, i.e. Figure 3 7×7Conv in the middle branch), and then through the second batch normalization layer BatchNorm (i.e. Figure 3 BatchNorm in the middle branch) to obtain an intermediate output tensor of dimension 1×C×W. Then, through the second activation layer (i.e. Figure 3 Sigmoid in the middle branch) converts the intermediate output tensor and the rotation tensor Generate attention weights and obtain the second branch result to be output. Finally, the second branch result to be output is passed through the fourth replacement layer Premute (i.e. Figure 3 The second Premute from the top in the middle branch is rotated 90° clockwise along the W axis to obtain the second branch result, thereby maintaining a shape consistent with the input.

[0048] In the third branch, the image is interactively connected between the H dimension and the W dimension. Figure 3 The rightmost branch. First, through the channel pooling layer (i.e. Figure 3The Channel Pool in the rightmost branch reduces the channel dimension of the input tensor to two dimensions. This simplified tensor of shape 2×H×W is passed through a standard convolutional layer with a kernel size of 7×7 (i.e., the third-dimensional convolutional layer, i.e., Figure 3 7×7Conv in the rightmost branch), and then through the third batch normalization layer BatchNorm (i.e. Figure 3 Batch Norm in the rightmost branch) to obtain an intermediate output tensor of dimension 1×H×W. Then, through the third activation layer (i.e. Figure 3 Sigmoid in the rightmost branch generates an attention weight with a shape of 1×H×W for the intermediate output tensor and input tensor X, and obtains the result of the third branch.

[0049] The results generated by the above three branches are averaged and aggregated to obtain a refined tensor y with a shape of C×H×W.

[0050] Figure 3 The Z-Pool in is used to reduce the C-dimensional tensor to two dimensions, connecting the average pooled features and the maximum pooled features on this dimension, thereby retaining the rich representation of the actual tensor while reducing its depth and reducing the amount of computation.

[0051] (2) On the basis of the three detection layers in the feature fusion part of the first model, a small-scale target detection layer P2 is added, and the large-scale target detection layer P5 and the feature extraction part corresponding to the large-scale target detection layer P5 in the feature fusion part of the first model are deleted to obtain a second model.

[0052] The added small-scale target detection layer P2 can solve the problem of insufficient capture of small target information features in the sample aerial images. In addition, in order to avoid the influence of noise introduced by the large-scale detection head on the small target detection accuracy and the occupation of computing resources, the large-scale target detection layer P5 in the feature fusion part of the first model and the feature extraction part corresponding to the large-scale target detection layer P5 are deleted, and more computing resources are used to process high-resolution feature maps, thereby enhancing the detection capability of small targets.

[0053] In some embodiments, the small-scale target detection layer P2 can perform the following processing: upsample the 80×80 feature map output by upsampling in the feature fusion part of the first model again to obtain shallower feature information; then concatenate and fuse the shallower feature information with the output of the shallow feature convolution layer of the corresponding scale in the feature extraction part, and output the target category and position of 160×160 scale after a 1×1 convolution operation.

[0054] In some embodiments, the output channel numbers of the three detection layers in the YOLOv8 improved model are 128, 256, and 512, respectively, 128 corresponds to the channel number of the small-scale target detection layer P2, 256 corresponds to the channel number of the target detection layer P3, and 512 corresponds to the channel number of the target detection layer P4.

[0055] In some embodiments, the pixel resolution of the small-scale object detection layer P2 is 160 pixels×160 pixels.

[0056] like Figure 4 As shown in the figure, the scales (i.e., pixel resolutions) of the feature maps output by the P3, P4, and P5 detection layers are 80pixel×80pixel, 40pixel×40pixel, and 20pixel×20pixel, respectively, corresponding to the detection of small, medium, and large targets. However, in reality, the scales of many small targets are smaller than 8pixel×8pixel, resulting in low accuracy in small target detection and serious missed detections. The leftmost part is based on the three detection layers in the feature fusion part of the first model, and adds a small-scale target detection layer P2. The pixel resolution of the small-scale target detection layer P2 is 160pixel×160pixel, which enables the capture of tiny targets and improves detection accuracy. Figure 4 The feature extraction part corresponding to the small-scale target detection layer P2 is added in the middle part. Figure 4 At the far right of , the large-scale target detection layer P5 and the feature extraction part corresponding to the large-scale target detection layer P5 in the feature fusion part of the first model are deleted, thereby greatly reducing the number of parameters, allocating computing resources to small target detection, and improving the detection accuracy of small targets. Figure 2 As shown in the figure, the improved YOLOv8 model first performs the flow of deep semantic information to shallow semantic information in the feature fusion part, realizing the feature fusion of 40pixel×40pixel, 80pixel×80pixel, and 160pixel×160pixel, and then performs the flow of shallow semantic information to deep semantic information, realizing the feature fusion of 160pixel×160pixel, 80pixel×80pixel, and 40pixel×40pixel, reducing the loss of small target feature information.

[0057] It should be noted that the small-scale target detection layer P2 represents a shallower layer in the feature pyramid network, has a smaller receptive field, and is suitable for capturing the detailed features of small-sized targets. Introducing it as a new detection layer can effectively improve the model's sensitivity and recognition ability for small targets. The small-scale target detection layer P2 is very suitable for feature extraction and detection of small targets due to its high resolution and richer texture information. Only adding the small-scale target detection layer P2 improves the small target detection performance, but the network has a large number of parameters and a large computational overhead. In order to reduce the computational overhead brought by the introduction of the small-scale target detection layer P2, the large-scale target detection layer P5 and its corresponding feature extraction part are removed, so that the focus and resources of the network are shifted from large target detection to finer-grained feature extraction, which helps to reduce the waste of resources that may be caused by focusing on large-sized targets, and give more computing resources to the detection of small targets. At the same time, it also avoids the influence of noise that may be introduced by the large-size detection head on the accuracy of small target detection.

[0058] (3) A full-dimensional dynamic convolution layer ODConv is set between the feature extraction part and the feature fusion part of the second model, and the number of output channels of the feature fusion part is set to 256 to obtain the YOLOv8 improved model.

[0059] The full-dimensional dynamic convolution layer ODConv can better balance the detail information in the high-resolution feature map with the semantic information in the low-resolution feature map, ensuring that each feature layer from P4 to P2 can be effectively fused in the same dimension, thereby improving the performance of the overall model.

[0060] Since the output channels of the three detection layers in the YOLOv8 improved model are 128, 256, and 512 respectively, high-level feature maps usually have fewer channels but contain richer semantic information, while low-level feature maps have more channels and contain more detailed information. This inconsistency in the number of channels may cause the low-level detailed information to be obscured by the high-level semantic information during feature fusion, or the high-level semantic information cannot be effectively used to enhance the low-level feature expression, thus affecting the detection effect. Figure 5 As shown in the figure, the number of output channels of the feature fusion part is set to 256, and the number of channels is unified, which ensures that each feature layer from P4 to P2 can be effectively fused in the same dimension, thereby improving the performance of the overall model.

[0061] It should be noted that the full-dimensional dynamic convolution layer ODConv is different from the traditional static convolution kernel. It also takes into account the dynamics of spatial domain, input channel, output channel and other dimensions. It can dynamically adjust the parameters of its convolution kernel according to the input features. It contains four types of attention, which are complementary, so that the convolution operation has differences in various dimensions for the input, providing better performance to capture rich contextual information. This dynamic adjustment makes the convolution kernel specially customized for the current input during each forward propagation, thereby enhancing the adaptability and flexibility of the network to the input data.

[0062] Step 102: Based on the sample aerial images and the sample aerial annotated images, train the YOLOv8 improved model to obtain a target detection model.

[0063] In some embodiments, based on the sample aerial images and the sample aerial annotated images, the YOLOv8 improved model is trained to obtain the target detection model, including:

[0064] Based on the sample aerial images and the sample aerial annotated images, the YOLOv8 improved model is trained to obtain the intermediate YOLOv8 improved model;

[0065] In response to the performance of the intermediate YOLOv8 improved model reaching a preset standard, the intermediate YOLOv8 improved model reaching the preset standard is used as a target detection model; in response to the performance of the intermediate YOLOv8 improved model not reaching the preset standard, the training parameters of the intermediate YOLOv8 improved model are adjusted until the performance of the intermediate YOLOv8 improved model reaches the preset standard, the training parameters comprising one or more parameters of an initial learning rate, a number of training iterations and a hyperparameter.

[0066] In some embodiments, before training the YOLOv8 improved model, the initial learning rate is set to 0.01, data enhancement is turned off in the last 10 rounds of training to prevent model overfitting, and the waiting rounds of early stopping training are set to 50.

[0067] In some embodiments, the preset standard is met that the training results do not change for 50 consecutive rounds, and the training results include accuracy, recall rate, mAP and parameter quantity.

[0068] In some embodiments, after training the YOLOv8 improved model based on the sample aerial images and the sample aerial annotated images to obtain the intermediate YOLOv8 improved model, it also includes calculating the loss function to measure the model performance, provide guidance for the adjustment of the training parameters, and ensure the robustness and efficiency of the model in various complex environments.

[0069] In some embodiments, in the YOLOv8 improved model, the regression results output by the model are subjected to feature reconstruction operations to obtain a more refined feature map, on which basis classification and regression operations are performed again to calculate the loss.

[0070] Table 1 Performance comparison between YOLOv8 improved model and YOLOv8 model

[0071]

[0072] As shown in Table 1, baseline represents the YOLOv8 model before improvement. Check marks represent the introduction of different modules respectively. No check mark means that the module is not introduced. "Triplet Attention" means the introduction of the triple attention mechanism module TripletAttention. "Add P2, delete P5" means the introduction of the small-scale target detection layer P2 and the deletion of the large target detection layer P5. "ODConv" means the introduction of the full-dimensional dynamic convolution layer ODConv. P represents the precision of the improved YOLOv8 model. R represents the recall rate and mAP. 50 、mAP 50-95 They represent the average precision under the 50% IoU threshold and the average precision in the 50%-95% IoU threshold range. See Table 1, the precision (P) increased by 6.7%, the recall rate (R) increased by 6.6%, and the mAP 50 Up 8.3%, mAP 50-95 It increased by 5.6%, and the number of parameters decreased from 11,129,454 to 3,912,068, a decrease of 64.8%. It should be noted that IoU is the overlap height between the predicted box and the true box.

[0073] In summary, the model training method provided by this application, the triple attention mechanism module TripletAttention can simultaneously focus on spatial, channel and pixel-level features, and enhance the network's feature representation ability by focusing on important feature areas to help the model distinguish multiple closely arranged objects, thereby enhancing the ability to detect small targets. The small-scale target detection layer P2 can solve the problem of insufficient capture of small target information features in aerial images. The full-dimensional dynamic convolution layer ODConv can better balance the detail information in the high-resolution feature map with the semantic information in the low-resolution feature map, ensuring that the feature maps can be effectively fused, enhancing the robustness and stability of the detection, and improving the small target detection performance of the model.

[0074] In a second aspect, the present application provides a method for detecting small targets in aerial images, the method comprising:

[0075] Obtaining the target detection model in the model training method of the first aspect;

[0076] The aerial image to be tested is input into the target detection model to obtain the target detection result, which includes the target category, target bounding box position, confidence and category probability.

[0077] In summary, the small target detection method for aerial images provided in this application can increase the level of attention to small targets, enhance the robustness and stability of detection, and improve the small target detection performance of the model.

[0078] In the present application, the terms "first" and "second" are used for descriptive purposes only and should not be understood as indicating or implying relative importance. The term "plurality" refers to two or more than two, unless otherwise clearly defined.

[0079] Those skilled in the art will readily appreciate other embodiments of the present application after considering the specification and practicing the present application disclosed herein. The present application is intended to cover any variations, uses or adaptations of the present application, which follow the general principles of the present application and include common knowledge or customary techniques in the art that are not disclosed in the present application. The specification and examples are intended to be exemplary only.

[0080] It should be understood that the present application is not limited to the precise structures that have been described above and shown in the drawings, and that various modifications and changes may be made without departing from the scope thereof. The scope of the present application is limited only by the appended claims.

Claims

1. A model training method, characterized in that: The method comprises: Obtain sample aerial images, sample aerial annotated images and a YOLOv8 improved model, wherein the YOLOv8 improved model includes a triple attention mechanism module TripletAttention, a small-scale target detection layer P2, and a full-dimensional dynamic convolution layer ODConv; Based on the sample aerial images and the sample aerial annotated images, a YOLOv8 improved model is trained to obtain a target detection model.

2. The model training method according to claim 1, characterized in that: The step of training the YOLOv8 improved model based on the sample aerial image and the sample aerial annotated image to obtain a target detection model includes: Based on the sample aerial images and the sample aerial annotated images, training a YOLOv8 improved model to obtain an intermediate YOLOv8 improved model; In response to the performance of the intermediate YOLOv8 improved model reaching a preset standard, the intermediate YOLOv8 improved model reaching the preset standard is used as the target detection model; in response to the performance of the intermediate YOLOv8 improved model not reaching the preset standard, the training parameters of the intermediate YOLOv8 improved model are adjusted until the performance of the intermediate YOLOv8 improved model reaches the preset standard, the training parameters comprising one or more parameters of an initial learning rate, a number of training iterations and a hyperparameter.

3. The model training method according to claim 1, characterized in that: The YOLOv8 improved model is a model that is improved on the basis of the YOLOv8 model. Both the YOLOv8 model and the YOLOv8 improved model include a feature extraction part and a feature fusion part. The improvement methods include: A triplet attention mechanism module is added after each upsampling layer Upsamp in the feature fusion part of the YOLOv8 model to obtain the first model; On the basis of the three detection layers in the feature fusion part of the first model, a small-scale target detection layer P2 is added, and the large-scale target detection layer P5 and the feature extraction part corresponding to the large-scale target detection layer P5 in the feature fusion part of the first model are deleted to obtain a second model; A full-dimensional dynamic convolution layer ODConv is set between the feature extraction part and the feature fusion part of the second model, and the number of output channels of the feature fusion part is set to 256 to obtain the YOLOv8 improved model.

4. The model training method according to claim 2, characterized in that: After the YOLOv8 improved model is trained based on the sample aerial image and the sample aerial annotated image to obtain the intermediate YOLOv8 improved model, the method further includes calculating a loss function.

5. The model training method according to claim 3, characterized in that: The numbers of output channels of the three detection layers in the YOLOv8 improved model are 128, 256, and 512, respectively.

6. The model training method according to claim 3, characterized in that: The triple attention mechanism module Triplet Attention has three parallel connected branches, the first branch includes a first permutation layer Premute, a first maximum pooling layer Z-Pool, a first dimensional convolution layer, a first batch normalization layer Batch Norm, a first activation layer, and a second permutation layer Premute connected in sequence, the output of the first permutation layer Premute is introduced between the first batch normalization layer BatchNorm and the second permutation layer Premute, the second branch includes a third permutation layer Premute, a second maximum pooling layer Z-Pool, a second dimensional convolution layer, a second batch normalization layer Batch Norm, a second activation layer, and a fourth permutation layer Premute connected in sequence, the output of the third permutation layer Premute is introduced between the second batch normalization layer Batch Norm and the fourth permutation layer Premute, the third branch includes a channel pooling layer, a third dimensional convolution layer, a third batch normalization layer BatchNorm, and a third activation layer connected in sequence, the first activation layer, the second activation layer, and the third activation layer all generate attention weights through a Sigmoid function.

7. The model training method according to claim 3 or 6, characterized in that: The pixel resolution of the small-scale target detection layer P2 is 160 pixels × 160 pixels.

8. The model training method according to claim 2, characterized in that: Before training the YOLOv8 improved model, the initial learning rate was set to 0.01, data enhancement was turned off in the last 10 rounds of training, and the waiting rounds of early stopping training were set to 50.

9. The model training method according to claim 2, characterized in that: The preset standard is met so that the training results do not change for 50 consecutive rounds, and the training results include precision, recall rate, mAP and parameter quantity.

10. A method for detecting small targets in aerial images, characterized in that: The method comprises: Obtaining a target detection model in the model training method as described in any one of claims 1 to 9; The aerial image to be tested is input into the target detection model to obtain a target detection result, which includes a target category, a target bounding box position, a confidence level, and a category probability.

Images