Infrared small target detection method based on three-way fusion attention

Abstract

The invention discloses an infrared small target detection method based on three-way fusion attention, and the method comprises the steps: extracting multi-scale semantic information in an infrared original image through a multi-scale feature extraction module, achieving the long-distance context information association, carrying out the enhancement of the multi-scale features through an effective three-way fusion attention feature enhancement module, and carrying out the recognition of the multi-scale features. And finally, outputting through a multi-scale feature fusion positioning module to complete the detection and positioning of the weak and small target. According to the method, the response of a small target in a deep network is enhanced, the target feature dispersion trend is reduced, and better fusion of multi-scale features is facilitated, so that relatively high mIoU and accuracy are achieved.

Description

Technical Field

[0001] The invention relates to target detection technology, in particular to an infrared small target detection method based on three-way fusion attention. Background Art

[0002] Infrared small target detection technology is a very challenging research direction in computer vision and has attracted extensive attention from scientists at home and abroad. With the development of deep learning technology, target detection methods have ushered in a new growth point, and many excellent detection algorithms have emerged. However, most of the current research and applications are designed for large-size targets, and the small and weak target characteristics of infrared small targets have not been fully considered. The information of small targets is easily lost in deep networks, resulting in detection mIoU and accuracy rates lower than those of ordinary visual tasks. Summary of the invention

[0003] The technical problem to be solved by the present invention is to provide an infrared small target detection method based on three-way fusion attention in view of the shortcomings of the existing technology, enhance the response of small targets in deep networks, reduce the diffusion trend of target features, and help multi-scale features to better fuse, thereby achieving higher mIoU and accuracy.

[0004] In order to solve the above technical problems, the technical solution adopted by the present invention is: a method for detecting small infrared targets based on three-way fusion attention, comprising the following steps:

[0005] S1. Obtaining an original infrared image and an annotated image, and constructing a training set using the original image and the annotated image;

[0006] S2. Using the training set as input of a multi-scale feature extraction module to obtain a multi-scale feature map;

[0007] S3, using a three-way fusion attention feature module to enhance the features of all scales in the multi-scale feature map;

[0008] S4. Use the multi-scale feature fusion positioning module to perform multi-scale feature fusion on the enhanced feature map, locate the target, and obtain the target detection result.

[0009] In step S1, the original infrared image and the annotated image are enhanced, and the enhancement process includes normalization, random flipping, random blurring and random cropping.

[0010] The specific implementation process of step S3 includes:

[0011] 1) For a feature map of any scale, the features of that scale are aggregated by channel-wise pooling to obtain the output M C ; After the features of this scale are pooled and aggregated in the height direction, the output M is obtained H; After the features of this scale are pooled and aggregated in the width direction, the output M is obtained W ;

[0012] 2) For M C 、M H 、M W Perform three-way fusion to generate three-way attention weight M;

[0013] 3) Use the three-way attention weight M to weight the original infrared image to obtain the output OUT,

[0014] That is, the enhanced feature map is obtained.

[0015] In step 2), the original infrared image is weighted using the Sigmoid function.

[0016] M=σ(M C )⊙σ(M H )⊙σ(M W ), where σ is the Sigmoid function and ⊙ represents the element-by-element multiplication of the broadcast mechanism.

[0017] The method of the present invention further comprises:

[0018] S5. Calculate the loss function and perform back propagation to update the weight of the small target detection network; the small target detection network includes a multi-scale feature extraction module, a three-way fusion attention feature module and a multi-scale feature fusion positioning module connected in sequence;

[0019] S6. Determine whether the iteration termination condition is met (the number of training rounds reaches the number of iteration rounds epoch, or the training loss loss no longer decreases). If so, obtain an infrared small target detector; otherwise, return to step S2.

[0020] The multi-scale feature fusion positioning module includes a feature fusion module, the feature fusion module is connected to a multi-scale feature cascade module, and the multi-scale feature cascade module is connected to an output convolution layer.

[0021] Compared with the prior art, the present invention has the following beneficial effects: the present invention proposes an effective three-way fusion attention infrared small target detection method, which enhances the response of weak small targets in deep networks, reduces the diffusion trend of target features, and helps multi-scale features to be better integrated. Experiments show that on the SIRST-NUDT dataset, the infrared small target detection method based on effective three-way fusion attention has better mIoU and detection rate than the baseline model. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] Figure 1 This is a flow chart of a target detection method according to an embodiment of the present invention;

[0023] Figure 2 This is a diagram of the infrared small target detection network structure of an embodiment of the present invention;

[0024] Figure 3 This is a structural diagram of the three-way fusion attention enhancement module of an embodiment of the present invention. DETAILED DESCRIPTION

[0025] In order to make the purpose, technical solution and advantages of the embodiments of the present invention clearer, the technical solution 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 part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0026] The embodiment of the present invention proposes a method for detecting small infrared targets based on three-way fusion attention ( Figure 1 ), the method mainly includes three modules: multi-scale feature extraction module, three-way fusion attention feature enhancement module and multi-scale feature fusion module. The multi-scale feature extraction module uses a powerful feature extractor as the backbone network, extracts multi-scale semantic information through multiple layers of neurons, and realizes long-distance context information association; the three-way fusion attention feature enhancement module enhances the extracted multi-scale features, aggregates the target features to reduce the diffusion trend in the deep network, and realizes the robustness of the target deep features. The multi-scale feature fusion module cascades the shallow features rich in spatial information and the deep features that conform to the semantic information to generate multi-layer global features, and fuses them for output.

[0027] The network structure of the embodiment of the present invention is as follows Figure 2 As shown in the figure, it is necessary to first extract the multi-scale semantic information in the infrared original image through the multi-scale feature extraction module to realize the long-distance context information association, and then enhance the multi-scale features through the three-way fusion attention feature enhancement module, and finally output it through the multi-scale feature fusion positioning module to complete the detection and positioning of weak targets.

[0028] The multi-scale feature extraction module is the encoder part of the small object detection network. It uses a powerful feature extractor as the backbone network, takes the original image as input, and outputs a multi-scale feature map. The deep feature map is upsampled to the same scale as the feature map of the previous layer and is spliced ​​from the channel dimension.

[0029] The three-way fusion attention feature enhancement module is embedded in the encoder part, which is mainly divided into a channel descriptor generator, a spatial height descriptor generator, a spatial width descriptor generator, and a three-way fusion (such as Figure 3As shown in Figure 1.3, the three generators represent the feature map effectively from the three dimensions of channel dimension, spatial height and width, aggregate them into three independent direction-aware feature maps, and encode them into three attention maps, each of which captures the long-range dependency of the input feature along one direction. These are the so-called channel dimension, height dimension and width dimension feature descriptors generated. Then, the three feature descriptors are input into the three-way fuser for fusion and multiplied with the input feature map to enhance its representation ability.

[0030] The channel descriptor generator is used to aggregate the channel dimension information of the feature map and generate the attention map in the channel direction. Its input is the feature map with a data dimension of [C*H*W] (C, H, W represent the number of channels, height and width of the feature map respectively), and its output is the descriptor in the channel direction with a data dimension of [C*1*1]. The channel descriptor generator contains 2 global pooling layers and 1 convolution layer, where the pooling layers are average pooling layers and maximum pooling layers respectively, and feature aggregation is performed in the C direction. The convolution layer is a one-dimensional convolution with selectable kernel size; the channel descriptor can be expressed as:

[0031] The spatial height descriptor generator is used to aggregate the spatial height dimension information of the feature map and generate an attention map in the height direction. Its input is the feature map with a data dimension of [C*H*W], and its output is a descriptor in the height direction with a data dimension of [1*H*1]. The spatial height descriptor generator contains 2 global pooling layers and 1 convolution layer, where the pooling layers are average pooling layers and maximum pooling layers, respectively, and feature aggregation is performed in the H direction. The convolution layer is a one-dimensional convolution with a selectable kernel size. The spatial height descriptor can be expressed as:

[0032] The spatial width descriptor generator is used to aggregate the spatial width dimension information of the feature map and generate an attention map in the width direction. Its input is the feature map with a data dimension of [C*H*W], and its output is a descriptor in the width direction with a data dimension of [1*1*W]. The spatial width descriptor generator contains two global pooling layers and one convolution layer, where the pooling layers are average pooling layers and maximum pooling layers, respectively. Feature aggregation is performed in the W direction, and the convolution layer is a one-dimensional convolution with a selectable kernel size. The spatial width descriptor can be expressed as:

[0033] The three-way fusion device is used to fuse the descriptor information in the three directions of channel, spatial height and width, and generate three-way attention weights to weight the input feature map. Its input is a three-way feature map with data dimensions of [C*1*1], [1*H*1], [1*1*W], and its output is a weighted enhanced feature map with data dimensions of [C*H*W]. The three-way fusion device contains one weight generation layer, which is Sigmoid. The three-way fusion attention weight is: M=σ(M C )⊙σ(M H )⊙σ(M W ). Where σ represents Sigmoid, and ⊙ represents the element-by-element multiplication of the broadcast mechanism.

[0034] The multi-scale feature fusion positioning module is the decoder part of the weak target detection network. After stretching the enhanced features of different scales to the same size, it performs feature fusion based on the feature pyramid architecture. The input is a multi-scale feature map, and the output is the result of target positioning detection. It is worth noting that the upsampling and stretching multiples of feature maps of different scales are different.

[0035] The method of this embodiment mainly includes three modules: a multi-scale feature extraction module, an effective three-way fusion attention feature enhancement module, and a multi-scale feature fusion positioning module. Among them, the multi-scale feature extraction module extracts multi-scale semantic information in the original infrared image to realize long-distance context information association, and then enhances the multi-scale features through the effective three-way fusion attention feature enhancement module, and finally outputs through the multi-scale feature fusion positioning module to complete the detection and positioning of weak targets. The specific implementation method is as follows:

[0036] Infrared small target detection network training based on effective three-way fusion attention

[0037] Step 1: Prepare data, including original images and annotated images, verification data and test data, and perform data enhancement processing, including normalization, random flipping, random blurring and random cropping;

[0038] Step 2: Select a powerful feature extraction baseline model, such as ResUNet, DNANet, etc., and build a multi-scale feature extraction module;

[0039] Step 3: Input the image into the multi-scale feature extraction module to obtain a multi-scale feature map;

[0040] Step 4: Construct an effective three-way fusion attention feature enhancement module;

[0041] Step 5: Input the features of each scale into the effective three-way fusion attention module for feature enhancement;

[0042] Step 6: Perform multi-scale feature fusion on each enhanced feature map, locate the target, and obtain the detection result;

[0043] Step 7: Calculate the loss, perform back propagation, and update the weights, where the loss function is SoftIoULoss;

[0044] Step 8: Determine the iteration termination condition. If the iteration is terminated, go to step 9; if the iteration continues, go to step 3;

[0045] Step 9: After training is completed, the infrared small target detector is obtained.

[0046] The computation process of an effective three-way fusion attention mechanism

[0047] (1) Input infrared feature map FEA, whose data dimension is [C*H*W] (C, H, W represent the number of channels, height and width of the feature map respectively)

[0048] (2) After FEA is pooled and aggregated in the channel direction, its output is M C , at this time the data dimension is [C*1*1], that is:

[0049]

[0050] (3) After FEA is pooled and aggregated in the spatial height direction, its output is M H , at this time the data dimension is [1*H*1], that is:

[0051]

[0052] (4) After FEA is pooled and aggregated in the spatial width direction, its output is M W , at this time the data dimension is [1*1*W], that is:

[0053]

[0054] (5) M C 、M H 、M W Perform three-way fusion to generate three-way attention weights, namely:

[0055] M=σ(M C )⊙σ(M H )⊙σ(M W )

[0056] Where σ is the Sigmoid function, ⊙ represents the element-by-element multiplication of the broadcast mechanism, and the data dimension M is [C*H*W].

[0057] (6) The FEA is weighted using the three-way attention weights, and the output is:

[0058] OUT=M*FEA

[0059] The data dimension of OUT is the same as that of FEA.

[0060] Although the preferred embodiments of the present application have been described, those skilled in the art may make other changes and modifications to these embodiments once they have learned the basic creative concept. Therefore, the appended claims are intended to be interpreted as including the preferred embodiments and all changes and modifications falling within the scope of the present application.

[0061] Obviously, those skilled in the art can make various changes and modifications to the present application without departing from the spirit and scope of the present application. Thus, if these modifications and variations of the present application fall within the scope of the claims of the present application and their equivalents, the present application is also intended to include these modifications and variations.

Claims

1. A method for detecting small infrared targets based on three-way fusion attention, characterized in that: The following steps are involved: S1. Obtaining an original infrared image and an annotated image, and constructing a training set using the original image and the annotated image; S2. Using the training set as input of a multi-scale feature extraction module to obtain a multi-scale feature map; S3, using a three-way fusion attention feature module to enhance the features of all scales in the multi-scale feature map; S4. Use the multi-scale feature fusion positioning module to perform multi-scale feature fusion on the enhanced feature map, locate the target, and obtain the target detection result.

2. The infrared small target detection method based on three-way fusion attention according to claim 1 is characterized in that: In step S1, the original infrared image and the annotated image are enhanced, and the enhancement process includes normalization, random flipping, random blurring and random cropping.

3. The infrared small target detection method based on three-way fusion attention according to claim 1 is characterized in that: The specific implementation process of step S3 includes: 1) For a feature map of any scale, the features of that scale are aggregated by channel-wise pooling to obtain the output M C ; After the features of this scale are pooled and aggregated in the height direction, the output M is obtained H ; After the features of this scale are pooled and aggregated in the width direction, the output M is obtained W ; 2) For M C 、M H 、M W Perform three-way fusion to generate three-way attention weight M; 3) The original infrared image is weighted using the three-way attention weight M to obtain an output OUT, that is, an enhanced feature map.

4. The infrared small target detection method based on three-way fusion attention according to claim 3 is characterized in that: In step 2), the original infrared image is weighted using the Sigmoid function.

5. The infrared small target detection method based on three-way fusion attention according to claim 3 is characterized in that: M=σ(M C )⊙σ(M H )⊙σ(M W ), where σ is the Sigmoid function and ⊙ represents the element-by-element multiplication of the broadcast mechanism.

6. The infrared small target detection method based on three-way fusion attention according to any one of claims 1 to 5, characterized in that: Also includes: S5. Calculate the loss function and perform back propagation to update the weight of the small target detection network; the small target detection network includes a multi-scale feature extraction module, a three-way fusion attention feature module and a multi-scale feature fusion positioning module connected in sequence; S6. Determine whether the iteration termination condition is met. If so, obtain an infrared small target detector; otherwise, return to step S2.

7. The infrared small target detection method based on three-way fusion attention according to claim 6 is characterized in that: The iteration termination condition is: the number of iterations reaches the set number of iteration rounds, or the loss function no longer decreases.

8. The infrared small target detection method based on three-way fusion attention according to claim 1 is characterized in that: The multi-scale feature fusion positioning module includes a feature fusion module, the feature fusion module is connected to a multi-scale feature cascade module, and the multi-scale feature cascade module is connected to an output convolution layer.

Images