Small Target Detection Method Based on Graph Attention Network

Abstract

The present invention proposes a small target detection method based on a graph attention network, and the implementation steps are as follows: obtaining a training sample set and a test sample set; constructing a small target detection model based on a graph attention network; initializing parameters; training the small target detection network; obtaining the trained small target detection network; and obtaining small target detection results. In the present invention, the pixel-to-graph attention module refines semantic features by utilizing the potential correlation between adjacent pixel points, and this potential correlation between pixels can be used to improve the detection performance of the model. The gated cross-attention feature aggregation module learns the multi-scale similarity of the target contour features obtained from different network layers through a cross-attention mechanism and outputs refined target contour features to constrain the prediction, which can better aggregate information features, effectively solves the problem that the existing methods gradually lose their structural details after the data passes through stacked convolutional and downsampling layers, and improves the accuracy of small target detection.

Description

Technical Field

[0001] The present invention belongs to the technical field of image processing, and particularly relates to a small target detection method based on a graph attention network, which can be used in the fields of maritime search and rescue, fire search and rescue, etc. Background Art

[0002] During maritime search and rescue and fire search and rescue, due to frequent interference from extreme weather and harsh environments such as heavy fog and heavy rain, and since the sensing system often images the target at a relatively long distance, the obtained targets basically appear in the form of small targets. In general visible light target detection, small targets are usually defined as those smaller than 1% of the image. In addition, since the proportion of positive samples in the entire image is very low, if a general target detection method is used to detect the image, serious class imbalance and sample imbalance problems will occur. This poses a huge challenge to detecting small targets and recognizing their shapes. The characteristics of traditional methods relying on manual and local information limit the performance of these methods. In existing algorithms for application scenarios mainly focused on small target detection, especially when the background noise accounts for a relatively large proportion, problems such as missed detection or false alarms are likely to occur, and there are still challenges in detecting small targets.

[0003] To solve these problems, researchers have continuously improved small target detection technologies. For example, in the patent document "A Small Target Detection Method Based on Improved Multi-Scale Feature Fusion of CenterNet" (Patent Application No.: CN202211459753.X, Publication No.: CN115631400A) applied by Harbin University of Science and Technology, a small target detection method based on improved multi-scale feature fusion of CenterNet is proposed. By replacing the high-resolution backbone network of the original model, high-quality feature maps are obtained. The high-resolution backbone network introduces a fine-grained multi-scale feature fusion module, which divides the input feature map into four groups with the same specifications. Except for the first group of feature maps, the output of each subsequent group of feature maps is added to the output of the previous group of feature maps to obtain the output, and then a convolution operation is performed. After the convolution is completed, the outputs of the four groups of feature maps are concatenated in the channel dimension, and then a coordinate attention mechanism and a Dropout operation are introduced. This method improves the model's ability to extract characteristic information by using a high-resolution feature extraction network; and increases the detection robustness of the model to multi-scale targets from two perspectives: hierarchical multi-scale feature fusion and fine-grained multi-scale feature fusion, thereby improving the detection performance of small targets. However, since the data gradually loses its structural details after passing through stacked convolutional and downsampling layers, the probability of missed detection or false alarms is still relatively high, affecting the further improvement of the detection accuracy under noise interference. Summary of the Invention

[0004] The object of the present invention is to overcome the defects existing in the above-mentioned prior art, and a small target detection method based on a graph attention network is proposed, aiming to improve the detection accuracy of small targets on the premise of ensuring the detection efficiency.

[0005] To achieve the above object, the technical solution adopted by the present invention includes the following steps:

[0006] (1) Obtain a training sample set and a test sample set:

[0007] Obtain N infrared small target images including multiple target categories, and label the small target objects in each infrared small target image. Then, a total of K small targets and their corresponding labels in more than half of each target category are formed into a training sample R, and the remaining N - K small target images and their corresponding labels are formed into a test sample set E, where N ≥ 500.

[0008] (2) Construct a small target detection network model O based on a graph attention network:

[0009] Construct a small target detection network model O including an encoder, a pixel-to-graph attention module PGA, and a decoder in series as the backbone network, and a gated cross-attention feature aggregation module GCFA including M cascades and a head block connected to each GCFA as the branch network; where the encoder includes M cascaded residual blocks and a max pooling layer loaded between the first and second residual blocks; the decoder includes M cascaded upsampling deconvolution modules and a head block; the branch network is loaded between the input end of the encoder and the input end of the head block in the backbone network; the first residual block, the max pooling layer, the m-th residual block, and the (M - 1)-th residual block in the encoder are respectively connected to the output ends of the M-th, (M - 1)-th, (M - m)-th, and the first upsampling deconvolution module in the decoder; M GCFAs are respectively connected to the output ends of the first to (M - 1)-th upsampling deconvolution modules in the PGA and the decoder.

[0010] (3) Initialize the parameters:

[0011] Initialize the number of iterations as t, the maximum number of iterations as T, T ≥ 1000, and the weights and bias parameters of the small target detection network model O in the t-th iteration are respectively w t 、b t ,and let t = 0, O t = O; t

[0012] (4) Train the small target detection network model:

[0013] Randomly select P training samples from the training sample set R with replacement as the input of the small target detection network model O for forward propagation to obtain P small target detection results, where 1 ≤ P ≤ K;

[0014] (5) Obtain the trained small target detection network:

[0015] Based on the P small target detection results obtained in step (4), update the weights and bias parameters w t of the small target detection network model O t and b t and determine whether t ≥ T holds. If so, obtain the trained small target detection network model O * , otherwise, set t = t + 1 and execute step (4);

[0016] (6) Obtain the small target detection results:

[0017] Use the test sample set E as the input of the trained small target detection network model O * for forward propagation to obtain the small target detection results corresponding to N - K test samples.

[0018] Compared with the prior art, the present invention has the following advantages:

[0019] In the process of training the small target detection network model and obtaining the small target detection results, the PGA module refines the semantic features by utilizing the potential correlation of adjacent pixel points to improve the detection performance of the model. The GCFA module learns the multi-scale similarity of the target contour features obtained from different network layers through the cross-attention mechanism and outputs refined target contour features to constrain the prediction, which can better aggregate information features; avoiding the defect of high missed detection rate caused by the gradual loss of its structural details when data passes through stacked convolutional and downsampling layers in the prior art, effectively improving the accuracy of small target detection. BRIEF DESCRIPTION OF THE DRAWINGS

[0020] Figure 1 is the implementation flowchart of the present invention;

[0021] Figure 2 is the structural schematic diagram of the small target detection network model adopted in the embodiment of the present invention;

[0022] Figure 3 is the structural schematic diagram of the pixel-to-graph attention module adopted in the embodiment of the present invention;

[0023] Figure 4 is the structural schematic diagram of the gated cross-attention feature fusion module adopted in the embodiment of the present invention. DETAILED DESCRIPTION OF THE INVENTION

[0024] The present invention will be further described in detail below in conjunction with the accompanying drawings and specific embodiments.

[0025] Referring to Figure 1 , the present invention includes the following steps:

[0026] Step 1) Obtain a training sample set and a test sample set:

[0027] Obtain N small target images of multiple target categories included in the IRSTD-1k dataset. The input image size is 512*512, and the number of channels is 3. And label the small target objects in each small target image. Then, a total of K small targets and their corresponding labels in more than half of each target category are composed into a training sample R, and the remaining N-K small target images and their corresponding labels are composed into a test sample set E. In this embodiment, N = 1000 and K = 600.

[0028] Step 2) Construct a small target detection network O as shown in Figure 2 :

[0029] Construct a small target detection network model O including a backbone network and a branch network; where the backbone network includes an encoder, a pixel-to-graph attention module PGA, and a decoder connected in series; the branch network includes three cascaded gated cross-attention feature aggregation modules GCFA; the encoder includes three cascaded residual blocks and a max pooling layer loaded between the first and second residual blocks; the decoder includes three cascaded upsampling transposed convolution modules and a head block; the first residual block in the encoder is connected to the output end of the third transposed convolution layer of the decoder, the max pooling layer of the encoder is connected to the output end of the second transposed convolution layer of the decoder, the second residual block in the encoder is connected to the output end of the first transposed convolution layer of the decoder, and the third residual block in the encoder is only connected to the PGA module; the three GCFA modules are respectively connected between the PGA and the first to second upsampling transposed convolution module layers in the decoder. The first GCFA module is also directly connected to the input image, and the last GCFA module is connected to the output end of the third transposed convolution layer of the decoder through a head block, where:

[0030] As shown in Figure 3 , the pixel-to-graph attention module PGA has a structure including a graph attention module and a spatial attention module arranged in parallel; the graph attention module uses a multi-layer graph convolutional neural network GCN to refine the high-level semantic features extracted by the encoder using the correlation between adjacent pixel points of the feature map output by the encoder; the spatial attention module includes a max pooling layer, an average pooling layer, and a convolutional layer stacked in sequence.

[0031] The three GCFA modules are used to obtain the multi-scale similarity of the target contour features and output refined target contour features to constrain the prediction, as shown in Figure 4As shown in the figure, the gated cross-attention feature aggregation module GCFA includes two branch structures arranged in parallel. Each branch structure contains multiple stacked 1×1 convolutional layers with normalization layers and GeLU activation functions, and a depth convolutional layer. After the depth convolutional layer, the obtained attention map enters the cross-attention module for processing, and the two branches respectively output gated features and contour features;

[0032] The specific parameter settings are as follows: the convolutional kernel sizes of the convolutional layers in the first, second, and third downsampling modules are 3*3, the convolutional layer sizes in the upsampling transposed convolutional module are all 3*3, and the convolutional kernel size of the convolutional layer in the PGA module is 3*3; the convolutional layer size of the GCFA module is 1*1. All normalization layers use layer normalization, and all activation functions use GeLU activation functions.

[0033] Step 3) Initialize parameters:

[0034] Initialize the number of iterations as t, the maximum number of iterations as T, T≥1000, and the weights and bias parameters of the graph attention network O t in the t-th iteration are w t and b t respectively. Set the learning rate α to 0.04, and let t = 0, O t = O;

[0035] Step 4) Train the small object detection network:

[0036] Randomly select P training samples from the training sample set R with replacement as the input of the graph attention network model O for forward propagation. In this embodiment, P = 32.

[0037] (4a) The encoder performs multi-layer downsampling on each training sample. The 32 images input to the first residual block of the encoder are subjected to downsampling feature extraction to obtain 32 local feature maps with a depth of 1, a size of 128*128, and a channel number of 16. After passing through the max-pooling layer, the spatial resolution of the feature map is reduced to obtain a local feature map with a size of 64*64, a channel number of 16, and a depth of 1. The second residual block extracts features from the feature map with a depth of 1 to obtain 32 local feature maps with a depth of 2, a size of 32*32, and a channel number of 32. The third downsampling residual block extracts features from the feature map with a depth of 2 to obtain 32 local feature maps with a depth of 3, a size of 16*16, and a channel number of 64;

[0038] (4b) The pixel-to-graph attention module PGA inputs 32 feature maps with a depth of 3 into two modules, the graph attention module and the spatial attention module. In the graph attention module, F4 is transformed into a graph structure composed of pixel points, and then the high-level semantic features are reshaped and further embedded into the corresponding node features. 32 feature maps are generated using multiple graph convolutional networks GCN:

[0039] In the spatial attention module, 32 feature maps are successively passed through a max pooling layer MaxPool, an average pooling layer AveragePool, and a convolutional layer Conv with a 3×3 convolutional kernel to extract 32 feature maps:

[0040] Finally, the features output by the two modules are fused one-to-one using the attention mechanism Figure 1 to output 32 refined feature maps:

[0041] (4c) The decoder upsamples the 32 refined feature maps and enhances the feature information through skip connections between layers of the encoder. After the 32 refined feature maps are input into the first deconvolutional layer of the decoder, they are combined with 32 feature maps with a depth of 2 and input into the second deconvolutional layer. The output result is then combined with 32 feature maps with a depth of 1 output by the max pooling layer and input into the third deconvolutional layer. The output result is then combined with 32 feature maps with a depth of 1 at the output end of the first residual block of the encoder to obtain the final 32 global feature maps;

[0042] (4d) Three GCFAs in the two branch structures of the branch network preprocess and perform depth convolution on their two-way inputs to obtain three groups of attention maps respectively output by the two branch structures; the cross-attention module performs cross-attention fusion on multiple attention maps of the two paths to obtain 3 refined contour features The head block fuses the M refined contour features to obtain the contour feature

[0043] The first GCFA module obtains the contour features extracted from the input samples using the Laplacian kernel and the refined contour features output by the PGA module, performs preprocessing and depth convolution to obtain three groups of attention maps respectively output by two branch structures, and then enters the cross-attention module for cross-attention fusion to obtain 32 contour features and 32 gating features to be delivered to the next GCFA module. The second GCFA module obtains the 32 gating features output by the first GCFA module, obtains multiple extracted 32 contour features from the output end of the first deconvolution layer, performs preprocessing and depth convolution to obtain three groups of attention maps respectively output by two branch structures, and then enters the cross-attention module for cross-attention fusion to obtain 32 gating features to be delivered to the next GCFA module and the current 32 contour features. The third GCFA module obtains the 32 gating features output by the second GCFA module, obtains multiple extracted 32 contour features from the output end of the second deconvolution layer, performs preprocessing and depth convolution to obtain three groups of attention maps respectively output by two branch structures, and then enters the cross-attention module for cross-attention fusion to obtain 32 gating features and the current 32 contour features; 3 GCFA modules obtain 3 refined contour features; the head block fuses the 3 refined contour features to obtain the contour features output by the branch structure.

[0044] (4e) The 32 global feature maps are multiplied by the 32 contour features obtained by the three GCFA modules and then added to themselves to obtain 32 fused feature maps with a size of 512*512 and 64 channels.

[0045] Step 5) Obtain the trained small object detection network:

[0046] Using the 32 small object detection results obtained in step (4), update the weights and bias parameters w t and b t of the small object detection network O t to obtain the network model O t of this iteration:

[0047] (5a) Adopt the Dice loss function L Dice , and calculate the contour loss value output by the branch structure using the predicted values and actual values of the output results of multiple GCFA modules obtained through the head block in the branch network At the same time, adopt the BCE loss function L BCE and calculate the contour loss value output by the branch structure using the predicted values and actual values of the output results of multiple GCFA modules obtained through the head block in the branch network and the weight coefficient Then add and to obtain the contour loss value L con; Then, calculate the Dice loss of the backbone network using the predicted and actual values of the global feature map obtained through the decoder head block. Use the BCE loss obtained from the predicted and actual values of the output result of the PGA module. The Dice loss of the backbone network The BCE loss of the output result of the PGA module And the contour loss L con Combine them to obtain the total loss L of the final training objective t ;

[0048]

[0049]

[0050]

[0051]

[0052]

[0053]

[0054] Among them, X1, X2, and X3 represent the predicted values of the output results of the branch network, the backbone network, and the PGA module respectively, Y1, Y2, and Y3 represent the actual values of the network output results, the backbone network output results, and the PGA module output results respectively, ∩ represents the intersection operation, ln(·) represents the natural logarithm function; ω represents the weight coefficient, which is taken as 1 in this embodiment; λ is the hyperparameter for balancing the two losses, which is taken as 5 in this embodiment; is the Dice loss of the predicted output result of the backbone network, L con is the final contour loss value output by the GCFA module, is the BCE loss obtained from the output end of the PGA module;

[0055] (5b) Calculate L through the chain rule t The partial derivative of the weight parameter w t and the bias parameter b t The partial derivative of and And according to For w t , b t Update:

[0056]

[0057]

[0058] Among them, w′ t , b′t Represents O t The weight value and bias parameter w t and b t The update result, α represents the learning rate, which is taken as 0.04 in this embodiment. It is judged whether t≥T holds. If so, the trained graph attention network O * is obtained. Otherwise, let t=t + 1 and execute step (4);

[0059] Step 6) Obtain the small target detection result:

[0060] Use the test sample set E as the input of the trained small target detection model O * to perform forward propagation and obtain the small target detection results corresponding to 400 test samples.

[0061] The technical effects of the present invention will be described below in combination with simulation experiments:

[0062] 1. Simulation conditions and content:

[0063] The hardware platform for the simulation experiment is: the processor is an Intel(R) Core i9-9900K CPU with a main frequency of 3.5GHz, the memory is 32GB, and the graphics card is an NVIDIA GeForce RTX 2080Ti. The software platform for the simulation experiment is: the Ubuntu 16.04 operating system, the python version is 3.7, and the Pytorch version is 1.7.1.

[0064] A comparative simulation of the intersection over union IoU between the present invention and an existing small target detection method based on the improvement of multi-scale feature fusion of CenterNet is carried out, and the comparison is made at the pixel level and the object level. The results are shown in Table 1. The calculation formula of IoU is:

[0065]

[0066] 2. Analysis of simulation results:

[0067] Referring to Table 1, the intersection over union of the prior art and the present invention on the IRSTD-1k dataset are 62.05% and 66.59% respectively. The experimental results show that compared with the prior art, the intersection over union IoU of the present invention has increased by 4.54%, indicating that the detection accuracy has been effectively improved.

[0068]

[0069] Table 1

Claims

1. A small target detection method based on graph attention network, characterized in that, Including the following steps: (1) Obtain a training sample set and a test sample set: Obtain N infrared small target images including multiple target categories, label the small target objects in each infrared small target image, and then form a training sample R by combining a total of K small targets and their corresponding labels in more than half of each target category. Form a test sample set E by combining the remaining N - K infrared small target images and their corresponding labels, where N ≥ 500. (2) Construct a small object detection network model O based on a graph attention network: Construct a small object detection network model O including an encoder, a pixel-to-graph attention module PGA, and a decoder in a sequential cascade as the backbone network, and a gated cross-attention feature aggregation module GCFA including M cascaded layers and a head block connected to each GCFA as the branch network; wherein the encoder includes M cascaded residual blocks and a max pooling layer loaded between the first and second residual blocks; the decoder includes M cascaded upsampling deconvolution modules and a head block; the branch network is loaded between the input end of the encoder and the input end of the head block in the backbone network; the first residual block, the max pooling layer, the m-th residual block, and the (M - 1)-th residual block in the encoder are respectively connected to the output ends of the M-th, (M - 1)-th, (M - m)-th, and the first upsampling deconvolution module in the decoder; M GCFAs are respectively connected to the output ends of the first to (M - 1)-th upsampling deconvolution modules in the PGA and the decoder; (3) Initialize the parameters: Initialize the number of iterations as t, the maximum number of iterations as T, where T ≥ 1000. The small target detection network model O in the t-th iteration t has weight and bias parameters w t and b t respectively. Let t = 0 and O t = O; (4) Train the small object detection network model: Randomly select P training samples from the training sample set R with replacement as the input of the small object detection network model O for forward propagation to obtain P small object detection results, where 1 ≤ P ≤ K; (5) Obtain the trained small object detection network: For the P small target detection results obtained through step (4), the weight and bias parameters w t of the small target detection network model O t , b t are updated, and it is judged whether t≥T holds. If so, the trained small target detection network model O * is obtained; otherwise, let t=t + 1 and execute step (4); (6) Obtain the small object detection results: Use the test sample set E as the input of the trained small object detection network model O * for forward propagation to obtain the small object detection results corresponding to N-K test samples.

2. The method according to claim 1, wherein The small object detection network model described in step (2), where: The pixel-to-graph attention module PGA includes a graph attention module and a spatial attention module arranged in parallel with their output ends connected; the graph attention module uses a multi-layer graph convolutional neural network GCN; the spatial attention module includes a max pooling layer, an average pooling layer, and a convolutional layer stacked in sequence; The gated cross-attention feature aggregation module GCFA includes two branch structures arranged in parallel, and a cross-attention module connected to the output ends of the two branches; each branch structure contains multiple stacked 1×1 convolutional layers with a normalization layer and a GeLU activation function and a depth convolutional layer.

3. The method according to claim 2, characterized in that, The training of the small object detection network model described in step (4) is achieved as follows: (4a) The encoder downsamples each training sample to obtain P feature maps F containing high-level semantic information p ; (4b) In the graph attention module and spatial attention module of the pixel-to-graph attention module PGA, the graph attention module and spatial attention module respectively perform graph feature extraction and spatial feature extraction on the feature map F p to obtain the graph feature map of F p and the spatial feature map The attention module performs feature fusion on and to obtain the refined semantic feature F of F p ', where: p ​ Wherein, GCN is a graph convolution operation, MaxPool is a max pooling operation, AveragePool is an average pooling operation, Conv is a convolution operation, and σ is a sigmoid activation function; (4c) The M upsampling transposed convolution modules in the decoder upsample each refined semantic feature F' p and perform upsampling on F', p and respectively fuse the M local feature maps with different depths obtained by upsampling F' p with the M local feature maps with different depths of the output F p of the encoder to obtain the global feature map F” p ; (4d) The branch network obtains the contour features of each input training sample (4e) The head block in the decoder pairs F' p with the global feature map F" p and the contour features to perform fusion, obtaining the fused feature map F''' p .

4. The method according to claim 3, characterized in that The branch network described in step (4d) obtains contour features The implementation steps are as follows: The M GCFAs in the two branch structures of the branch network preprocess and perform depth convolution on their two-way inputs to obtain multiple attention maps respectively output by the two branch structures; the cross-attention module performs cross-attention fusion on the two-way multiple attention maps; M refined contour features are obtained from M GCFA The head block fuses the M refined contour features to obtain a contour feature 5. The method according to claim 3, wherein The weight and bias parameters w t of the small target detection network model O t , b t are updated, and the implementation steps are as follows: (5a) Calculate the loss value L of the small target detection network model O t : Adopt the Dice loss function \(L\) Dice , and calculate the contour loss value output by the branch structure using the predicted values and actual values of the output results of multiple GCFA modules obtained through the head block in the branch network Meanwhile, adopt the BCE loss function \(L\) BCE , and calculate the contour loss value output by the branch structure using the predicted values, actual values of the output results of multiple GCFA modules obtained through the head block in the branch network, and the weight coefficient Then and are added together to obtain the contour loss value \(L\) con ; Then calculate the Dice loss of the backbone network using the predicted values and actual values of the global feature map obtained through the decoder head block Use the BCE loss obtained from the predicted values and actual values of the output results of the PGA module The Dice loss of the backbone network The BCE loss of the output results of the PGA module and the contour loss \(L\) con are combined to obtain the total loss \(L\) of the final training objective t ; Among them, X1, X2, and X3 respectively represent the predicted values of the output results of the branch network, the output results of the backbone network, and the output results of the PGA module. Y1, Y2, and Y3 respectively represent the actual values of the output results of the branch network, the output results of the backbone network, and the output results of the PGA module. ∩ represents the intersection operation, ln(·) represents the natural logarithm function, ω represents the weight coefficient, and λ is a hyperparameter that balances the two losses; (5b) Calculate L by the chain rule t The partial derivatives of t the weight parameter w t and the bias parameter b and and update w and b t according to t as follows: Among them, w t ', b t ' represent the weights and bias parameters w t and b t of O t 's updated results, and α represents the learning rate.

Images