An infrared small target detection method based on YOLOv4 multi-scale feature fusion

Abstract

The application discloses an infrared small target detection method based on YOLOv4 multi-scale feature fusion, which improves the problems of poor robustness, feature coupling and easy feature loss of traditional infrared small target detection methods in different backgrounds. The application contains, 1, obtaining an infrared small target detection image for pretreatment; 2, extracting the global features of the image, embedding the features after dilated convolution of the target region into the global branch; 3, using a sliding window to sample and down-sample the original image to reduce the resolution of the original image, obtaining a generated adversarial network model trained by the input of the adversarial network model; 4, fusing the local enhanced features output by the backbone network, the obtained global features and the context features by learning different weights; 5, inputting the fused module into a spatial self-attention mechanism; 6, calculating a loss function to obtain a final detection result. The technology can adapt to the demand of infrared small target detection in different environments.

Description

TECHNICAL FIELD

[0001] The present application relates to the field of image processing technology and pattern recognition technology, and particularly relates to an infrared small target detection method based on YOLOv4 multi-scale feature fusion. BACKGROUND

[0002] The infrared imaging capability can penetrate obstacles and provide relatively clear images without illumination, so the infrared small target detection plays an irreplaceable role in the fields of early warning systems, remote sensing and target tracking. However, on the one hand, small targets lack appearance information, shape and texture features that can be distinguished from them, and locking the position of small objects will produce more uncertainty. On the other hand, under the influence of various factors such as noise, humidity and distance in a complex background, the current method still has a high false alarm rate and even misses detection. Therefore, there are still great difficulties and challenges in the current infrared small target detection.

[0003] The traditional method assumes that the background is static or the target image in adjacent frames is consistent, and detects by using the difference between the background and the infrared small target. This is easy to cause feature loss due to the lack of obvious contrast, leading to missed detection and false positive detection results. At the same time, the dynamic changes of the background image caused by external conditions affect the final detection results, so the traditional method does not have good robustness to the diversity of changes.

[0004] In recent years, with the rapid development of convolutional neural networks in deep learning, some infrared small target detection methods based on convolutional neural networks have made great progress. Convolutional neural networks can learn the deep semantic features of infrared small target images in an end-to-end manner, so the target detection network is applied to many fields. Related work includes: Yang Qili et al. proposed an infrared weak small target detection algorithm based on a fully convolutional network, constructed a fully convolutional network based on a recursive convolutional layer, and enhanced the contrast between the target and the background in a complex environment (see reference: Yang Qili, Zhou Binghong, Zheng Wei, et al. Infrared weak small target detection algorithm based on fully convolutional network [J]. Infrared technology, 2021, 43(4): 349-356); Cai Yunze designed a double-channel feature extraction module to enrich target features by feature enhancement to solve the problems of high false alarm rate and missed detection rate in long-distance infrared target detection (see reference: Cai Yunze, Zhang Yanjun. Infrared weak small target detection method based on double-channel feature enhancement integrated attention network [J]. Space and air defense, 2021, 4(4): 14-22); Wang Longxin et al. used conditional adversarial generative network and YOLOv3-tiny detector to construct an infrared small target detection and tracking system, which optimized the problems of complex infrared image background, low signal-to-clutter ratio and small target size (see reference: Wang Longxin, Zeng Dan, Zhu Xiaojun. Airborne infrared small target detection and tracking system [J]. Industrial control computer, 2022, 35(5): 39-41). SUMMARY

[0005] The present application aims at the problems of poor robustness in different backgrounds, feature coupling and feature loss in traditional infrared small target detection methods, and provides an infrared small target detection method based on YOLOv4 multi-scale feature fusion with high accuracy of detection results and good robustness in different scenes.

[0006] The technical solution of the present application is to provide an infrared small target detection method based on YOLOv4 multi-scale feature fusion with the following steps:

[0007] Step 1, acquire an infrared small target detection image, and pre-process the image to adjust it to a uniform size;

[0008] Step 2, acquire global feature information of the input image, perform dilated convolution on the target region, embed the features obtained by dilated convolution into the global branch, and learn the context information of the target region;

[0009] Step 3, sample the infrared small target detection image using a sliding window, downsample the sampled image block sequence to reduce the resolution of the original image, and obtain a low-resolution image for training a generative adversarial network model, input the image block sequence sampled by the sliding window into the trained generative adversarial network model, and acquire an image with enhanced local features;

[0010] Step 4, input the high-resolution image obtained in step 3 into the backbone network of YOLOv4, and fuse the locally enhanced features output by the backbone network, the global features obtained in step 2, and the context features through learning different weights;

[0011] Step 5, input the fused features into the spatial self-attention mechanism to enhance the dependent response of the infrared small target in the spatial structure;

[0012] Step 6, calculate the loss function, regress the position of the infrared small target through the head network of YOLOv4, and obtain the final detection result.

[0013] Preferably, in step 1, the input image to be detected is resized to the same size using the resize() function in Opencv.

[0014] Preferably, in step 2, obtaining the global feature information of the input image includes the following steps:

[0015] Step 2.1, input the image to be detected into a convolutional neural network for feature extraction to abstract the global feature representation of the target object;

[0016] Step 2.2, introduce dilated convolution to learn context information and enhance the representation ability of global feature extraction;

[0017] Step 2.3, embed the context information into the global branch to obtain more complete feature information and prior representation.

[0018] Preferably, in step 3, constructing a generative adversarial network model to obtain a high-resolution image with enhanced local features includes the following steps:

[0019] Step 3.1, set the size of the sliding window to 16x16 and the sliding step to 12;

[0020] Step 3.2, downsample the sliding window image sequence using a bicubic interpolation method to obtain a low-resolution image;

[0021] Step 3.3, input the obtained image with the horizontal and vertical directions scaled to half the original size into the generator, and the discriminator is a normal image sequence. Introduce random noise, calculate the loss function through the adversarial learning of the generator and the discriminator, and train the generative adversarial network model;

[0022] Step 3.4, input the image block sequence sampled by the sliding window into the trained generative adversarial network model to obtain a high-resolution image with enhanced local features.

[0023] Preferably, the step 4 of fusing the global features and the local enhanced features comprises the following steps: inputting the obtained high-resolution image into the backbone network of YOLOv4, outputting different scale features through the CSPDarknet53, setting different weights for the output different scale features and the global features obtained in step 2, and fusing and adding the different scale features after convolution.

[0024] Preferably, the step 5 of inputting the fused features into the spatial self-attention mechanism comprises the following steps: step 5.1, setting the output feature map of the feature fusion module as [C×H×W], wherein H, W and C respectively represent the height, width and channel of the feature map;

[0025] Step 5.2, inputting the same into three 1×1 convolutions to generate three feature maps F1, F2 and F3;

[0026] Step 5.3, measuring the influence between different positions of the feature map by reconstructing F1, F2 and F3, first transposing F1 and performing matrix multiplication on F2; then obtaining F4 through the Softmax normalization function;

[0027] Step 5.4, finally, performing matrix multiplication operation on F3 and F4 to obtain the final output [C×H×W].

[0028] Preferably, in the step 6, the total loss function is defined as:

[0029] Loss=L GAN +L Class +L Box

[0030] Wherein, L GAN , L Class , L Box are the loss functions of the generative adversarial network model, the classification loss function and the bounding box regression loss function respectively,

[0031] The classification loss function is specifically defined as follows:

[0032]

[0033] Wherein, represents the dot product result of the feature vector and the weight vector corresponding to the jth class of the ith sample,

[0034] and represent the predicted probability and the real probability of the object belonging to the class c in the jth bounding box of the ith grid respectively;

[0035] The bounding box regression loss function is specifically defined as follows:

[0036]

[0037] wherein IOU represents the intersection of the union of the predicted bounding box and the ground truth bounding box, w gt , h gt are the ground truth width and height of the bounding box respectively, w and h are the predicted bounding box width and height respectively, p 2 (b, b gt ) represents the Euclidean distance between the center points of the predicted bounding box and the ground truth bounding box, c represents the minimum diagonal distance between the predicted bounding box and the ground truth bounding box, according to the true class of the target object and the model prediction class, the sample data can be divided into four categories: true positive (TP), false positive (FP), true negative (TN) and false negative (FN), and finally the detection accuracy of the generative adversarial network is judged by the recall rate, the precision rate and the F1, and the specific formula is as follows:

[0038]

[0039] wherein F represents the proportion of the predicted target being the real target, P represents the proportion of the predicted target being the real target, R represents the proportion of being successfully predicted, and F1 represents the accuracy of the infrared small target detection network.

[0040] Compared with the prior art, the infrared small target detection method based on the YOLOv4 multi-scale feature fusion has the following advantages:

[0041] 1. The existing feature fusion technology mainly applies only one convolutional neural network or obtains multiple features through different processing methods in the same step, then simply fuses the information of the multiple features, and does not carefully consider the coupling in the feature fusion process and the influence degree of different features on the final detection structure, the present application fully considers the global feature, the context information and the enhanced local feature, extracts different scale features, and then learns different weights for different features according to the influence of different features on the detection result, which not only avoids the influence between features in the same step, but also can effectively improve the accuracy of the detection result by adjusting the weight.

[0042] 2. Firstly, the extraction of global features, the enhancement of local features and the context features are designed as a pre-processing process, and multi-scale features are extracted through different network structures and methods, which reduces the feature coupling and feature missing problem in the same step processing process; then, by assigning different learning weights to the local enhanced features, the global features and the context information, the accuracy of the final detection result is accurately improved, and the robustness of the method in different environments is improved; finally, the spatial self-attention mechanism is introduced to obtain the relative dependence of small targets in space, and the final detection result is output.

[0043] Compared with the existing infrared small target detection technology, it is more comprehensive and reliable, and can adapt to the needs of infrared small target detection in different environments. BRIEF DESCRIPTION OF DRAWINGS

[0044] Figure 1 is a workflow schematic diagram of the present application;

[0045] Figure 2 is a module framework diagram of the present application;

[0046] Figure 3 is a partial feature enhancement module schematic diagram of the present application;

[0047] Figure 4 is a feature fusion module schematic diagram of the present application;

[0048] Figure 5 is a spatial self-attention mechanism module schematic diagram of the present application. DETAILED DESCRIPTION

[0049] In order for those skilled in the art to better understand the present application, the technical solutions in the embodiments of the present application will be described clearly and completely below in conjunction with the drawings in the embodiments of the present application. Obviously, the described embodiments are only a part of the embodiments of the present application, not all. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without creative labor should be within the scope of protection of the present application.

[0050] The infrared small target detection method based on YOLOv4 multi-scale feature fusion of the present application will be further described below in conjunction with the drawings and specific embodiments: as shown in the figure, the present embodiment is divided into image initialization module, model initialization module, feature enhancement module, global scene learning module, feature fusion module, spatial self-attention mechanism module, and classification module. The first four modules mainly provide the prerequisite for the feature fusion module, and the last two modules regress and classify the infrared small target position by enhancing the position response design multi-task loss function. Specifically, the steps include: Figure 1

[0051] S1. Obtain the infrared small target detection image, and pre-process the image to adjust to a uniform size.

[0052] In S1, the original image is pre-processed, specifically using the resize() function in Opencv, which can output images of the same size and reduce the impact of image size on network performance and structure.

[0053] cv.resize(src,dsize,dst=None,f x =None,f y ​=None,interpolation=None)

[0054] wherein src represents the original input image, dsize represents the size of the output image, dst represents the output image, f x represents the scaling factor along the horizontal axis, f y represents the scaling factor along the vertical axis, and interpolation represents the interpolation method.

[0055] S2. Obtain the global feature information of the input image, further perform dilated convolution on the target region, embed the features obtained by the dilated convolution into the global branch, and thus learn the context information of the target region.

[0056] S2 mainly includes the following steps:

[0057] S21. Input the image to be detected into a 1x1 convolutional neural network for feature extraction, and abstract the global feature representation of the target object.

[0058] S22. In order to learn the context information, a 3x3 dilated convolution is introduced, and the input image is subjected to convolution operation with different dilation rates, and the hole rates are set to 2, 4 and 8 respectively, effectively expanding the receptive field of the image.

[0059] S23. Embed the features learned by the scene into the global branch, input into a 3x3 convolutional layer, and after the ELU activation function, the gradient vanishing problem is alleviated, and the BatchNormalization (batch normalization) is used to normalize the features, and more complete feature information and prior representation are extracted.

[0060]

[0061] wherein x>0 part alleviates the gradient vanishing problem, and x≤0 part can be more robust to input changes or noise.

[0062] S3. Use a sliding window to sample the infrared small target detection image, downsample the sampled image block sequence to reduce the resolution of the original image, and obtain a low-resolution image for training a generative adversarial network model. The image block sequence sampled by the sliding window is input into the trained generative adversarial network model to obtain an image with enhanced local features.

[0063] S3 mainly includes the following steps:

[0064] S31. The size of the sliding window is set to 16x16, and the sliding step number is set to 12.

[0065] S32. The sliding window image sequence is down-sampled using the bicubic interpolation method to obtain a low-resolution image, which not only increases the smoothing processing of the pixels, but also provides more accurate information by referencing the surrounding sixteen pixels. The size of the training sample image A is MxN, and after scaling K times, the image B has a size of mxn, K=M / m. The BiCubic(double cubic interpolation) function is used to calculate the weight W(x) of the pixel point:

[0066]

[0067] where x represents the distance of the adjacent pixel point to the image B, and a is-0.5.

[0068] Since the position of the image B is related to the weight of the above-mentioned pixel point, the value of the image B at point (i, j) can be calculated:

[0069]

[0070] S33. The generated generative adversarial network model is composed of a generator composed of two up-sampling convolution layers and a discriminator composed of three convolution layers. In order to be closer to the real situation, random noise is added to the input picture. The image obtained in S32 is input into the generator, and the discriminator is a normal image sequence. Through the adversarial learning of the generator and the discriminator, the loss function of the generative adversarial network model is calculated:

[0071]

[0072] where G represents the generator, D represents the discriminator, E(*) represents the expected value of the distribution function, P data (x) represents the distribution of the real sample, and P noise (z) is the noise distribution defined in low dimension. The structure of the generative adversarial network is shown in Figure 3 .

[0073] S34. The image block sequence sampled by the sliding window is input into the trained generative adversarial network model to obtain a local feature enhanced image block sequence.

[0074] S4. The high-resolution image obtained in S34 is input into the backbone network of YOLOv4, including CSPDarknet53, Mish activation function, Dropblock. The local enhanced features output by the backbone network, the global features obtained in S2, and the context features are fused with each other through learning different weights.

[0075] In the process of S4 feature fusion, the to-be-fused features are divided into three scales of features of two channels, one part is the global features through S2 and the context features, and the other part is the local enhanced features obtained by inputting the enhanced pictures of S3 into the backbone network of YOLOv4. The features are divided into two channels for separate processing, which not only can reduce the coupling effect between different scale features, but also can fully exert the characteristics of different networks to a great extent to obtain single-scale features with more complete information.

[0076] In S4, the features of different scales output by CSPDarknet53 are input into a 3x3 convolution module together with the global features obtained by S2, and different learnable weights W1, W2 and W3 are set according to the contribution of the receptive field of different features to the output results to construct a feature fusion module. By designing different weights, the influence degree of different scale features on the final detection results can be measured, and the network structure of the feature fusion module is as shown in Figure 4 .

[0077] S5. The fused features are input into a spatial self-attention mechanism to enhance the dependent response of the infrared small target in the spatial structure, and the network structure of the spatial self-attention mechanism module is as shown in Figure 5 .

[0078] Due to the large scene span and long imaging distance of the infrared small target, the infrared small target often occupies only a few pixels, and there is great uncertainty in the process of small target position regression and classification. In the spatial self-attention mechanism, the influence of the position relationship in different spaces can be considered by reconstructing the feature maps of different scales to enhance the regression position condition of the infrared small target.

[0079] S5 mainly includes the following steps:

[0080] S51. The output feature map of the feature fusion module is set to [CxHxW], wherein H, W and C represent the height, width and channel of the feature map respectively.

[0081] S52. It is input into three 1x1 convolutions to generate three feature maps F1, F2 and F3, wherein F1, F2, F3∈[CxHxW].

[0082] S53. The influence between different positions of the feature map is measured by reconstructing F1, F2 and F3. First, the transpose of F1∈(HxW)xC and F2∈Cx(HxW) are multiplied; then F4∈(HxW)x(HxW) is obtained by the Softmax normalization function.

[0083] S54. Finally, F3∈[CxHxW] and F4 are multiplied to obtain the final output [CxHxW].

[0084] S6. Calculate the loss function, regress the position of the infrared small target through the head network of YOLOv4, and obtain the final detection result.

[0085] The main task loss functions in the design of the entire network are three parts, namely: the loss function of the generative adversarial network in S3, the classification loss function and the bounding box regression loss function in S6. In S6, a total loss function based on the infrared small target detection method of YOLOv4 multi-scale feature fusion is defined: Loss = L GAN +L Class +L Box

[0086] Among them, L GAN , L Class , L Box are the loss function of the generative adversarial network model, the classification loss function and the bounding box regression loss function, respectively.

[0087] The classification loss function is specifically defined as follows:

[0088]

[0089] Among them, represents the dot product result of the feature vector of the i-th sample and the weight vector corresponding to the j-th class,

[0090] and represent the predicted probability and the true probability of the object belonging to the c-th class in the i-th grid and the j-th bounding box, respectively.

[0091] The bounding box regression loss function is specifically defined as follows:

[0092]

[0093] Among them, IOU represents the intersection of the union of the predicted bounding box and the true value bounding box, w gt , h gt are the true value width and height of the bounding box, w and h are the predicted bounding box width and height, respectively, ρ 2 (b,b gt ) represents the Euclidean distance between the center points of the predicted bounding box and the true value bounding box, and c represents the minimum diagonal distance between the predicted bounding box and the true value bounding box.

[0094] According to the true class and the model predicted class of the target object, the sample data can be divided into four categories: true positive (TP), false positive (FP), true negative (TN) and false negative (FN), and finally the detection accuracy of the generative adversarial network is evaluated through the recall rate, precision rate and F1. The specific formula is as follows:

[0095]

[0096] Wherein F represents the proportion of the real target predicted in the target, P represents the proportion of the real target predicted accurately by the prediction target, R represents the proportion successfully predicted, and F1 represents the accuracy of the infrared small target detection network.

[0097] The above merely describes the preferred embodiments of the present application, and it should be noted that those skilled in the art can make several improvements and refinements without departing from the principles of the present application, and these improvements and refinements should also be considered within the protection scope of the present application.

Claims

1. An infrared small target detection method based on YOLOv4 multi-scale feature fusion, characterized in that: The method comprises the following steps: Step 1, obtaining an infrared small target detection image, pre-processing the image, and adjusting the image to a uniform size; Step 2, obtaining global feature information of the input image, performing dilated convolution on a target region, embedding the features obtained by the dilated convolution into a global branch, and learning context information of the target region; Step 3, using a sliding window to sample the infrared small target detection image, performing down-sampling on the sampled image block sequence to reduce the resolution of the original image, obtaining a low-resolution image for training a generative adversarial network model, inputting the image block sequence sampled by the sliding window into the trained generative adversarial network model, and obtaining an image with enhanced local features; Step 4, inputting the high-resolution image obtained in step 3 into the backbone network of YOLOv4, fusing the local enhanced features output by the backbone network, the global features obtained in step 2 and the context features by learning different weights; Step 5, inputting the fused features into a spatial self-attention mechanism to enhance the dependent response of the infrared small target in the spatial structure; Step 6, calculating a loss function, regressing the position of the infrared small target through the head network of YOLOv4, and obtaining a final detection result; the overall loss function is defined as: Loss = L GAN + L Class + L Box wherein, L GAN , L Class , L Box are a generative adversarial network model loss function, a classification loss function and a bounding box regression loss function, respectively, The classification loss function is specifically defined as follows: wherein, denotes the dot product of the feature vector of the i-th sample and the weight vector corresponding to the j-th class, and respectively represent the predicted probability and the true probability of the object belonging to class c in the i-th grid and the j-th bounding box. The bounding box regression loss function is specifically defined as follows: where IOU denotes the intersection of the union of the predicted bounding box and the ground truth bounding box, w gt , h gt are the ground truth width and height of the bounding box, respectively, w and h are the predicted bounding box width and height, respectively, p 2 (b, b gt ) denotes the Euclidean distance between the centers of the predicted bounding box and the ground truth bounding box, and c denotes the minimum diagonal distance between the predicted bounding box and the ground truth bounding box.

2. The YOLOv4 multi-scale feature fusion-based infrared small target detection method according to claim 1, characterized in that: In step 1, the input image to be detected is adjusted to the same size by using the resize() function in Opencv.

3. The YOLOv4 multi-scale feature fusion-based infrared small target detection method according to claim 1, characterized in that: In step 2, the global feature information of the input image includes the following steps: Step 2.1, inputting the image to be detected into a convolutional neural network for feature extraction, and abstracting global feature representation of the target object; Step 2.2, introducing dilated convolution to learn context information and enhance the representation ability of global feature extraction; Step 2.3, embedding the context information into the global branch to obtain more complete feature information and prior representation.

4. The YOLOv4 multi-scale feature fusion-based infrared small target detection method according to claim 1, characterized in that: In step 3, the generative adversarial network model is constructed to obtain high-resolution images with enhanced local features, including the following steps: Step 3.1, setting the size of the sliding window to 16x16 and the sliding step to 12; Step 3.2, using a bicubic interpolation method to down-sample the sliding window image sequence to obtain a low-resolution image; Step 3.3, inputting the obtained image with the horizontal direction and the vertical direction scaled to half of the original size into the generator, inputting the normal image sequence into the discriminator, introducing random noise, calculating the loss function through the adversarial learning of the generator and the discriminator, and training the generative adversarial network model; Step 3.4, inputting the image block sequence sampled by the sliding window into the trained generative adversarial network model to obtain high-resolution images with enhanced local features.

5. The YOLOv4 multi-scale feature fusion-based infrared small target detection method according to claim 1, characterized in that: In step 4, the global feature and the local enhanced feature are fused, including the following steps: inputting the obtained high-resolution image into the backbone network of YOLOv4, setting different weights for the different scale features output by the backbone network CSPDarknet53 and the global features obtained in step 2, and fusing and adding the different scale features after convolution.

6. The YOLOv4 multi-scale feature fusion-based infrared small target detection method according to claim 1, characterized in that: The step 5 inputs the fused feature into a space self-attention mechanism, including the following steps: step 5.1, setting the output feature map of the feature fusion module as [C×H×W], H, W and C respectively representing the height, width and channel of the feature mapping; Step 5.2, inputting it into three 1×1 convolution to generate three feature mappings F1, F2 and F3; Step 5.3, measuring the influence between different positions of the feature map by reconstructing F1, F2 and F3, first performing matrix multiplication on the transpose of F1 and F2; then obtaining F4 through the Softmax normalization function operation; Step 5.4, finally, performing matrix multiplication operation on F3 and F4 to obtain the final output [C×H×W].

Images