An unmanned aerial vehicle image detection method based on an improved RT-DETR model

Abstract

The application belongs to the technical field of unmanned aerial vehicle image detection, and in particular to an unmanned aerial vehicle image detection method based on an improved RT-DETR model. The method comprises the following steps: selecting a public unmanned aerial vehicle image dataset, dividing the training set, constructing a small dataset, training and evaluating the improved RT-DETR model on the small dataset, and then implementing the model on the complete dataset to obtain the final data. The application enhances the information of small targets in the fusion process without increasing the number of model parameters. While introducing S2, a new fusion module mainly composed of DSConv is used to replace the fusion module of the benchmark model, which balances the parameter quantity of the model while improving the small target detection accuracy. Compared with the previous algorithm, the improved model has higher accuracy and smaller parameter quantity, and can provide protection for accurate detection of unmanned aerial vehicle images.

Description

Technical Field

[0001] This invention relates to the field of UAV image detection technology, specifically to a UAV image detection method based on an improved RT-DETR model. Background Technology

[0002] With the rapid development of deep learning, machine learning, and other technologies, computer vision technology has been widely applied in various fields. The emergence of numerous detection models has helped solve many complex problems and significantly saved human and financial resources. These detection models are mainly divided into single-stage and two-stage detection models. Generally, single-stage detection models are more conducive to engineering deployment and real-time detection, while two-stage models focus more on detection stability. To address the challenges of complex drone image scenes, numerous small targets, and difficulty in detection, many improved models have also been applied to the field of drone imagery.

[0003] Existing UAV image target detection technologies suffer from several problems, including insufficient representation of small target features, easy occlusion of dense targets, difficulty in distinguishing targets from backgrounds in complex environments, and low detection accuracy or insufficient real-time performance of some detection models. To address these issues, this invention proposes a UAV image detection method based on an improved RT-DETR model. Summary of the Invention

[0004] The purpose of this section is to outline some aspects of the embodiments of the present invention and to briefly describe some preferred embodiments. Simplifications or omissions may be made in this section, as well as in the abstract and title of this application, to avoid obscuring the purpose of these documents; however, such simplifications or omissions should not be construed as limiting the scope of the invention.

[0005] To address the aforementioned technical problems, according to one aspect of the present invention, the present invention provides the following technical solution: A drone image detection method based on an improved RT-DETR model includes the following steps: selecting a publicly available drone image dataset, dividing the training set, constructing a small dataset, training and evaluating the improved RT-DETR model on the small dataset, and then implementing it on the complete dataset to obtain the final data; The improved RT-DETR model includes: introducing the S2 layer output from the backbone network into the fusion module, adding two fusion operations between adjacent layers; designing an S2 enhancement module, selecting a method of generating foreground and background images using high-level feature maps to perform complementary enhancement of the foreground and background; designing a multi-scale query selection module to handle situations with a large number of high-quality tokens and situations with a small number of high-quality tokens; designing a multi-scale fusion module, whose feature maps are sourced from F5, S4, S3, and F2, where S4 and S3 are the output features of the backbone network itself, F5 is constructed using the AIFI module proposed by RT-DETR itself, mainly performing self-attention operations on the S5 feature map; F2 is constructed using the designed enhancement module of S2; and designing a lightweight fusion module DSF to replace the original RT-DETR Fusion module.

[0006] As a preferred embodiment of the UAV image detection method based on the improved RT-DETR model described in this invention, the lightweight fusion module DSF has the following specific structure:

[0007]

[0008]

[0009] in, This is represented as the result of superimposing two input feature maps, forming two parallel branches. and , These are branches that preserve basic features, formed by 1×1 convolutions. One is an enhanced branch consisting of N consecutive DSBlocks and 1×1 convolutions.

[0010] As a preferred embodiment of the UAV image detection method based on the improved RT-DETR model described in this invention, the DSBlock is represented as follows:

[0011] DSBlock by DepthWise Convolution, PointWise Convolution and residual connections are used to construct the system, which is then calculated using the SiLU activation function to achieve spatial information extraction and channel mixing with low parameter overhead.

[0012] As a preferred embodiment of the UAV image detection method based on the improved RT-DETR model described in this invention, the method enhances the learnable scaling coefficients in the residual connections within the enhanced branches and the outer layer fusion. This helps to achieve a balance between feature enhancement capabilities and information fidelity, specifically as follows:

[0013] in Represents the final output. The branch output represents the branch that retains the basic features. This indicates the output of the merged branch. This represents the learnable scaling factor.

[0014] As a preferred embodiment of the UAV image detection method based on the improved RT-DETR model described in this invention, S2 is enhanced using F5 by the foreground / background enhancement module before entering the multi-scale fusion stage. The specific operation is as follows:

[0015]

[0016]

[0017]

[0018] As shown in Equation 8, F5 generates a foreground guidance map M through a simple attention map generator structure, extracts neighborhood information through two 3×3 convolutions, and performs channel alignment using a 1×1 convolution. Represents the Sigmoid function. and These represent 1×1 convolution and 3×3 convolution, respectively. This represents an upsampling operation, using 1-M as the background image, where S2 is multiplied by the foreground and background images respectively to form the foreground flow. and background stream , This represents element-wise multiplication, merging the two streams, and then generating a residual with S2.

[0019] As a preferred embodiment of the UAV image detection method based on the improved RT-DETR model described in this invention, the attenuation rate is calculated according to the sorting of token classification scores to determine the quality of different token segments. 300, 800, 1200, and 2000 are used as the number of queries. The attenuation rate D is calculated by comparing the difference in scores between each level of queries with the score of the previous segment, and then filtering is performed sequentially from 300 to 2000 according to a certain threshold.

[0020] As a preferred embodiment of the UAV image detection method based on the improved RT-DETR model described in this invention, the detection process of the improved RT-DETR model includes the following steps: Step 1: Input the UAV image into the ResNet50 backbone network for feature extraction, and select S2-S5 of the backbone network; Step 2: The S5 feature map is processed by the AIFI module for self-attention, and the S2 feature map is processed by the foreground and background enhancement module for feature enhancement. Together with S3 and S4, they are fed into the multi-scale fusion module MSFM. Step 3: After multi-scale module fusion, four feature maps of the same size as the backbone network are obtained. These feature maps are then flattened into tokens and fed into the decoder. Step 4: The decoder calculates the corresponding category score for each token and sorts the tokens according to the category score; Step 5: Enter the multi-scale query selection module, select the corresponding number of queries based on the token score decay rate, and select a fixed number of queries to enter the decoder; Step 6: After the decoder's calculation, the query is matched with the ground truth box using Hungarian matching to obtain the final detection result, and the loss function is calculated.

[0021] Compared with existing technologies, the advantages of this invention are as follows: This invention aims to obtain real-time and accurate UAV image detection results. To achieve a balance between algorithm model size and real-time performance, RT-DETR is chosen as the baseline model. To enhance the information of small targets during the fusion process without excessively increasing the number of model parameters, a new fusion module with DSConv as its main component replaces the baseline model's fusion module while introducing S2, improving the accuracy of small target detection while balancing the number of model parameters. Simultaneously, S2 is "cleaned" before being fed into the fusion part to avoid introducing noise into the fusion module, thus improving the fusion quality. Before being fed into the decoder, a multi-scale query selection is designed to ensure that high-quality queries are not missed when entering the decoder for different types of images, thereby improving detection accuracy. Compared with previous algorithms, the improved model has higher accuracy and fewer parameters, providing a guarantee for accurate detection of UAV images. Attached Figure Description

[0022] To more clearly illustrate the technical solutions of the embodiments of the present invention, the present invention will be described in detail below with reference to the accompanying drawings and detailed embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort. Wherein: Figure 1 The flowchart below shows a UAV image detection method based on an improved RT-DETR model according to the present invention. Figure 2 This is a diagram of the improved RT-DETR model of the present invention; Figure 3 This is a structural diagram of the multi-scale fusion module of the present invention; Figure 4 This is a structural diagram of the fusion module of the present invention; Figure 5 This is a structural diagram of the background enhancement module of the present invention; Figure 6 The RT-DETR inference effect in the embodiments of the present invention Figure 1 ; Figure 7 This invention aims to improve the RT-DETR inference performance in its embodiments. Figure 1 ; Figure 8 The RT-DETR inference effect in the embodiments of the present invention Figure 2 ; Figure 9 This invention aims to improve the RT-DETR inference performance in its embodiments. Figure 2 ; Figure 10 The reasoning effect in the embodiments of the present invention Figure 2 Local contrast 1; Figure 11 The reasoning effect in the embodiments of the present invention Figure 2 Local contrast 2; Figure 12 The reasoning effect in the embodiments of the present invention Figure 2 Local comparison 3. Detailed Implementation

[0023] To make the above-mentioned objects, features and advantages of the present invention more apparent and understandable, the specific embodiments of the present invention will be described in detail below with reference to the accompanying drawings.

[0024] Secondly, the present invention is described in detail with reference to the schematic diagrams. When detailing the embodiments of the present invention, for ease of explanation, the cross-sectional views illustrating the device structure may be partially enlarged, not according to the usual scale. Furthermore, the schematic diagrams are merely examples and should not limit the scope of protection of the present invention. In addition, actual fabrication should include three-dimensional spatial dimensions of length, width, and depth.

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

[0026] This paper presents a small target detection method based on an improved RT-DETR model, achieving accurate and rapid target detection in UAV images. The method includes the following steps: First, a publicly available UAV image dataset is selected, and the training set is divided into smaller datasets. The improved RT-DETR model is then trained and evaluated on these smaller datasets. Finally, the model is implemented on the complete dataset to obtain the final data. For detailed implementation procedures, please refer to [reference needed]. Figure 1 By improving existing advanced detection models, the detection of small targets in UAV images is enhanced. Applying deep learning to UAV image detection improves the detection method and saves manpower and time.

[0027] Please see Figure 2-5 This invention uses RT-DETR as the baseline model. RT-DETR itself is a relatively advanced single-stage detector and a Transformer-type detector. The advantage of this type of detector is its high accuracy, but it also has disadvantages such as high complexity and many parameters. RT-DETR proposes an AIFI module and a CCFF fusion module, performing self-attention only on the high-semantic S5 layer, and performing an FPN to PAN-style fusion process on S3-S5 layers. This reduces the complexity brought by attention while increasing scale interaction, resulting in good detection results. For the detection results, a token-based query selection method is used. The fused feature map is flattened into tokens. Each token is processed by the encoder to obtain a classification score. Based on the classification score, the top 300 tokens are used as queries and enter the decoder. These queries then undergo Hungarian matching and ground truth matching. However, RT-DETR's detection performance is weak for small targets. Therefore, the model is improved as follows: (1) The S2 layer output from the backbone network is introduced into the fusion module, and the fusion of two adjacent layers is increased, so that more fine-grained spatial details can be seen in the fusion stage. Since the S2 layer is introduced for fusion, the amount of computation increases. Therefore, a lighter fusion module is designed to form a new multi-scale fusion module with the previous overall fusion architecture.

[0028] (2) Design the S2 enhancement module. S2 is the core level of small target detection and needs to be enhanced. Although the fusion module combines high-level semantic information, S2 itself is a low-level feature and has a certain noise. Therefore, the method of generating the foreground and background images using high-level feature maps is chosen to perform complementary enhancement of the foreground and background.

[0029] (3) Design a multi-scale query selection module to handle situations with a large number of high-quality tokens (dense graph) and a small number of high-quality tokens (sparse graph). In addition, after introducing S2, the number of tokens will increase significantly, so it is appropriate to expand the number of queries accordingly.

[0030] Because small targets have a small pixel ratio in the image, after continuous downsampling by the backbone network, S3 can still barely retain some structure, but the information of S4 and S5 is already very blurry. The introduction of S2 can enhance the extraction of small target information. Therefore, a multi-scale fusion module MSFM (Multi-Scale Fusion Module) is designed. The feature maps of this module are sourced from F5, S4, S3 and F4, respectively. S4 and S3 are the output features of the backbone network itself, and F5 is constructed by the AIFI (Attention-based Intra-scale Feature Interaction) module proposed by RT-DETR itself. It mainly performs self-attention operation on the S5 feature map.

[0031] S2, S3, S4, and S5 are the names of the feature maps output by the custom backbone network, and F5 is the name of the feature map S5 after passing through the AIFI module. The backbone network outputs feature maps S1-S5 with resolution from high to low. S2-S5 are selected as feature map inputs to the fusion network, and the lowest layer feature map (with the highest resolution) is named S2, and the highest layer feature map (with the lowest resolution) is named S5.

[0032] However, since the size of the feature map itself is several times or even more than ten times that of the high-level feature map, it also brings a large amount of computation and parameters. Therefore, a lightweight fusion module DSF was designed to replace the original RT-DETR fusion module. DSF (Depthwise Separable Fusion) was designed as a new fusion module, and its specific structure can be represented as follows:

[0033]

[0034]

[0035] in, Represented as two feature maps as input The result of (feature map overlay) is that two parallel branches are formed. and , These are branches that preserve basic features, formed by 1×1 convolutions. An enhancement branch consists of N consecutive DSBlocks (Depthwise Separable Convolution Blocks) and 1×1 convolutions. A DSBlock can be represented as:

[0036] DSBlock by DepthWise Convolution, PointWise The structure consists of convolutions and residual connections, then computed using the SiLU activation function to achieve spatial information extraction and channel mixing with lower parameter overhead. Compared to the original heavily convolutional RepConv+ structure, it is more lightweight and better suited for high-resolution feature fusion scenarios. Furthermore, it enhances the learnable scaling coefficients in the residual connections within the enhanced branches and in the outer fusion layers. This helps to achieve a balance between feature enhancement capabilities and information fidelity, specifically as follows:

[0037] in Represents the final output. The branch output represents the branch that retains the basic features. This indicates the output of the merged branch. This represents the learnable scaling factor.

[0038] As a newly introduced backbone network layer, S2 has the greatest impact on small targets. Therefore, before entering the multi-scale fusion stage, S2 is enhanced using F5 through the Dual Stream Enhancement Module (DSEM). The specific operation is as follows:

[0039]

[0040]

[0041]

[0042] As shown in Equation 8, F5 generates a foreground guidance map M through a simple attention map generator structure, extracts neighborhood information through two 3×3 convolutions, and performs channel alignment using a 1×1 convolution. Represents the Sigmoid function. and These represent 1×1 convolution and 3×3 convolution, respectively. This represents an upsampling operation. Using 1-M as the background image, S2 is multiplied by the foreground and background images respectively to form the foreground flow. and background stream , This represents element-wise multiplication, merging the two streams, and then generating a residual with S2.

[0043] RT-DETR itself selects 300 best tokens for queries. However, due to the introduction of S2 as a fusion component, the number of output tokens is approximately four times the original, making 300 tokens insufficient. To handle images of different types and complexities, a multi-scale query selection method is chosen. Based on the token classification scores, a decay rate is calculated to determine the quality of different token segments, using 300, 800, 1200, and 2000 tokens as the query quantity. The decay rate (D) is calculated by comparing the score difference between queries at each level with the score of the previous segment, and then filtering sequentially from 300 to 2000 according to a certain threshold. For example, in formulas 8 and 9, the average score of tokens ranked 1-300 is used. Subtract the average score of tokens ranked 301-800 Then remove by Obtain the decay rate of tokens ranked 301-800 ,in This represents taking the average value. constant To prevent the denominator from being 0, subsequent query selections will follow the same principle.

[0044]

[0045]

[0046] The testing process consists of the following main steps: Step 1: Input the UAV image into the ResNet50 backbone network for feature extraction, and select S2-S5 of the backbone network; Step 2: The S5 feature map is processed by the AIFI module for self-attention, and the S2 feature map is processed by the foreground and background enhancement module for feature enhancement. Together with S3 and S4, they are fed into the multi-scale fusion module MSFM. Step 3: After multi-scale module fusion, four feature maps of the same size as the backbone network are obtained. These feature maps are then flattened into tokens and fed into the decoder. Step 4: The decoder calculates the corresponding category score for each token and sorts the tokens according to the category score; Step 5: Enter the multi-scale query selection module, select the corresponding number of queries based on the token score decay rate, and select a fixed number of queries to enter the decoder; Step 6: After the decoder's calculation, the query is matched with the ground truth box using Hungarian matching to obtain the final detection result, and the loss function is calculated.

[0047] Example: 1. Experimental Setup: All experiments were trained on an NVIDIA GeForce RTX 4090, following the default parameters of RT-DETR until convergence. The backbone network was ResNet50, with 200 training epochs.

[0048] 2. Dataset: VisDrone-2019-DET was selected as the dataset for the experiment. It contains 6471 training images and 548 validation images. These images were taken by drones at different altitudes and locations. They contain ten categories, and most of the targets in the dataset are small targets.

[0049] 3. Experimental data indicators: Standard COCO indicators AP and AP2 were used. 50 and AP S AP (Average Precision) represents the average precision of the model after taking into account both precision and recall. AP 50 This represents the average accuracy when the predicted bounding box and the ground truth bounding box have a 50% overlap, which is considered a successful match. S This represents the average index obtained by calculating AP on a small target sample set. In the COCO index, a small target sample refers to a target box with a box area of ​​less than 1024 pixels.

[0050] 4. Comparative Experiment: The improved model and the baseline model RT-DETR were trained in the same experimental environment for the same number of epochs, and the best results were compared. The results are shown in Table 1. The data shows that with a slightly smaller number of parameters, AP increased by 2.9%. 50 It grew by 3.5%, AP S It increased by 4%.

[0051] Table 1. Results of the comparative experiment

[0052] 5. Comparison of Results: like Figure 6 , Figure 7 As shown in the comparison, the number of small targets in this figure is relatively small, and most of them are medium-sized targets. The improved RT-DETR model ensures the correct detection of small targets while guaranteeing the correct detection of medium-sized targets.

[0053] like Figure 8 , 9 As shown, the target in this image is quite dense. A local comparison is performed on three densely packed areas of the image, as follows: Figure 10 , Figure 11 , Figure 12 The image on the right shows the effect of the improved RT-DETR. Clearly, in many parts, the improved RT-DETR performs relatively well in detection.

[0054] Although the present invention has been described above with reference to embodiments, various modifications can be made and components can be replaced with equivalents without departing from the scope of the invention. In particular, as long as there is no structural conflict, the features in the disclosed embodiments can be combined with each other in any manner. The lack of an exhaustive description of these combinations in this specification is merely for the sake of brevity and resource conservation. Therefore, the present invention is not limited to the specific embodiments disclosed herein, but includes all technical solutions falling within the scope of the claims.

Claims

1. A UAV image detection method based on an improved RT-DETR model, characterized in that, Includes the following steps: We selected a publicly available drone image dataset, divided the training set to construct a small dataset, trained and evaluated the improved RT-DETR model on the small dataset, and then implemented it on the complete dataset to obtain the final data. The improved RT-DETR model includes: introducing the S2 layer output from the backbone network into the fusion module, adding two fusion operations between adjacent layers; designing an S2 enhancement module, selecting a method of generating foreground and background images using high-level feature maps to perform complementary enhancement of the foreground and background; designing a multi-scale query selection module to handle cases with a large number of high-quality tokens and cases with a small number of high-quality tokens; designing a multi-scale fusion module, whose feature maps are sourced from F5, S4, S3, and F2, where S4 and S3 are the output features of the backbone network itself, F5 is constructed using the AIFI module proposed by RT-DETR itself, mainly performing self-attention operations on the S5 feature map; F2 is constructed using the designed enhancement module of S2; and designing a lightweight fusion module DSF to replace the original RT-DETR Fusion module.

2. The UAV image detection method based on the improved RT-DETR model according to claim 1, characterized in that, The lightweight fusion module DSF has the following specific structure: ； ； ； in, This is represented as the result of superimposing two input feature maps, forming two parallel branches. and , These are branches that preserve basic features, formed by 1×1 convolutions. One is an enhanced branch consisting of N consecutive DSBlocks and 1×1 convolutions.

3. The UAV image detection method based on the improved RT-DETR model according to claim 2, characterized in that, The DSBlock is represented as: ； DSBlock by DepthWise Convolution, PointWise Convolution and residual connections are used to construct the system, which is then calculated using the SiLU activation function to achieve spatial information extraction and channel mixing with low parameter overhead.

4. The UAV image detection method based on the improved RT-DETR model according to claim 2, characterized in that, Enhance the learnable scaling factor in internal residual connections of branches and outer fusion. This helps to achieve a balance between feature enhancement capabilities and information fidelity, specifically as follows: ； in Represents the final output. The branch output represents the branch that retains the basic features. This indicates the output of the merged branch. This represents the learnable scaling factor.

5. The UAV image detection method based on the improved RT-DETR model according to claim 1, characterized in that, Before entering the multi-scale fusion stage, S2 is enhanced using the foreground / background enhancement module with F5. The specific operation is as follows: ； ； ； ； As shown in Equation 8, F5 generates a foreground guidance map M through a simple attention map generator structure, extracts neighborhood information through two 3×3 convolutions, and performs channel alignment using a 1×1 convolution. Represents the Sigmoid function. and These represent 1×1 convolution and 3×3 convolution, respectively. This represents an upsampling operation, using 1-M as the background image, where S2 is multiplied by the foreground and background images respectively to form the foreground flow. and background stream , This represents element-wise multiplication, merging the two streams, and then generating a residual with S2.

6. The UAV image detection method based on the improved RT-DETR model according to claim 1, characterized in that, Based on the sorting of token classification scores, the decay rate is calculated to determine the quality of different token segments. 300, 800, 1200, and 2000 are used as the number of queries. The decay rate D is calculated by comparing the difference in scores between queries at each level with the score of the previous segment, and then filtering is performed sequentially from 300 to 2000 according to a certain threshold.

7. The UAV image detection method based on the improved RT-DETR model according to claim 1, characterized in that, The improved RT-DETR model detection process includes the following steps: Step 1: Input the UAV image into the ResNet50 backbone network for feature extraction, and select S2-S5 of the backbone network; Step 2: The S5 feature map is processed by the AIFI module for self-attention, and the S2 feature map is processed by the foreground and background enhancement module for feature enhancement. Together with S3 and S4, they are fed into the multi-scale fusion module MSFM. Step 3: After multi-scale module fusion, four feature maps of the same size as the backbone network are obtained. These feature maps are then flattened into tokens and fed into the decoder. Step 4: The decoder calculates the corresponding category score for each token and sorts the tokens according to the category score; Step 5: Enter the multi-scale query selection module, select the corresponding number of queries based on the token score decay rate, and select a fixed number of queries to enter the decoder; Step 6: After the decoder's calculation, the query is matched with the ground truth box using Hungarian matching to obtain the final detection result, and the loss function is calculated.

Images