Infrared weak and small target detection method based on double-branch attention mechanism

Abstract

The invention belongs to the field of target detection, and particularly relates to an infrared weak and small target detection method based on a double-branch attention mechanism. The method comprises the following specific steps: step 1, preparing data sets: performing data enhancement on a first infrared data set and a second infrared data set for pre-training, and performing model fine adjustment on a third infrared data set; 2, constructing a network model, wherein the whole model is composed of a target analysis path, a splicing path and a target reconstruction path; according to the method, a global semantic enhancement large-kernel attention module is introduced into a coding module of a target analysis path, global information is obtained and utilized from four angles of a channel, a transverse direction, a longitudinal direction and a receptive field, and a double-branch single-head attention module is introduced into a splicing path to obtain association degrees among different scales; the purpose of eliminating the semantic difference between the target analysis path and the target reconstruction path is achieved, the difference between the weak and small target and the background is expanded, and the detection precision of the infrared weak and small target is improved.

Description

Technical Field

[0001] The present invention relates to the technical field of target detection, and in particular to an infrared dim small target detection method based on a dual-branch attention mechanism. Background Art

[0002] With the development of infrared imaging technology, infrared systems have been widely used in civil, military, security and other fields. They have the advantages of all-weather operation, strong concealment, long detection distance and strong anti-interference ability. However, due to the limitations of the imaging mechanism, infrared images often present problems such as low contrast, blurred imaging and noise interference. In particular, the pixel ratio of small and weak targets is low, the resolution is poor and the feature details are not obvious, which makes target detection more difficult. In recent years, infrared small and weak target detection based on deep learning has become a hot topic of research. Compared with traditional methods, deep learning has a flexible network structure and powerful automatic feature extraction capabilities, strong robustness, wide versatility and high detection accuracy, and is more suitable for infrared small and weak target detection in complex backgrounds. Although the classic detection network performs well on medium and large targets without overlapping occlusion, the accuracy of infrared small and weak target detection in complex backgrounds still has room for improvement, which shows that this technology has the potential for further optimization in this field.

[0003] The Chinese patent publication number is "CN118691916A", and the name is "Infrared weak target detection method based on multi-dimensional information cross enhancement". This method first inputs the data set into the infrared weak target detection network, uses the multi-dimensional complementary Transformer module to extract the local and global features of the target and the residual full-dimensional dilated convolution module to maintain the existence of the target features in the deep layer of the network, and then uses the global joint feature fusion module and the channel adjustment module in the decoder to fuse feature information from different scales, and outputs the final detection image of the network in combination with the channel adjustment module. Although this method uses a multi-level Transformer to obtain local and global information, since the self-attention mechanism is only performed on feature maps of the same scale, there is a lack of mutual relations between different receptive fields, which will lead to a reduction in the information correlation between information of different scales, making it susceptible to environmental interference, leading to misjudgment and low detection accuracy. Summary of the invention

[0004] 1. Technical issues to be resolved

[0005] In view of the deficiencies in the prior art, the present invention provides an infrared dim small target detection method based on a dual-branch attention mechanism, which solves the problem that the existing infrared dim small target detection method is easily affected by environmental interference, resulting in misjudgment and low detection accuracy.

[0006] (II) Technical solution

[0007] In order to achieve the above-mentioned purpose, the present invention specifically adopts the following technical solutions:

[0008] An infrared dim small target detection method based on a dual-branch attention mechanism comprises the following steps:

[0009] Step 1, prepare data sets: prepare three infrared dim target data sets, merge the first infrared data set with the second infrared data set for data enhancement, i.e. rotation operation, for model pre-training, and use the third infrared data set for model fine-tuning and model evaluation;

[0010] Step 2, build the network model: the entire infrared dim target detection network based on the dual-branch attention mechanism consists of a target parsing path, a splicing path and a target reconstruction path. The target parsing path contains a total of 5 encoding modules, the splicing path consists of a channel splicing operation and a dual-branch single-head attention module, and the target reconstruction path consists of a decoding module and a prediction module;

[0011] Step 3, determine the loss function; select a combination of loss functions suitable for target segmentation and target positioning, and only use the loss that meets the target segmentation in the pre-training stage, and use the loss of target segmentation and target positioning in the fine-tuning stage. When the loss reaches the threshold, the model training is considered complete.

[0012] Step 4, training the network model: input the pre-training data set described in step 1 into the network for training, so that it has the relevant capabilities of infrared weak target detection, and use the model fine-tuning data set for target positioning training after achieving the expected results.

[0013] Step 5, evaluate the network model; evaluate the fine-tuned network model using two evaluation indicators to obtain the performance evaluation of the model from two perspectives: pixel-level prediction and coordinate position, which are used as the criteria for selecting the best model.

[0014] Step 6, solidify the network model parameters; determine the final segmentation model. For example, when performing infrared small target image detection, the infrared image to be detected can be directly input into the trained end-to-end network model to obtain the final infrared small target detection image and coordinate position information.

[0015] Furthermore, the data set used in the pre-training in step 1 adopts a merged data set of the first infrared image data set and the second infrared image data set, and data enhancement is achieved through a rotation operation. By pre-training in a large number of infrared images, the network model has the ability to process infrared images.

[0016] Furthermore, the target parsing path in step 2 includes a total of five encoding modules, encoding module 1 and encoding module 2 are both composed of a residual module, a global semantic enhancement large core attention module and a convolutional downsampling module; encoding module 3 and encoding module 4 are composed of a residual module, a three-dimensional interactive global attention module and an M-pooling downsampling layer; encoding module 5 is composed of a residual module and a three-dimensional interactive global attention module; the residual module is used to extract shallow image feature information and fuse the basic information of each layer in the residual module; the global semantic enhancement large core attention module is used to increase the receptive field, the three-dimensional interactive global attention module realizes three-dimensional feature modeling to obtain the relationship between the three dimensions, and the multi-dimensional feedforward network module integrates the output data obtained from the first two parts; the convolutional downsampling module is used to aggregate shallow feature information and increase the number of channels; the M-pooling downsampling layer is used to extract the most representative data and increase the number of channels of the image.

[0017] Furthermore, the residual module in the target parsing path in step 2 is composed of convolution layer 1 and convolution layer 2, and the residual module of the decoder part is the same as that of the encoder part, each convolution layer is composed of ordinary convolution, batch normalization and leaky linear activation function, and the convolution kernel size is unified to n×n; the global semantic enhancement large core attention module is composed of 3 modules, namely, local feature attention module, three-dimensional interactive global attention module and multi-dimensional feedforward network module; the local feature attention module is composed of 6 large core convolutions of different scales and 1×1 convolution; the three-dimensional interactive attention module is composed of dimensional transformation, global average pooling, channel standard pooling, 1×3 2D convolution and S activation function; the multi-dimensional feedforward network module is composed of 1×1 convolution and Gaussian linear activation function;

[0018] Furthermore, the dual-branch single-head attention module in the splicing path in step 2 is composed of a batch normalization layer, a channel patch embedding layer, a deep convolution, a dual-branch single-head attention layer, and a complementary feedforward network. The channel patch embedding layer obtains the input for the self-attention mechanism according to the number of channels of the feature map; the dual-branch single-head attention layer groups the features and performs self-attention operations on the feature embeddings of different scales, and obtains the attention weights between different scales according to the weights;

[0019] Furthermore, in step 2, the channel patch embedding layer is a 16×16 ordinary convolution; the complementary feedforward network consists of layer normalization, 1×1 ordinary convolution, 3×3 depth convolution, 5×5 depth convolution, linear rectification function, global average pooling layer and S activation function;

[0020] Furthermore, the target reconstruction path in step 2 is composed of five decoding modules and a prediction module 6, wherein decoding module 1 is composed of a residual module and a prediction module; decoding modules 2 to 5 are composed of a residual module, an upsampling layer and a prediction module; the residual module in the decoding module functions as the residual module in the encoding module; the prediction module is used to process multi-channel features into a prediction image with only one channel; the coordinate calculation module performs coordinate calculation on the obtained prediction image;

[0021] Furthermore, the composition of the residual module in the encoding module in step 2 is the same as the residual module in the encoding module; the prediction module is composed of 1×1 convolution and S activation function; the coordinate calculation module is composed of two parts, predicting the number of target pixels and predicting the coordinate pixel position statistics.

[0022] Furthermore, the loss function design in step 3 consists of a binary cross entropy loss and a position sensitive loss to jointly constrain network training. The binary cross entropy loss is used to measure the difference between the true mask and the predicted image pixels, and the position sensitive loss realizes the difference between the true mask coordinates and the predicted image coordinates.

[0023] (III) Beneficial effects

[0024] Compared with the prior art, the present invention provides an infrared dim small target detection method based on a dual-branch attention mechanism, which has the following beneficial effects:

[0025] The present invention designs a new infrared dim small target detection technology that combines a segmentation method with a positioning method, constructs a detection network based on a dual-branch attention mechanism of a U-shaped structure, establishes the connection between local features and global features from four angles of channel, horizontal, vertical and receptive field, and uses a new composite loss function to realize the association between the segmentation function and the positioning function. Currently, there are few network structures that can realize segmentation and positioning at the same time. A new implementation method is provided for it, which achieves the purpose of using positioning to better guide segmentation and improves the detection accuracy of dim small targets in complex backgrounds.

[0026] The present invention introduces a dual-branch single-head attention module in the splicing path between the target parsing path and the target reconstruction path, extracts useful information of the same level of the target parsing path in the channel dimension and splices it into the corresponding target reconstruction path, further improving the network's ability to extract image feature information and reducing the semantic difference between the target parsing path and the target reconstruction path.

[0027] The present invention introduces a global semantic enhancement large core attention module in the encoder module to enhance the correlation between local information and global information, solve the information fusion barrier between local features, effectively expand the receptive field corresponding to each spatial position, and provide more refined data for the self-attention module in the splicing path.

[0028] The present invention proposes a composite loss composed of binary cross entropy loss and position sensitive loss, which enables the network to use position loss to enhance the accuracy of target shape prediction while completing pixel prediction, establish a connection between pixel prediction and position coordinates, and improve the accuracy of infrared small target detection. BRIEF DESCRIPTION OF THE DRAWINGS

[0029] Figure 1 This is the flow chart of infrared dim small target detection based on dual-branch attention mechanism;

[0030] Figure 2 This is the structure diagram of infrared dim small target detection based on dual-branch attention mechanism;

[0031] Figure 3 The structure diagram of the encoding module 1 and the encoding module 2 in the target parsing path of the present invention;

[0032] Figure 4 The structure diagram of the encoding module 3 and the encoding module 4 in the target parsing path of the present invention;

[0033] Figure 5 This is a structural diagram of the encoding module 5 in the target parsing path of the present invention;

[0034] Figure 6 Schematic diagram of the residual module in the encoding module of the present invention;

[0035] Figure 7 Schematic diagram of the convolutional layer in the residual module in the encoding module of the present invention;

[0036] Figure 8 This is a structural diagram of the global semantic enhancement large core attention module of the present invention;

[0037] Fig. 9 This is a schematic diagram of the structure of the local feature attention module in the global semantic enhancement large core attention module;

[0038] Fig.10 This is a schematic diagram of the structure of a three-dimensional interactive global attention module in a global semantic enhanced large core attention module of the present invention;

[0039] Fig.11 Schematic diagram of the structure of the convolution downsampling module in the encoding module of the present invention;

[0040] Fig.12 This is a flow chart of the dual-branch single-head attention module of the present invention;

[0041] Fig.13 This is a schematic diagram of the structure of the dual-branch single-head attention module of the present invention;

[0042] Fig.14 Schematic diagram of the complementary feedforward network structure of the present invention;

[0043] Fig.15 It is a schematic diagram of the characteristic recovery layer structure of the present invention;

[0044] Fig.16 This is a schematic diagram of the structure of the target reconstruction path encoding module of the present invention;

[0045] Fig.17 This is a data chart comparing this method with existing methods. DETAILED DESCRIPTION

[0046] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only 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.

[0047] Embodiment 1

[0048] like Figure 1-17 As shown, an infrared dim small target detection method based on a dual-branch attention mechanism is provided in an embodiment of the present invention, and the method specifically comprises the following steps:

[0049] Step 1, prepare the data set: the data sets used for pre-training are the NUDT-SIRST data set and the NUAA-SIRST data set. The NUDT-SIRST data set and the NUAA-SIRST data set are combined into one and rotated to achieve data enhancement and expand the amount of data for use in the pre-training stage. The data set used in the fine-tuning stage is IRSTD-1K for model fine-tuning and testing.

[0050] Step 2, build a network model; the entire infrared weak target detection network based on the dual-branch attention mechanism consists of a target parsing path, a splicing path and a target reconstruction path, and a dual-branch single-head self-attention module is introduced in the splicing path; the target parsing path includes five encoding modules, encoding module 1 and encoding module 2 are both composed of a residual module, a global semantic enhancement large core attention module and a convolutional downsampling module; encoding module 3 and encoding module 4 are composed of a residual module, a three-dimensional interactive global attention module and an M-pooling downsampling layer; encoding module 5 is composed of a residual module and a three-dimensional interactive global attention module; in the target parsing path, the residual module is used to extract shallow feature information of the image and fuse the basic information of each layer in the residual module, the global semantic enhancement large core attention module is used to increase the receptive field and enhance useful information, and the three-dimensional interactive global attention module realizes the modeling of three-dimensional features to obtain the relationship between the three dimensions; the dual-branch single-head attention module realizes the modeling of fine global information in the channel; the target reconstruction path realizes the decoding of input information and prediction. The dual-branch single-head attention module consists of a batch normalization layer, a channel embedding patch layer, a deep convolution, a dual-branch single-head attention layer, and a complementary feed-forward network;

[0051] Step 3, determine the loss function; in the pre-training stage, select binary cross entropy loss as a constraint to continuously optimize the detection network, effectively separate deeper infrared targets and backgrounds, and enhance the network's global semantic modeling ability. When the network loss reaches the threshold, the model pre-training stage is considered to be completed and the pre-training parameters are saved; in the fine-tuning stage, the loss function is modified to a composite loss function consisting of binary cross entropy loss and position-sensitive loss, so that the network can train positioning capabilities on real infrared data sets, reduce the difference between data enhancement and real infrared data, and enable the network to have the ability to handle complex backgrounds while enriching its positioning capabilities.

[0052] Step 4: Train the network model. Use the pre-training dataset and fine-tuning dataset to provide training data for the model, and select the appropriate optimizer, learning rate, and threshold to complete the iteration of the model parameters.

[0053] Step 5: Evaluate the network model. In the fine-tuning stage, the indicators used are detection probability and false alarm rate as the evaluation criteria for the evaluation model. The two indicators are used to determine the accuracy of the model at the pixel level, providing a basis for selecting the optimal network model parameters.

[0054] Step 6: Save the network weights; select the most appropriate model parameters according to the indicators and use the weight parameters to achieve end-to-end infrared dim target detection tasks.

[0055] Embodiment 2

[0056] like Figure 1-17As shown, an infrared dim small target detection method based on a dual-branch attention mechanism is provided in an embodiment of the present invention, and the method specifically comprises the following steps:

[0057] The dataset prepared in step 1 is: the entire data enhancement process is a pre-trained dataset consisting of the NUDT-SIRST dataset and the NUAA-SIRST dataset, and the pre-trained data is obtained by rotating the data enhancement operation. The dataset used in the fine-tuning stage is IRSTD-1K for model fine-tuning and testing. Among them, the NUAA-SIRST dataset has a total of 427 infrared images containing 480 small targets, many of which are very blurry and hidden in complex backgrounds. The NUDT-SIRST dataset has a total of 1327 infrared images, which contain many target categories, various target sizes, different clutter backgrounds and accurate annotations; the IRSTD-1K dataset has a total of 1000 real images with different targets, including different types of small targets, such as drones, organisms, ships and vehicles, which are captured from long imaging distances at different locations. The dataset covers many different scenes, and the background includes oceans, rivers, fields, mountains, cities and clouds, which are messy and noisy.

[0058] Step 2: Build a network model; Figure 2 As shown in the figure, the entire infrared dim target detection network based on the dual-branch attention mechanism consists of a target parsing path, a splicing path, and a target reconstruction path. The target parsing path consists of five encoding modules, and the dual-branch single-head attention module in the splicing path consists of a batch normalization layer, a deep convolution, a dual-branch single-head attention layer, and a complementary feedforward network; the target reconstruction path consists of five decoding modules, a prediction module, and a coordinate calculation module.

[0059] There are five encoding modules in the target parsing path, among which encoding module 1 and encoding module 2 have the same structure, encoding module 3 and encoding module 4 have the same structure, and encoding module 5 is as follows: Figure 3 , Figure 4 and Figure 5 As shown, the output 1 of each encoding module, that is, x1, is (16, 256, 256), (32, 128, 128), (64, 64, 64), (128, 32, 32) and (256, 16, 16), and the output 2, that is, x2, is (32, 128, 128), (64, 64, 64), (128, 32, 32) and (256, 16, 16). The modules used in the target parsing path are the residual module, the global semantic enhancement large core attention module, the three-dimensional interactive global attention module, the convolutional downsampling layer and the M-pooling downsampling layer. The residual module extracts shallow information and fuses the information of each layer through residual connection. The module consists of two identical convolutional layers, as shown in Figure 6As shown in , the convolution used is unified as a 3×3 convolution kernel, a step size of 1, and a 2D ordinary convolution with an image padding of 1, so that the input and output feature maps are of the same size after convolution. The input feature map passes through ordinary 2D convolution, batch normalization, leaky linear activation function, batch normalization, jump connection, ordinary convolution, batch normalization, leaky linear activation function and jump connection in the residual module in sequence, as shown in Figure 7 The final output feature map is shown, and the residual module formula is as follows:

[0060] x i =Leaky ReLU(BN(conv(Leaky ReLU(BN(conv(x)))))),

[0061] res n =x1+x2, where x i where i=(1,2), res n Where n=(1, 2, 3, 4, 5, 6, 7, 8, 9, 10), LeakyReLU activation function is leaky linear activation function, BN is batch normalization, and conv is a 3×3 ordinary 2D convolution with a stride of 1 and image padding of 1.

[0062] The modules used in encoding module 1 and encoding module 2 are residual module, global semantic enhancement large core attention module and convolution downsampling module. The residual module is as described above, and the global semantic enhancement large core attention module is as follows: Figure 8 As shown in Figure 1, the module consists of a local feature attention module, a three-dimensional interactive global attention module, and a multi-dimensional feedforward network module. The input first enters a 3×3 deep convolution with a step size of 1 and an image padding of 1, and the layer is normalized to obtain O i , respectively input into the local feature attention module and the three-dimensional interactive global attention module, and finally the information of global features and local features is integrated through the multi-dimensional feedforward network module; the local feature attention module is as follows Fig. 9 As shown in the figure, the local feature attention module promotes information interaction between different channels by expanding the feature map channel and using three sets of large kernel convolution layers of different scales to fuse the information between different receptive fields. iThe input is a 1×1 convolution with a step size of 1 that expands the channel three times, and the three-fold expanded feature is evenly divided into three equal parts F1, F2 and F3, which are respectively input into the large kernel convolution layer 1, large kernel convolution layer 2 and large kernel convolution layer 3 to obtain three feature maps L1, L2 and L3. The large kernel convolution layer 1 is a 7×7 depth convolution with a step size of 1 and an image padding of 3, a 9×9 depth convolution with a step size of 1 and an image padding of 16 and a dilation rate of 4, and a 1×1 depth convolution with a step size of 1. The large kernel convolution layer 2 is a 5×5 depth convolution with a stride of 1 and an image padding of 2, a 7×7 depth convolution with a stride of 1 and an image padding of 9 and a dilation rate of 3, and a 1×1 ordinary convolution with a stride of 1. The large kernel convolution layer 3 is a 3×3 depth convolution with a stride of 1 and an image padding of 1, a 5×5 depth convolution with a stride of 1 and an image padding of 4 and a dilation rate of 2, and a 1×1 ordinary convolution with a stride of 1. The calculation formula in the local feature attention module is as follows:

[0063]

[0064] In the formula, is a 7×7 depthwise convolution with stride 1 and image padding 3, 9×9 depthwise convolution with stride 1, image padding 16, and dilation rate 4, 5×5 depthwise convolution with stride 1 and image padding 2, 7×7 depthwise convolution with stride 1, image padding 9, and dilation rate 3, 3×3 depthwise convolution with stride 1 and image padding 1, 5×5 depthwise convolution with stride 1, image padding 4, and dilation rate 2, conv 1×1 is a 1×1 convolution with a stride of 1.

[0065] The second part of the global semantic enhancement large core attention module is the three-dimensional interactive global attention module, such as Fig.10 As shown in Figure 1, this module is used in decoding modules 1 to 5. By splitting the spatial and channel attention into three dimensions of attention, the three-dimensional weights are extracted and recombined independently, so that the weights of the three dimensions are used more effectively. The input of this module is also O i , in the channel attention part, there are two branches, branch one is global average pooling and S activation function, and its output is Avg and Avg weight , branch two-channel standard deviation pooling and S activation function, whose output is Std and Std weight , the four data are weighted to get Weight, and the weight is combined with the input feature O iThe channel attention is adjusted by multiplication to obtain CA, and CA is input into branches three and four for horizontal and vertical attention. In the horizontal attention of branch three, the dimension transformation is performed in turn to change (channel, height, width) to (width, height, channel), and then the global average pooling and channel standard pooling are performed and the two outputs are weighted to obtain (width, 1, 1), and then the dimension transformation is performed to obtain (1, 1, width), and then the ordinary 1×3 2D convolution with a step size of 1 and image filling of (0, 1) is performed, and then the dimension change is obtained to obtain (width, 1, 1). After the S activation function, the last dimension transformation is performed to obtain the horizontal attention weight CW of (1, 1, width). Similarly, the same operation is performed in branch four. The difference from branch three is that the dimension transformation is to interchange the height and channel dimensions to obtain LW. Then the CW and LW matrices are multiplied to obtain the feature tensor CL of (1, height, width), and then the channel attention CA and CL are multiplied by their built-in broadcast mechanism to obtain (channel, height, width), and the final attention weight TD is obtained after the S-type activation function. attention , the calculation formula of this module is as follows:

[0066]

[0067] CA=Weight×O i ,

[0068]

[0069] CW=permute (0,3,2,1) Sigmoid(permute (0,3,2,1) (conv2D 1×3 (permute (0,3,2,1) (F avg&std (permute (0,3,2,1) (CA))))

[0070] LW=permute (0,2,1,3) Sigmoid(permute (0,2,1,3) (conv2D 1×3 (permute (0,2,1,3) (F avg&std (permute (0,2,1,3) (CA))))

[0071]

[0072] TD weight =Sigmoid(CA×CL),

[0073] Where permute represents dimension transformation, conv2D represents 2D convolution, and Sigmoid represents S activation function. Represents matrix multiplication.

[0074] After obtaining the outputs L1, L2, L3 and TD of the local feature attention module and the three-dimensional interactive global attention module weight Then, matrix multiplication is performed on L2 and L3 and then matrix multiplication is performed on L1 to obtain L attention , and then after a 1×1 ordinary 2D convolution with a step size of 1 and a Gaussian linear activation function, it is combined with TD attention Multiply and add to O i The residual is added to the final output GELA, which is used as the input TP of the dual-branch single-head attention module i .

[0075] L attention =L2⊙L3⊙L1,

[0076] GELA=(GELU(f 1×1 (L attention ))⊙TD attention )+O i ,

[0077] Where ⊙ represents the multiplication of the main elements, + represents the addition of the corresponding elements, and GELU represents the Gaussian linear activation function.

[0078] The outputs GELA1 and GELA2 of encoding module 1 and encoding module 2 are respectively input into the convolution downsampling of the two decoding modules to obtain the inputs of decoding module 2 and decoding module 3. The convolution downsampling module is as follows: Fig.11 As shown in the figure, the input features of the convolution downsampling module are sequentially passed through: 1×1 convolution, batch normalization, linear rectification function, 3×3 convolution, batch normalization and 1×1 convolution. The first 1×1 convolution has a step size of 1, which increases the channel to twice the input. The 2D ordinary convolution with a convolution kernel size of 3×3, a step size of 2, and an image padding of 1 realizes the downsampling function. The second 2D ordinary convolution with a convolution kernel size of 1×1 and a step size of 1 does not modify the number of channels, and performs channel information interaction.

[0079] Coding module 3 and coding module 4, such as Figure 4 As shown, the residual module used is the same as encoding module 1 and encoding module 2, and the three-dimensional interactive global attention module used is the three-dimensional interactive global attention module of the global semantic enhancement large core attention module in encoding module 1 and encoding module 2. The attention weight is multiplied by the residual module to obtain the input IP of the dual-branch single-head attention module. i , and then the M pooling downsampling layer gets the input of the next level;

[0080] In the encoding module 5, the residual module and the three-dimensional interactive global attention module are the same as above. The weights of the residual module and the three-dimensional interactive global attention module are multiplied and added to the module input to obtain the input IP of the dual-branch single-head attention module. i .

[0081] After getting the output IP of the five encoding modules 1-5 As the input of the dual-branch single-head attention module, the structure diagram of the dual-branch single-head attention module is as follows Fig.12 As shown in the figure, firstly, batch normalization is performed, deep convolution and residual connection are performed, and the obtained feature map is input into the dual-branch single-head attention layer to realize the exchange of local spatial features and full-level global channel information, and finally the complementary feedforward network is used to fuse the channel information;

[0082] The dual-branch single-head attention layer is as follows Fig.13 As shown, the input channel of this level is divided into channel I in a ratio of 1:3 (i,1) and Channel II I (i,2) , where i in (i, 1) represents the number between the decoder and the encoder and the number 1 represents branch 1. After channel 1 is layer normalized, it passes through convolutional layers 1 to 3 to generate Q (i,1) Feature matrix, K (i,1) The feature matrix and V (i,1) Feature matrix, convolution blocks 1 to 3 are composed of 1×1 convolution with stride 1 and 3×3 depthwise separable convolution with stride 1 and image padding 1. The dual-branch single-head attention layer operation for channel 1 is named branch 1, in which Q (i,1) Feature Matrix and K (i,1) Multiply the feature matrix and use the normalized exponential function to calculate the probability. The probability and V (i,1) The feature matrix is ​​multiplied to obtain the self-attention map of branch one. A 1×1 convolution with a step size of 1 is used to fuse the channel features to obtain the self-attention weight one. Channel one is multiplied by the self-attention weight and input into the complementary feedforward network to obtain the attention output Attention in branch one. (i,1) (Q,K,V), the formula is as follows:

[0083]

[0084] In the formula, softmax is the normalized exponential function and conv is a 1×1 convolution.

[0085] In branch 2, the self-attention operation is performed on the feature map of channel 2 obtained by segmentation and the concatenation of features at all levels. The feature map of channel 2 is input into a convolution layer consisting of a 1×1 convolution with a step size of 1 and a 3×3 depthwise separable convolution with a step size of 1 and an image padding of 1, and Q is obtained. (i,2)The feature matrix, the feature maps of each level of feature concatenation are input into the convolutional layers 5 and 6 respectively. The convolutional layer consists of a 1×1 convolution with a step size of 1 and an output channel the same as channel 2, a 3×3 depth convolution with a step size of 1 and an image padding of 1, and K is obtained. (i,2) The feature matrix and V (i,2) Characteristic matrix, in this branch, Q (i,2) Feature Matrix and K (i,2) Multiply the feature matrix and use the normalized exponential function to calculate the probability. The probability and V (i,2) Multiply the feature matrices, and use a 1×1 convolution with a step size of 1 to fuse the channel features to get the self-attention weight. Use a 1×1 convolution with a step size of 1 to fuse the channel features to get the self-attention weight 2. Add channel 2 to the self-attention weight and input it into the complementary feedforward network to get the attention output of branch 2. (i,2) (Q,K,V), the formula is as follows:

[0086]

[0087] In the formula, Softmax is the normalized exponential function and conv is a 1×1 convolution.

[0088] The complementary feedforward network structure used in branch 1 and branch 2 is as follows: Fig.14 As shown in the figure, the output X is obtained by layer normalization, a 1×1 convolution with a stride of 1 and an output channel that is twice the input channel, and a 3×3 depth convolution with a stride of 1, image padding of 1 and an output channel that is half the input channel. 3×3 Similarly, the input is a 5×5 depth convolution with a stride of 1, image padding of 2, and an output channel of half the input channel to obtain the output X 5×5 , perform linear rectification function on the two branches respectively, concatenate the two outputs of the nonlinear activation, and perform 1×1 convolution with the step size to obtain the output X sc , X sc Input the global average pooling layer and S activation function to get C1, and convert X sc Do a dot multiplication with C1 to Y, and input Attention i Make a residual connection with Y to get X0, completing the information fusion of the local and channel directions of the attention map. The formula is as follows:

[0089]

[0090] C1=Sigmoid(GAP(X sc )),

[0091] Y=C1⊙X sc ,

[0092] X0=Y+Attention i ,

[0093] Where Attention is the output of the dual-branch single-head attention layer, δ is the linear activation function, Sigmoid is the S activation function, Chunk is the channel grouping, GAP is the global average pooling, and The convolution kernel is 3×3 depth convolution and 5×5 depth convolution, concat is channel concatenation, LN is layer normalization, f 1×1 For a 1×1 convolution that reduces the channel by a factor of two, A 1×1 convolution that doubles the channel.

[0094] The obtained X0 input feature recovery layer is restored to the size of the feature map at this level. The structure of this layer is as follows: Fig.15 As shown in Figure 1, X0 is sequentially subjected to bilinear interpolation upsampling, 1×1 ordinary convolution with a step size of 1, batch normalization, and linear rectification function, with multiples of 1, 2, 4, 8, and 16 respectively, and is added to the current level features as the input of the decoding module.

[0095] out=ReLU(BN(conv(upsimple(X0))))+IP i , where ReLU is the linear rectification function, BN is batch normalization, upsimple is bilinear interpolation upsampling, IP i is the original input of this module.

[0096] The target reconstruction path consists of five decoding modules, which are composed of five decoding modules and a prediction module 6. Decoding module 1 consists of a residual module and a prediction module; decoding modules 2 to 5 consist of a residual module, an upsampling layer and a prediction module; the residual module is the same as the residual module in the decoding module. The input obtained by the residual is upsampled by the upsampling layer and the channel is reduced and input into the upper decoding module and the prediction module for prediction. The five sets of feature channels are spliced ​​and input into the prediction module 6 to obtain the final prediction prediction. The prediction formula is as follows:

[0097] prediction=up(Sigmoid(conv(up i ))),

[0098] In the formula, Sigmoid is the S activation function, conv is 1×1 convolution, and up is upsampling.

[0099] The final prediction image passes through the coordinate calculation module to obtain the coordinate information and complete the entire process. This module does not work during the pre-training stage.

[0100] Step 3: Determine the loss function to measure the difference between the predicted target and the true target, and then continuously optimize the network parameters to make it more accurate in predicting the target. During the training process, the loss function used is composed of binary cross entropy loss and position sensitive loss, as follows:

[0101]

[0102] L BCE is the binary cross entropy loss, L local Position sensitive loss.

[0103] For pixel prediction accuracy, we use binary cross entropy loss:

[0104]

[0105] where y i is the probability of predicting correctly, p i is the probability of prediction error.

[0106] The binary cross entropy loss uses the pixel difference between the predicted image and the real image for the infrared small target detection task, which helps to detect more accurate binary images.

[0107] To ensure the accuracy of the predicted target position, we use position-sensitive loss:

[0108]

[0109] Where d p With θ gt To predict the distance and angle of the target in the polar coordinate system, d gt and θ gt The distance and angle of the real target in the polar coordinate system are respectively, taking the predicted target as an example, the calculation formula is as follows:

[0110]

[0111] Where x p and p Predict the target's location coordinate information (x, y).

[0112] The position-sensitive loss is calculated based on the center point of the predicted target and the true target. The position-sensitive loss effectively distinguishes most different positioning errors, making the detector sensitive to different types of positioning errors and locating the target more accurately.

[0113] Step 4: Train the network model. First, put the data set used for pre-training in step 1 into the model for training. The optimizer used in the iterative optimization process of the model is the Adam optimizer. The total number of pre-training rounds is 500 rounds. The number of training rounds in the fine-tuning stage is 1000 rounds. The initial learning rate is 0.001, and the cosine annealing strategy is used to gradually reduce the learning rate to 1×10 -5 If the indicator does not improve during the training process and has reached the loss threshold, the final detection network will be obtained.

[0114] Step 5: Evaluate the network model; select detection probability and false alarm rate as indicators to evaluate the quality of the segmentation model. The indicators of this method and the comparison method data are as follows: Fig.17 As shown, the best model is determined based on the indicator parameters, and its detection probability and false alarm rate are defined as follows:

[0115]

[0116] Where P d is the detection probability, F a is the false alarm rate, N pred is the number of correct predictions, N false is the number of prediction errors, N all is the number of all correct targets, the deviation of the target centroid is less than 3, and we consider the target prediction to be correct.

[0117] Step 6: Save the network weights. Save the network model parameters with the best evaluation indicators. You can input pictures of any size to achieve end-to-end target detection.

[0118] The activation function used in the present invention is as follows:

[0119] (1) Leaky linear activation function:

[0120]

[0121] (2) S activation function:

[0122]

[0123] (3) Normalized exponential function:

[0124]

[0125] (4) Linear rectification function:

[0126]

[0127] (5) Gaussian linear activation function:

[0128]

[0129] The present invention constructs an infrared dim small target detection method based on a dual-branch attention mechanism that combines position information and pixel segmentation. The infrared image is input into the model, and the infrared dim small target detection task can be completed without going through other steps, thereby greatly reducing the complexity of the task.

[0130] Finally, it should be noted that the above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art can still modify the technical solutions described in the aforementioned embodiments or replace some of the technical features therein by equivalents. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present invention shall be included in the protection scope of the present invention.

Claims

1. An infrared dim small target detection method based on a dual-branch attention mechanism, characterized by: The method comprises the following steps: Step 1, prepare data sets: prepare three infrared data sets, perform data enhancement on the first infrared data set and the second infrared data set for model pre-training, and use the third infrared data set for model fine-tuning and model testing; Step 2, building a network model; the entire infrared small target detection network includes a target parsing path, a splicing path, and a target reconstruction path; The target parsing path includes five encoding modules in total. Encoding module 1 and encoding module 2 are both composed of a residual module, a global semantic enhancement large core attention module and a convolutional downsampling module; encoding module 3 and encoding module 4 are composed of a residual module, a three-dimensional interactive global attention module and an M-pooling downsampling layer; encoding module 5 is composed of a residual module and a three-dimensional interactive global attention module; The residual module is used to extract shallow feature information of the image and fuse the basic information of each layer in the residual module; The global semantic enhancement large core attention module includes a local feature attention module, a three-dimensional interactive global attention module and a multi-dimensional feedforward network module, which are used to increase the receptive field and enhance useful information. The local feature attention module uses a large core deep convolution to expand the receptive field, the three-dimensional interactive global attention module models the three-dimensional features to obtain the relationship between the three dimensions, and the multi-dimensional feedforward network module realizes the information utilization and fusion of the above two parts; The splicing path performs a channel splicing operation and a dual-branch single-head attention module. The dual-branch single-head attention module obtains more diverse attention weights by performing different attention operations through channel segmentation, and greatly reduces the amount of calculation by performing patch embedding on each channel. The target reconstruction path is composed of five decoding modules, a prediction module six and a coordinate calculation module to realize the target reconstruction operation and perform binary classification to obtain the final prediction map; the decoding module one is composed of a residual module and a prediction module; the decoding modules two to five are composed of a residual module, an upsampling layer and a prediction module; Step 3, determine the loss function; the detection network is used to complete the object segmentation and positioning tasks, so a single binary cross entropy loss cannot meet the needs, so the binary cross entropy loss is combined with the position sensitive loss, and the physical meaning of the target center point represented by the position sensitive loss is used to further correct the pixel prediction accuracy of the target and achieve more complex constraints; Step 4: training the network model; first, input the pre-training data set in step 1 into the network for training, then fine-tune the model with the fine-tuning data set, and set the preset values ​​such as loss threshold and number of training times; Step 5: Evaluate the network model. The two indicators used in the fine-tuning stage determine the accuracy of the model at the pixel level, providing a basis for selecting the optimal network model parameters. Step 6, solidify the network model parameters; determine the final detection model. For example, when performing infrared small target image detection, the infrared image to be detected can be directly input into the trained end-to-end network model to obtain the final infrared small target detection image and coordinate data information.

2. According to claim 1, a method for infrared dim small target detection based on a dual-branch attention mechanism is characterized in that: In step 1, the first infrared dataset is a NUDT-SIRST dataset, and the second infrared dataset is a NUAA-SIRST dataset. A rotation operation is used to perform image augmentation on the first infrared dataset and the second infrared dataset to expand the datasets. The third infrared dataset is the IRSTD-1K dataset.

3. According to claim 1, a method for infrared dim small target detection based on a dual-branch attention mechanism is characterized in that: In the step 2, the global semantic enhancement large core attention module is composed of a local feature attention module, a three-dimensional interactive global attention module and a multi-dimensional feedforward network module, wherein the local feature attention module is composed of 3 groups of different convolutional layers to realize the function of expanding the receptive field, the three-dimensional interactive global attention module is composed of branch one and branch two for realizing channel attention and branch three and branch four for realizing horizontal attention and vertical attention, branch one is global average pooling and S activation function, branch two is channel standard deviation pooling and S activation function, in branch three and branch four, dimensional transformation and 1×3 2D convolution are used to realize horizontal and vertical attention functions, and matrix multiplication is used to realize global attention functions in channels, horizontal and vertical directions, and the multi-dimensional feedforward network module realizes data processing of the local feature attention module and the three-dimensional interactive global attention module.

4. The infrared dim small target detection method based on the dual-branch attention mechanism according to claim 1 is characterized in that: The dual-branch single-head attention module in the splicing path in step 2 consists of a batch normalization layer, a channel patch embedding layer, a deep convolution, a dual-branch single-head attention layer, and a complementary feedforward network. The dual-branch single-head attention layer groups features and performs self-attention operations on feature embeddings of different scales, and obtains attention weights between different scales according to the weights.

5. The infrared dim small target detection method based on the dual-branch attention mechanism according to claim 1 is characterized in that: The loss function designed in step 3 is composed of binary cross entropy loss and position sensitive loss. The binary cross entropy loss is used to measure the difference between the real mask and the predicted image pixels, and the position sensitive loss realizes the difference between the real mask coordinates and the predicted image coordinates.

Images