Fast Target Detection Method for Remote Sensing Images Based on Attention Mechanism and Pruning

Abstract

The present invention provides a fast target detection method for remote sensing images based on the attention mechanism and pruning. First, a simple attention mechanism CBAM is integrated into GhostNet (this attention mechanism can be integrated into any convolutional neural network, effectively improving the detection accuracy of the model without bringing a large increase in the model's parameters) to obtain the GhostNet-CBAM network; then, the heavy feature extraction network in YOLOv5 is replaced by the GhostNet-CBAM network, greatly reducing the number of parameters and the computational amount of the model and accelerating the inference speed of the model. Subsequently, a rotation angle term is added to the output part and the loss function of YOLOv5, so that the model outputs a five-parameter rotated bounding box. Compared with the commonly used four-parameter rectangular target detection, the five-parameter rotated rectangular box can more accurately reflect the true size and aspect ratio of the target; finally, a model pruning algorithm based on the L1 norm is used to prune the above-improved target detection model, significantly reducing the number of parameters, the computational amount, and the inference time of the model, and enabling it to be deployed to some edge devices with limited storage and computing resources to achieve the effect of real-time detection.

Description

Technical Field

[0001] The present invention relates to the technical field of remote sensing image information processing, and particularly to a fast target detection method for remote sensing images based on an attention mechanism and pruning. Background Art

[0002] Target detection refers to detecting whether there are targets of interest (such as pedestrians, vehicles, animals, etc.) in an image, as well as information such as the type, location, and quantity of the targets. Remote sensing image target detection refers to detecting targets of interest (such as buildings, roads, ships, etc.) in remote sensing images. The difference from natural image target detection is that remote sensing images are affected by the shooting angle, and the same target may present different rotation angles. Moreover, compared with natural images, remote sensing images contain more small targets, which makes remote sensing image target detection more difficult than natural image target detection. In recent years, the earth observation technology has developed rapidly, and the data volume and resolution of remote sensing images have been significantly improved, greatly increasing the difficulty of remote sensing image target detection.

[0003] Up to now, remote sensing image target detection technology has mainly gone through two stages: traditional detection methods and deep learning-based detection methods. Remote sensing image target detection algorithms can be roughly divided into the following categories: template matching-based methods, prior knowledge-based methods, object-based image analysis methods, traditional machine learning-based methods, and deep learning-based methods. Template matching-based methods are complex because they require manual feature design, and have problems such as slow running speed and low detection accuracy. Prior knowledge-based methods need to establish various knowledge and rules, and it is difficult to grasp the definition of the rules. If the rules are too strict, many targets cannot be detected, and if they are too loose, some non-target objects may be misidentified as targets. Object-based image analysis methods need to segment the image into homogeneous regions, and the clarity of the contour will directly affect the detection accuracy. Traditional machine learning-based methods usually use an exhaustive method based on sliding windows. The large amount of calculation leads to low efficiency. The sizes and dimensions of targets in remote sensing images are diverse, and usually multiple-scale sliding windows need to be used, further increasing the amount of calculation and with low accuracy. Deep learning-based methods have a fast detection speed and high accuracy, and are currently the mainstream detection methods. However, the number of model parameters and the amount of calculation are usually large, and edge devices usually have difficulty meeting their operating conditions.

[0004] Traditional remote sensing image target detection methods have certain bottlenecks in terms of accuracy or detection speed. Deep learning-based methods have high detection accuracy and have become the mainstream in the current target detection field. However, their parameter and computational amounts are large, making it difficult to deploy them to edge devices (such as drones) with limited storage and computing capabilities. Therefore, reducing the parameter and computational amounts of deep learning target detection models so that they can be deployed to resource-constrained platforms and improving the detection speed is of great significance. Summary of the Invention

[0005] In view of this, to reduce the parameter and computational amounts of deep learning target detection models for remote sensing images, the present invention provides a fast target detection method for remote sensing images based on the attention mechanism and pruning. The method includes the following steps:

[0006] S1: Obtain remote sensing image data, organize the data label format to obtain a data set, and divide the data set into a training set and a validation set. Perform Mosaic data augmentation and adaptive anchor box calculation on the training set remote sensing images;

[0007] S2: Select the lightweight feature extraction network GhostNet, fuse the attention mechanism module CBAM in this network to obtain the lightweight feature extraction network GhostNet-CBAM with the attention mechanism. Subsequently, replace the feature extraction network in the target detection model YOLOv5 with the lightweight feature extraction network GhostNet-CBAM to obtain a lightweight target detection network that integrates the lightweight feature extraction network and the attention mechanism;

[0008] S3: Add a rotation angle to the output part of the lightweight target detection model obtained in step S2, and add a loss of a rotation angle to the loss function of the lightweight target detection model to obtain a lightweight rotated bounding box target detection network;

[0009] S4: Train the lightweight rotated bounding box target detection model obtained in step S3 with the training set in step S1 until the model converges to obtain a trained lightweight target detection model;

[0010] S5: Extract all the convolutional kernels in each layer of convolution in the trained lightweight target detection model obtained in step S4, calculate their L1 norms respectively, remove the convolutional kernels with L1 norms less than the threshold, and then fine-tune the pruned model to obtain a lightweight target detection model after pruning;

[0011] S6: Replace the original NMS algorithm in the lightweight target detection model after pruning in step S5 with an NMS algorithm based on RIoU to remove duplicate targets in the target detection results to obtain the final model;

[0012] S7: Use the final model to perform fast target detection on newly acquired remote sensing images.

[0013] Furthermore, in step S1, the specific process of performing mosaic data enhancement on each remote sensing image in the training set is to randomly scale, randomly crop, and randomly arrange the four initial remote sensing images, and then splice them to obtain an image, which is used as a training set image.

[0014] Furthermore, the specific process of adaptive anchor frame calculation for each remote sensing image in the training set in step S1 is to use the K-means clustering algorithm to calculate several anchor frames, and then use the genetic algorithm to mutate the target length and width. If the mutation effect is better, the mutated length and width are set as the anchor frame value, otherwise it is skipped.

[0015] Furthermore, in step S1, the image labels in the dataset are sorted in the following format:

[0016] x,y,w,h,θ,class

[0017] Among them, x and y are the coordinates of the center point of the remote sensing image rectangle, w, h, and θ are the width, height, and rotation angle of the remote sensing image rectangle, respectively, and class is the category of the target in the remote sensing image.

[0018] Furthermore, the method further comprises step S61 after step S6:

[0019] S61: Use the final model to detect the validation set remote sensing images in step 1 and obtain the detection accuracy of the final model.

[0020] Furthermore, the total loss function of the lightweight target detection model after adding the rotation angle loss in step S3 is:

[0021] L total =λ box L box +λ obj L obj +λ cls L cls +λ theta L theta

[0022] Among them, L total is the total loss function, L box is the bounding box regression loss, L cls is the classification loss, L obj is the confidence loss, L theta is the angle classification loss, λ box is the proportion of border loss, λ obj is the proportion of confidence loss, λ clsThe proportion of the classification loss, λ theta is the proportion of the rotation angle loss.

[0023] The beneficial effects of a fast target detection method for remote sensing images based on the attention mechanism and pruning in the present invention are as follows:

[0024] By adding a rotation angle term to the output part and the loss function of the target detection model YOLOv5, the final model obtained by this method has the ability to detect rotated targets in remote sensing images; compared with typical general target detection algorithms for horizontal rectangular boxes, the final model of the present invention can better reflect the directional diversity characteristics of targets in remote sensing images.

[0025] This method replaces the heavy feature extraction network in the YOLOv5 target detection model with a lightweight feature extraction network incorporating the attention mechanism, greatly reducing the number of parameters and computational amount of the target detection model and accelerating the model inference speed.

[0026] This method adopts a model pruning algorithm based on the L1 norm, further significantly reducing the number of parameters and computational amount of the model, enabling the model to be deployed on edge devices with limited computing and storage capabilities. Description of the Drawings

[0027] Att Figure 1 is a schematic flow chart of a fast target detection method for remote sensing images based on the attention mechanism and pruning.

[0028] Att Figure 2 is a schematic diagram of the final model structure and processing flow of the present invention.

[0029] Att Figure 3 is an example diagram of Mosaic data augmentation.

[0030] Att Figure 4 is a schematic diagram showing the difference between conventional two types of rectangular box target detection and rotated box target detection.

[0031] Att Figure 5 is a schematic diagram of the convolutional kernel pruning algorithm based on the L1 norm used in the present invention.

[0032] Att Figure 6 is a schematic diagram of the graphical illustration of the RIoU algorithm used by the present invention when calculating the IoU of rotated boxes. Detailed Embodiments

[0033] To make the objectives, technical solutions, and advantages of the present invention clearer, the embodiments of the present invention will be further described below in conjunction with the accompanying drawings.

[0034] Refer to Figure 1, a fast target detection method for remote sensing images based on attention mechanism and pruning. The method includes the following steps:

[0035] S1: Obtain remote sensing image data, organize the data label format to obtain a dataset, and divide the dataset into a training set and a validation set. Perform data augmentation and adaptive anchor box calculation on the training set remote sensing images;

[0036] Among them, the specific processes of performing data augmentation, adaptive anchor box calculation, and label format arrangement on each remote sensing image in the training set in step S1 are as follows:

[0037] Data augmentation: Perform Mosaic data augmentation on the remote sensing image, that is, randomly crop four initial remote sensing images and then splice them to obtain an image, which is used as a training set image (refer to Figure 3 ).

[0038] Adaptive anchor box calculation: Use the K-means clustering algorithm to calculate several anchor boxes, and then use the genetic algorithm to mutate the target length and width. If the effect after mutation is better, set the mutated length and width as the anchor box values, otherwise skip.

[0039] Label format arrangement: Since the detection result of the present invention is the result of a rotated box, the dataset needs to be converted into the corresponding format. The label format of the dataset is as follows:

[0040] x, y, w, h, θ, class

[0041] Among them, x and y are the center point coordinates of the rectangle box of the remote sensing image, w, h, and θ are the width, height, and rotation angle of the rectangle box of the remote sensing image respectively, and class is the category of the target in the remote sensing image.

[0042] S2: Select the lightweight feature extraction network GhostNet, fuse the attention mechanism module CBAM in this network to obtain the lightweight feature extraction network GhostNet-CBAM with fused attention mechanism. Then replace the feature extraction network of the target detection model YOLOv5 with the lightweight feature extraction network GhostNet-CBAM to obtain a lightweight object detection network that fuses the lightweight feature extraction network and the attention mechanism;

[0043] S3: Add a rotation angle to the output part of the lightweight object detection model obtained in step S2, and add a loss of the rotation angle to the loss function of the lightweight object detection model to obtain a lightweight rotated box object detection network;

[0044] The prediction result of the YOLOv5 object detection model only has four parameters: xywh. Considering the diversity of object directions in remote sensing images, this invention adds a rotation angle to the output part of YOLOv5 and adds a parameter θ to the final prediction result. The difference between conventional rectangular box object detection and the rotated rectangular box object detection in this invention is as shown in Appendix Figure 4 This parameter represents the rotation angle of the object relative to the position of the positive rectangular box in the final prediction result. The value range of θ is [-90, 90), and θ is an integer. On the basis of step 2, the detection head of the model is improved so that the output of the model has an additional detection angle, and a loss of the rotation angle is added to the loss function of YOLOv5.

[0045] The loss function of YOLOv5 includes three parts: bounding box regression loss, object classification loss, and confidence loss. For the rotated box object detection task, this invention adds a classification loss of the rotation angle to the loss function. The total loss function is as follows:

[0046] L total = λ box L box + λ obj L obj + λ cls L cls + λ theta L theta

[0047] Among them, L box is the bounding box regression loss, λ box is a hyperparameter representing the proportion of the bounding box loss in the total loss; L cls is the classification loss, λ cls is the proportion of the classification loss in the total loss; L obj is the confidence loss, λ obj is the proportion of the confidence loss in the total loss; L theta is the angle classification loss, λ theta is the proportion of the rotation angle loss in the total loss. The detailed introduction of each loss function is as follows.

[0048] 1) Bounding box regression loss

[0049] In this invention, the rotation angle θ is not regressed together with the rectangular box parameters. Instead, the rotation angle is regarded as a classification task with 180 categories, and the bounding box loss and the rotation angle loss are calculated separately. The bounding box loss function uses the original loss function:

[0050]

[0051]

[0052] Among them, ρ represents the Euclidean distance between the two center points, b and b gt respectively represent the center points of the predicted box and the ground truth box, c represents the diagonal distance of the smallest closed area that can simultaneously contain the predicted box and the ground truth box, α is the weight function, and v is used to measure the similarity of the aspect ratio.

[0053] 2) Object classification loss and confidence loss

[0054] Both the object classification loss and the confidence loss in the present invention adopt the binary cross-entropy loss function BCE (Binary Cross Entropy) used in YOLOv5:

[0055]

[0056] 3) Angle loss

[0057] For the object detection of rotated boxes in the present invention, a loss of the rotation angle is added. The angle is regarded as a classification task with 180 categories, and the BCE loss function is used. The calculation formula is as follows:

[0058]

[0059] After this step, a rotated box object detection network based on the GhostNet lightweight feature extraction network and fused with the attention mechanism is obtained. Compared with YOLOv5, the number of parameters and the computational complexity are smaller, and the output result has an additional rotation angle, which can better reflect information such as the position, size, and aspect ratio of the detected object.

[0060] S4: Use the training set in step S1 to train the lightweight object detection model obtained in step S3 until the model converges (that is, the loss function of the model no longer continues to decrease), and obtain the trained lightweight object detection model;

[0061] S5: Extract all the convolutional kernels in each layer of convolution in the trained lightweight object detection model in step S4, calculate their L1 norms respectively, remove the convolutional kernels with L1 norms less than the threshold, and then fine-tune the pruned model to obtain the lightweight object detection model after pruning;

[0062] The value of the threshold is determined according to the pruning rate p. Sort the L1 norms of all convolutional kernels from small to large, and take the L1 norm of the N×p-th convolutional kernel as the pruning threshold, where N is the number of model layers. The pruning is shown in the appendix Figure 5Subsequently, the pruned model is fine-tuned. This step mainly removes a large number of redundant parameters and computational amounts existing in the model obtained in step S4. After this step, a fast rotated bounding box object detection model based on the L1 norm pruning algorithm can be obtained, which has a smaller number of parameters and computational amount compared with the model in step S4.

[0063] S6: Replace the NMS algorithm in the lightweight object detection model pruned in step S5 with an NMS algorithm based on RIoU (Rotation-Robust Intersection over Union) to remove duplicate objects in the object detection results and obtain the final model;

[0064] The NMS algorithm of the original YOLOv5 network is for processing regular rectangular boxes, while the output of the present invention is a rotated rectangular box. Therefore, the original NMS processing is no longer applicable to the improved model. In order for the present invention to be applicable to detection boxes with rotation angles, the RIoU (Rotation-Robust Intersection over Union) algorithm is adopted in NMS. The RIoU approach is that for two rotated rectangular boxes p and g, first keep g stationary, rotate p from the initial position gradually until it is perpendicular to g, and calculate the IoU of the two perpendicular boxes. Then keep p stationary, rotate g to be perpendicular to p, and calculate the IoU of the two perpendicular boxes. Take the smaller value of the two IoU values as the IoU of the two rotated boxes. The diagram of RIoU is as shown in the appendix Figure 6 as shown. The results output in step 5 contain a large number of duplicate objects. Therefore, the NMS algorithm based on RIoU is used to remove the duplicate objects in the detection results and obtain the results without duplicate detection boxes.

[0065] S61: Use the final model to detect the remote sensing images in the validation set in step 1 and obtain the detection accuracy of the final model;

[0066] S7: Use the final model to perform fast object detection on newly acquired remote sensing images.

[0067] In this article, the orientation terms such as front, back, up, and down are defined based on the positions of the components in the drawings and the positions of the components relative to each other, only for the sake of clarity and convenience in expressing the technical solution. It should be understood that the use of the orientation terms should not limit the scope of protection requested by this application.

[0068] Without conflict, the above-mentioned embodiments and the features in the embodiments in this article can be combined with each other.

[0069] The above are only the preferred embodiments of the present invention and are not intended to limit the present invention. Any modifications, equivalent replacements, improvements, etc. made within the spirit and principle of the present invention shall be included within the protection scope of the present invention.

Claims

1. A fast target detection method for remote sensing images based on the attention mechanism and pruning, characterized in that The method includes the following steps: S1: Obtain remote sensing image data, organize the data label format to obtain a dataset, and divide the dataset into a training set and a validation set. Perform Mosaic data augmentation and adaptive anchor box calculation on the training set remote sensing images; S2: Select the lightweight feature extraction network GhostNet, fuse the attention mechanism module CBAM in this network to obtain the lightweight feature extraction network GhostNet-CBAM with the attention mechanism. Subsequently, replace the feature extraction network in the object detection model YOLOv5 with the lightweight feature extraction network GhostNet-CBAM to obtain the lightweight object detection network that integrates the lightweight feature extraction network and the attention mechanism; S3: Add the rotation angle to the output part of the lightweight object detection network obtained in step S2, and add a loss of the rotation angle to the loss function of the lightweight object detection network to obtain the lightweight rotated bounding box object detection network; S4: Train the lightweight rotated bounding box object detection network obtained in step S3 using the training set in step S1 until the model converges to obtain the trained lightweight object detection model; S5: Extract all the convolutional kernels in each layer of convolution in the trained lightweight object detection model in step S4, calculate their L1 norms respectively, remove the convolutional kernels with L'1 norms less than the threshold, and then fine-tune the pruned model to obtain the pruned lightweight object detection model; S6: Replace the original NMS algorithm in the pruned lightweight object detection model in step S5 with an NMS algorithm based on RIoU to remove duplicate objects in the object detection results to obtain the final model; S7: Use this final model to perform fast object detection on newly acquired remote sensing images.

2. The fast target detection method for remote sensing images based on the attention mechanism and pruning according to claim 1, wherein The specific process of performing Mosaic data augmentation on each remote sensing image in the training set in step S1 is to randomly scale, randomly crop, and randomly arrange four initial remote sensing images, and then splice them to obtain an image, which is used as an image in the training set.

3. The fast target detection method for remote sensing images based on the attention mechanism and pruning according to claim 1, characterized in that The specific process of performing adaptive anchor box calculation on each remote sensing image in the training set in step S1 is to use the K-means clustering algorithm to calculate several anchor boxes, and then use the genetic algorithm to mutate the target length and width. If the mutation effect is better, set the mutated length and width as the anchor box values, otherwise skip.

4. The fast target detection method for remote sensing images based on the attention mechanism and pruning according to claim 1, characterized in that, In step S1, the image labels in the dataset are organized in the following format: x,y,w,h,θ,class Among them, x and y are the center point coordinates of the rectangular box of the remote sensing image, w, h, and θ are the width, height, and rotation angle of the rectangular box of the remote sensing image respectively, and class is the category of the object in the remote sensing image.

5. The fast target detection method for remote sensing images based on attention mechanism and pruning according to claim 1, wherein After step S6, the method further includes step S61:

6. The fast target detection method for remote sensing images based on the attention mechanism and pruning according to claim 1, characterized in that S61: Use the final model to detect the remote sensing images in the validation set in step 1 and obtain the detection accuracy of the final model. The total loss function of the lightweight object detection network with the rotation angle loss added in step S3 is: L total = λ box L box + λ obj L obj + λ cls L cls + λ theta L theta Among them, L total is the total loss function, L box is the bounding box regression loss, L cls is the classification loss, L obj is the confidence loss, L theta is the angle classification loss, λ box is the proportion of the bounding box loss, λ obj is the proportion of the confidence loss, λ cls is the proportion of the classification loss, λ theta is the proportion of the rotation angle loss.

Images