A UAV infrared small target detection method based on attention mechanism

Abstract

The present invention relates to the field of infrared small target detection technology, and specifically to a method for detecting infrared small targets on unmanned aerial vehicles (UAVs) based on an attention mechanism, comprising the following steps: detection model determination, model revision, and infrared small target detection. This application improves upon the YOLO11 model by utilizing deep separable convolution and DSEAM modules to achieve standardized data distribution and accelerate model training, thereby significantly reducing unnecessary forward propagation information while enhancing infrared small target feature information during operation. Furthermore, the S_RepViTBlock module utilizes a larger receptive field to enhance the ability to extract details of background and target features while capturing infrared small targets, thereby substantially reducing errors and improving the detection quality of infrared small targets.

Description

Technical Field

[0001] The present invention relates to the technical field of infrared small target detection, and in particular to a method for detecting infrared small targets of unmanned aerial vehicles (UAVs) based on an attention mechanism. Background Art

[0002] The rapid development of deep learning and visual inspection technologies has led to the increasing application of drone-mounted visual inspection systems in both civil and military fields. Small drones offer advantages such as low cost, portability, high flight flexibility, good concealment, and high efficiency, making them suitable for applications such as traffic monitoring, accident search and rescue, medical imaging, and military intelligence gathering. Currently, most drones use visible light vision systems, achieving promising results. However, limited computing resources severely impact detection tasks in long-range detection, inclement weather, low light conditions at night, and indistinguishable from surrounding conditions. Infrared imaging technology is unaffected by external light sources and offers strong anti-interference capabilities, long-range detection capabilities, and stable imaging and detection capabilities in all weather conditions.

[0003] However, when using UAV infrared imaging, the targets often appear as densely packed infrared targets due to their flight altitude. This is limited by low resolution and a lack of color and texture features, making detection more difficult. Therefore, improving the performance of infrared small target detection in complex environments using UAV platforms with limited computing resources has attracted widespread research attention.

[0004] Thanks to the continuous efforts of researchers at home and abroad, deep learning-based infrared small target detection algorithms have gradually achieved breakthroughs, moving from low-level to high-level, and from localized to universal. The infrared small target detection algorithm based on multi-layer convolutional fusion proposes a deep learning model that combines multi-layer convolution with multi-receptive field fusion. This model characterizes infrared small targets through multi-layer feature extraction and fusion, adapts to their complex morphology using a multi-receptive field module, combines residual learning and CBAM to adjust network importance, and finally outputs a predicted image of matching size through a full convolutional module.

[0005] Infrared weak infrared small target detection method based on cross-connection and fusion attention mechanism A single-stage infrared weak infrared small target detection algorithm based on cross-connection and fusion attention mechanism is proposed, which improves the detection ability of infrared small weak targets under various complex backgrounds;

[0006] The target detection network with dynamic selection of infrared and visible light image features proposes to add dynamic fusion layer and dynamic selection layer to the network structure, perform feature fusion on multi-source image feature maps multiple times, and use three attention mechanisms to enhance the multi-scale feature maps to screen useful features.

[0007] Although the above algorithm has improved the accuracy of infrared small target detection in specific environments, in complex backgrounds, the inspection target may overlap or be obscured by other objects such as buildings and trees, and lacks the ability to distinguish in dense places, making it difficult to accurately detect infrared small targets; in addition, drones will obtain a large amount of dynamic data in real time during flight, which will cause large delays when computing resources are limited, reducing the timeliness of detection and making it difficult to meet the real-time requirements of certain usage scenarios, so improvements are needed. Summary of the Invention

[0008] The purpose of this invention is to address the shortcomings of the existing technology and propose a method for detecting small infrared targets of drones based on the attention mechanism.

[0009] In order to achieve the above object, the present invention adopts the following technical solutions:

[0010] A method for detecting small infrared targets of drones based on an attention mechanism includes the following steps:

[0011] S1. Detection model determination: The YOLO11 model is used, which consists of three parts: the backbone network Backbone network Neck network and the head network Head;

[0012] S2. Model revision: Add the efficient dual-channel feature enhancement module DSEAM module to the third and sixth layers of the YOLO11 backbone network, which can enhance the feature information of small infrared targets while significantly reducing unnecessary forward propagation information; replace the standard convolution in the 19th and 22nd layers of the original head structure with the S_RepViTBlock module; and use the NWD loss function throughout the model;

[0013] S3. Infrared small target detection: The input image is divided into an S×S grid, and a fixed number of bounding boxes, their confidence scores, and class probabilities are predicted for each grid. The confidence score represents the probability that the bounding box contains an object and is composed of the probability of the target existing and the degree of overlap between the predicted box and the true box. The output vector contains the bounding box coordinates, confidence scores, and class probabilities. The improved YOLO11 model uses a comprehensive loss function to measure prediction errors. Finally, non-maximum suppression (NMS) removes overlapping bounding boxes and retains the boxes with the highest confidence scores. Detection of small infrared targets is then performed.

[0014] Compared with the existing technology, this application improves on the basis of the YOLO11 model, namely, using depthwise separable convolution and DSEAM modules to achieve standardized data distribution, accelerate model training, and make it possible to enhance the feature information of small infrared targets while significantly reducing unnecessary forward propagation information during operation; at the same time, using the S_RepViTBlock module, while capturing small infrared targets, it can use a larger receptive field to enhance the ability to extract details of background and target features, so as to fully reduce errors and improve the detection quality of small infrared targets.

[0015] Preferably, the design concept of the efficient dual-channel feature enhancement module DSEAM module in S2 includes the following steps:

[0016] First, add depthwise separable convolution before the input of ECA, perform depthwise convolution first, and perform batch normalization on the output. This operation can standardize the data distribution and accelerate the training of the model. The batch normalization calculation formula is as follows:

[0017]

[0018] Where: μ is the mean of the current batch, σ 2 is the current batch variance, γ and β are learnable scaling and translation parameters, and ∈ is the anti-zero parameter;

[0019] Second, after output, the result is activated by ReLU to increase the nonlinear expression of the model. The formula of ReLU activation function is as follows:

[0020] x=ReLU(x)=max(0,x)

[0021] ReLU maps all negative values ​​to 0, and keeps positive values ​​unchanged, which can reduce the gradient vanishing phenomenon. Then, point-by-point convolution is performed to adjust the number of channels to achieve linear combination of channel information. The output result is batch normalized and ReLU activated again, and the current feature tensor x is saved as x_orig for subsequent residual connection.

[0022] Third, the information processing results are input into the ECA module for global average pooling to obtain the global features of each channel; the formula is:

[0023] y=AvgPool(x)

[0024] The output tensor y has a shape of (B, c2, 1, 1). Then, y is transformed to a format suitable for one-dimensional convolution input. One-dimensional convolution is performed on y to capture local dependencies between channels. Then, y is restored to its original dimension so that it can be multiplied with the input feature x. A temperature-adjusted Sigmoid activation is performed on y to generate the channel attention weights. The formula for the temperature-adjusted Sigmoid activation function is as follows:

[0025]

[0026] Where T is the temperature parameter, which is set to 0.5 by default. The shape of the final output attention channel weight CA is (B, c2, 1, 1); CA is expanded to the shape (B, c2, H, W) by element-by-element multiplication with the feature x; the input feature x (B, C, H, W) is average pooled AvgOut to obtain (B, 1, H, W) and the maximum pooled MaxOut to obtain (B, 1, H, W); the formula is:

[0027] AvgOut=Mean(x,dim=1,keepdim=True)

[0028] MaxOut=Max(x,dim=1,keepdim=True)

[0029] The output average pooling result AvgOut has a shape of (B, 1, H, W), and the maximum pooling result MaxOut has a shape of (B, 1, H, W); the former can extract global information, and the latter can extract significant features within the channel;

[0030] Fourth, the two pooled feature maps are fused with learnable weight parameters α and β; the fusion formula is as follows:

[0031] Fusion=α×AvgOut+β×MaxOut

[0032] The initial values ​​of α and β are both 0.5. The fused feature Fusion is convolved to extract the spatial attention feature. The convolution kernel size is set to 5, which can well take into account the extraction of features at all scales. The convolution result is activated by Sigmoid, and the weight is compressed between 0 and 1. This allows the network to suppress unimportant areas and highlight important areas, generating weights with spatial attention. The formula is:

[0033]

[0034] The output SA shape is (B, 1, H, W), and then it is expanded to make its shape become (B, c2, H, W); the learnable parameters α and β are used to perform weighted fusion with the twice expanded tensors. In order to ensure the normalization of the weights, the Softmax function is used to normalize [α, β]; formula:

[0035]

[0036] CombineOut=ω α ×x ca +ω β ×xsa

[0037] Among them, xca and xsa are the output tensors of the previous feature x after channel weight expansion and spatial weight expansion, and the shape of the fused feature CombinedOut is (B, c2, H, W);

[0038] Fifth, the fused features are added to the previously saved features x_orig to implement residual connections, which can enhance information flow and prevent gradient vanishing. The final output feature OutPut has a shape of (B, c2, H, W). In practical applications, the DSEAM module can enhance the features of small and densely detected targets in complex environments and reduce the computational complexity and parameter count of the model.

[0039] Furthermore, the features of small and densely detected targets can be enhanced in complex environments, and the computational complexity and parameter amount of the model can be reduced.

[0040] Preferably, the design steps of the S_RepViTBlock module in S2 include:

[0041] First, RepViT achieves a good balance between accuracy and latency on mobile devices by integrating the efficient Vit architecture into CNN. However, in the case of drone infrared detection images, images may become noisy or pixel quality may degrade due to factors such as long distance, lens contamination, and complex mission scenarios, making it difficult to separate objects from the background. The RepViT module was improved and added to YOLO11.

[0042] Second, dilated convolution is introduced in the input part to increase the receptive field. While capturing small infrared targets, the larger receptive field can be used to enhance the ability to extract details of background and target features. In order to enhance the generalization ability of the model and improve its ability to cope with noise and fuzzy features, DropBlock regularization is introduced. DropBlock randomly blocks continuous rectangular areas on the feature map, allowing the model to learn features in a wider context. The formula is:

[0043]

[0044] Where B is the batch size, C is the number of channels, H and W are the height and width respectively, p is the zero probability, and the occlusion size is k×k;

[0045] Third, the improved RepViT module is named S_RepViT.

[0046] Furthermore, the S_RepViT module can extract small infrared targets in the drone's infrared detection image, and enhance the ability to extract details of background and target features through a larger receptive field.

[0047] Preferably, the error in S3 includes position error MSE calculation, confidence error and category error Cross-Entropy metric.

[0048] Furthermore, the measurement error is fully reduced to improve the quality of subsequent inspections.

[0049] The beneficial effects of the present invention are:

[0050] 1. The YOLO series of models is a multi-scale, single-stage target detection algorithm based on regression. Based on a fully convolutional network, it divides the input image into an S×S grid. Each grid predicts a fixed number of bounding boxes, their confidence scores, and class probabilities. This helps reduce the impact of other values ​​on the detection of small infrared targets, thereby improving the quality of detection.

[0051] 2. Through depthwise separable convolution, the output is batch normalized. This operation can standardize the data distribution and accelerate model training. It can achieve feature enhancement of small and dense detection targets in complex environments and reduce the computational complexity and parameter count of the model.

[0052] 3. The S_RepViT module can extract small infrared targets in the drone's infrared detection image, and enhance the ability to extract details of background and target features through a larger receptive field. BRIEF DESCRIPTION OF THE DRAWINGS

[0053] Figure 1 This is the YOLO11 model structure diagram in the UAV infrared small target detection method based on the attention mechanism proposed in this invention;

[0054] Figure 2 This is a diagram of the depth-separable convolution structure in the attention mechanism-based UAV infrared small target detection method proposed in this invention;

[0055] Figure 3 This is the DSEAM module structure diagram of the UAV infrared small target detection method based on the attention mechanism proposed in this invention;

[0056] Figure 4 This is the S_RepViTBlock structure diagram of the attention mechanism-based UAV infrared small target detection method proposed in this invention. DETAILED DESCRIPTION

[0057] The technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, rather than all the embodiments.

[0058] Reference Figure 1-4 ,A method for detecting small infrared targets of UAV based on attention mechanism, includes the following steps:

[0059] S1. Detection model determination: The YOLO11 model is used, which consists of three parts: the backbone network Backbone network Neck network and the head network Head;

[0060] S2. Model revision: Add the efficient dual-channel feature enhancement module DSEAM module to the third and sixth layers of the YOLO11 backbone network, which can enhance the feature information of small infrared targets while significantly reducing unnecessary forward propagation information; replace the standard convolution in the 19th and 22nd layers of the original head structure with the S_RepViTBlock module; and use the NWD loss function throughout the model;

[0061] S3. Infrared small target detection: The input image is divided into an S×S grid, and each grid predicts a fixed number of bounding boxes and their confidence and category probabilities. The confidence indicates the probability that the bounding box contains an object, which is composed of the probability of the target existing and the degree of overlap between the predicted box and the true box. The output vector contains the bounding box coordinates, confidence, and category probability. The improved YOLO11 model uses a comprehensive loss function to measure the prediction error, which includes the position error MSE calculation, confidence error, and category error Cross-Entropy metric. Finally, non-maximum suppression (NMS) removes overlapping bounding boxes and retains the box with the highest confidence. Detection of infrared small targets is performed.

[0062] In the present invention, the design concept of the efficient dual-channel feature enhancement module DSEAM module in S2 includes the following steps:

[0063] First, add depthwise separable convolution before the input of ECA, perform depthwise convolution first, and perform batch normalization on the output. This operation can standardize the data distribution and accelerate the training of the model. The batch normalization calculation formula is as follows:

[0064]

[0065] Where: μ is the mean of the current batch, σ 2 is the current batch variance, γ and β are learnable scaling and translation parameters, and ∈ is the anti-zero parameter;

[0066] Second, after output, the result is activated by ReLU to increase the nonlinear expression of the model. The formula of ReLU activation function is as follows:

[0067] x=ReLU(x)=max(0,x)

[0068] ReLU maps all negative values ​​to 0, and keeps positive values ​​unchanged, which can reduce the gradient vanishing phenomenon. Then, point-by-point convolution is performed to adjust the number of channels to achieve linear combination of channel information. The output result is batch normalized and ReLU activated again, and the current feature tensor x is saved as x_orig for subsequent residual connection.

[0069] Third, the information processing results are input into the ECA module for global average pooling to obtain the global features of each channel; the formula is:

[0070] y=AvgPool(x)

[0071] The output tensor y has a shape of (B, c2, 1, 1). Then, y is transformed to a format suitable for one-dimensional convolution input. One-dimensional convolution is performed on y to capture local dependencies between channels. Then, y is restored to its original dimension so that it can be multiplied with the input feature x. A temperature-adjusted Sigmoid activation is performed on y to generate the channel attention weights. The formula for the temperature-adjusted Sigmoid activation function is as follows:

[0072]

[0073] Where T is the temperature parameter, which is set to 0.5 by default. The shape of the final output attention channel weight CA is (B, c2, 1, 1); CA is expanded to the shape (B, c2, H, W) by element-by-element multiplication with the feature x; the input feature x (B, C, H, W) is average pooled AvgOut to obtain (B, 1, H, W) and the maximum pooled MaxOut to obtain (B, 1, H, W); the formula is:

[0074] AvgOut=Mean(x,dim=1,keepdim=True)

[0075] MaxOut=Max(x,dim=1,keepdim=True)

[0076] The output average pooling result AvgOut has a shape of (B, 1, H, W), and the maximum pooling result MaxOut has a shape of (B, 1, H, W); the former can extract global information, and the latter can extract significant features within the channel;

[0077] Fourth, the two pooled feature maps are fused with learnable weight parameters α and β; the fusion formula is as follows:

[0078] Fusion=α×AvgOut+β×MaxOut

[0079] The initial values ​​of α and β are both 0.5. The fused feature Fusion is convolved to extract the spatial attention feature. The convolution kernel size is set to 5, which can well take into account the extraction of features at all scales. The convolution result is activated by Sigmoid, and the weight is compressed between 0 and 1. This allows the network to suppress unimportant areas and highlight important areas, generating weights with spatial attention. The formula is:

[0080]

[0081] The output SA shape is (B, 1, H, W), and then it is expanded to make its shape become (B, c2, H, W); the learnable parameters α and β are used to perform weighted fusion with the twice expanded tensors. In order to ensure the normalization of the weights, the Softmax function is used to normalize [α, β]; formula:

[0082]

[0083] CombineOut=ω η ×x ca +ω β ×x sa

[0084] Among them, xca and xsa are the output tensors of the previous feature x after channel weight expansion and spatial weight expansion, and the shape of the fused feature CombinedOut is (B, c2, H, W);

[0085] Fifth, the fused features are added to the previously saved features x_orig to implement residual connections, which can enhance information flow and prevent gradient vanishing. The final output feature OutPut has a shape of (B, c2, H, W). In practical applications, the DSEAM module can enhance the features of small and densely detected targets in complex environments and reduce the computational complexity and parameter count of the model.

[0086] In the present invention, the design steps of the S_RepViTBlock module in S2 include:

[0087] First, RepViT achieves a good balance between accuracy and latency on mobile devices by integrating the efficient Vit architecture into CNN. However, in the case of drone infrared detection images, images may become noisy or pixel quality may degrade due to factors such as long distance, lens contamination, and complex mission scenarios, making it difficult to separate objects from the background. The RepViT module was improved and added to YOLO11.

[0088] Second, dilated convolution is introduced in the input part to increase the receptive field. While capturing small infrared targets, the larger receptive field can be used to enhance the ability to extract details of background and target features. In order to enhance the generalization ability of the model and improve its ability to cope with noise and fuzzy features, DropBlock regularization is introduced. DropBlock randomly blocks continuous rectangular areas on the feature map, allowing the model to learn features in a wider context. The formula is:

[0089]

[0090] Where B is the batch size, C is the number of channels, H and W are the height and width respectively, p is the zero probability, and the occlusion size is k×k;

[0091] Third, the improved RepViT module is named S_RepViT.

[0092] In the present invention, a drone is used to take aerial photos of the ground situation. After obtaining the aerial photos, they are transmitted to the improved YOLO11 model to process the aerial photos. First, the YOLO11 model can divide the aerial photos to form an S×S grid. At the same time, the staff revised the YOLO11 model. For example, the DSEAM module can enhance the features of small and dense detection targets in complex environments, and reduce the calculation amount and parameter amount of the model to improve the calculation efficiency.

[0093] At the same time, the S_RepViT module uses a larger receptive field to enhance the ability to extract details of background and target features, so as to improve the ability to capture small infrared targets and obtain small infrared targets more comprehensively from drone aerial images.

[0094] The above description is only a preferred specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any technician familiar with the technical field, within the technical scope disclosed by the present invention, who makes equivalent replacements or changes based on the technical solution and inventive concept of the present invention, should be covered by the scope of protection of the present invention.

Claims

1. A method for detecting small infrared targets of drones based on attention mechanism, characterized in that: The following steps are involved: S1. Detection model determination: The YOLO11 model is used, which consists of three parts: the backbone network Backbone network Neck network and the head network Head; S2. Model revision: Add the efficient dual-channel feature enhancement module DSEAM module to the third and sixth layers of the YOLO11 backbone network, which can enhance the feature information of small infrared targets while significantly reducing unnecessary forward propagation information; replace the standard convolution in the 19th and 22nd layers of the original head structure with the S_RepViTBlock module; and use the NWD loss function throughout the model; The design idea of ​​the efficient dual-channel feature enhancement module DSEAM module in S2 includes the following steps: First, add depthwise separable convolution before the input of ECA, perform depthwise convolution first, and perform batch normalization on the output. This operation can standardize the data distribution and accelerate the training of the model. The batch normalization calculation formula is as follows: Where: μ is the mean of the current batch, σ 2 is the current batch variance, γ and β are learnable scaling and translation parameters, and ∈ is the anti-zero parameter; Second, after output, the result is activated by ReLU to increase the nonlinear expression of the model. The formula of ReLU activation function is as follows: x = ReLU(x) = max(0, x) ReLU maps all negative values ​​to 0, and keeps positive values ​​unchanged, which can reduce the gradient vanishing phenomenon. Then, point-by-point convolution is performed to adjust the number of channels to achieve linear combination of channel information. The output result is batch normalized and ReLU activated again, and the current feature tensor x is saved as x_orig for subsequent residual connection. Third, the information processing results are input into the ECA module for global average pooling to obtain the global features of each channel; the formula is: y=AvgPool(x) The output tensor y has a shape of (B, c2, 1, 1). Then, y is transformed to a format suitable for one-dimensional convolution input. One-dimensional convolution is performed on y to capture local dependencies between channels. Then, y is restored to its original dimension so that it can be multiplied with the input feature x. A temperature-adjusted Sigmoid activation is performed on y to generate the channel attention weights. The formula for the temperature-adjusted Sigmoid activation function is as follows: Where T is the temperature parameter, which is set to 0.5 by default. The shape of the final output attention channel weight CA is (B, c2, 1, 1); CA is expanded to the shape (B, c2, H, W) by element-by-element multiplication with the feature x; the input feature x (B, C, H, W) is average pooled AvgOut to obtain (B, 1, H, W) and the maximum pooled MaxOut to obtain (B, 1, H, W); the formula is: AvgOut=Mean(x,dim=1,keepdim=True) MaxOut=Max(x,dim=1,keepdim=True) The output average pooling result AvgOut has a shape of (B, 1, H, W), and the maximum pooling result MaxOut has a shape of (B, 1, H, W); the former can extract global information, and the latter can extract significant features within the channel; Fourth, the two pooled feature maps are fused with learnable weight parameters α and β; the fusion formula is as follows: Fusion=α×AvgOut+β×MaxOut The initial values ​​of α and β are both 0.

5. The fused feature Fusion is convolved to extract the spatial attention feature. The convolution kernel size is set to 5, which can well take into account the extraction of features at all scales. The convolution result is activated by Sigmoid, and the weight is compressed between 0 and 1. This allows the network to suppress unimportant areas and highlight important areas, generating weights with spatial attention. The formula is: The output SA shape is (B, 1, H, W), and then it is expanded to make its shape become (B, c2, H, W); the learnable parameters α and β are used to perform weighted fusion with the twice expanded tensors. In order to ensure the normalization of the weights, the Softmax function is used to normalize [α, β]; formula: [oh α ,oh β ]=Softmax([a,b]) CombineOut=ω α ×x ca +oh β ×x sa Among them, xca and xsa are the output tensors of the previous feature x after channel weight expansion and spatial weight expansion, and the shape of the fused feature CombinedOut is (B, c2, H, W); Fifth, the fused features are added to the previously saved features x_orig to implement residual connections, which can enhance information flow and prevent gradient vanishing. The final output feature OutPut has a shape of (B, c2, H, W). In practical applications, the DSEAM module can enhance the features of small and densely detected targets in complex environments and reduce the computational complexity and parameter count of the model. The design steps of the S_RepViTBlock module in S2 include: First, RepViT achieves a good balance between accuracy and latency on mobile devices by integrating the efficient Vit architecture into CNN. However, in the case of drone infrared detection images, images may become noisy or pixel quality may degrade due to factors such as long distance, lens contamination, and complex mission scenarios, making it difficult to separate objects from the background. The RepViT module was improved and added to YOLO11. Second, dilated convolution is introduced in the input part to increase the receptive field. While capturing small infrared targets, the larger receptive field can be used to enhance the ability to extract details of background and target features. In order to enhance the generalization ability of the model and improve its ability to cope with noise and fuzzy features, DropBlock regularization is introduced. DropBlock randomly blocks continuous rectangular areas on the feature map, allowing the model to learn features in a wider context. The formula is: Where B is the batch size, C is the number of channels, H and W are the height and width respectively, p is the zero probability, and the occlusion size is k×k; Third, the improved RepViT module is named S_RepViT; S3. Infrared small target detection: The input image is divided into an S×S grid, and a fixed number of bounding boxes, their confidence scores, and class probabilities are predicted for each grid. The confidence score represents the probability that the bounding box contains an object and is composed of the probability of the object existing and the degree of overlap between the predicted box and the true box. The output vector contains the bounding box coordinates, confidence scores, and class probabilities. The improved YOLO11 model uses a comprehensive loss function to measure prediction error. Finally, non-maximum suppression (NMS) removes overlapping bounding boxes and retains the boxes with the highest confidence scores. Infrared small target detection is performed. The errors in S3 include position error MSE calculation, confidence error and category error Cross-Entropy metric.

Images