Infrared dim small target detection method and device based on YOLOv8

Abstract

The application discloses a weak and small target detection method and device based on Yolov8, which comprises the following steps: pre-processing a to-be-detected image to obtain a feature image; inputting the feature image into an optimized target detection model; and detecting the weak and small target by using the target detection model; wherein the optimization of the target detection model comprises the following steps: fusing a perception large kernel convolution in a spatial pyramid pooling layer of a feature extraction network to obtain multi-scale sparse features of a target; introducing a triple attention mechanism layer at an output end of a feature fusion network to adaptively focus feature information integrated with the sparse features; and constructing a multi-scale target detection network to detect the adaptively focused feature information. The application captures sparse features of a target by combining a large kernel convolution layer in a spatial pyramid, and realizes cross-latitude interaction by using a triple attention mechanism added at an output end of a feature fusion network, so that the detection accuracy of the model for infrared weak and small targets is improved as a whole.

Description

TECHNICAL FIELD

[0001] The present disclosure relates to the field of target detection, and in particular, to an infrared small target detection method and device based on YOLOv8. BACKGROUND

[0002] Infrared small target recognition technology is an important part of an infrared search and track (IRST) system and can also be flexibly deployed to other platforms, and has extremely important engineering application value in the fields of infrared early warning, medium and long-range precision strike, space attack and defense confrontation, etc.

[0003] Due to the limitations of the detection instrument itself and the fact that the target is far away from the detector, the infrared image has low contrast, the infrared target in the image has low brightness, the texture shape feature is weak, the occupied pixel points are few, and the background is complex, and when a YOLO (You Only Look Once) model is used for recognition, the recognition of the infrared small target lacks pertinence and has low recognition accuracy.

[0004] Therefore, there is a need for a detection method that can more accurately recognize infrared small targets. SUMMARY

[0005] The purpose of the present disclosure is to provide an infrared small target detection method and device based on YOLOv8, which can solve at least one of the technical problems mentioned above. The specific scheme is as follows:

[0006] According to the specific embodiment of the present disclosure, in a first aspect, an infrared small target detection method based on YOLOv8 is provided, comprising: pre-processing a to-be-detected image to obtain a feature image;

[0007] inputting the feature image into an optimized target detection model, and performing small target detection using the optimized target detection model;

[0008] The optimization of the target detection model comprises:

[0009] fusing a perception large kernel convolution in a spatial pyramid pooling layer of a feature extraction network of the target detection model to obtain multi-scale sparse features of a target;

[0010] introducing a triple attention mechanism layer at an output end of a feature fusion network of the target detection model, for adaptively focusing feature information integrated with the sparse features;

[0011] constructing a multi-scale target detection network to detect the adaptively focused feature information.

[0012] Preferably, the pre-processing of the image to be detected obtains a feature image, comprising:

[0013] Adaptive size adjustment is performed on the input image to be detected;

[0014] The pixel value of the image to be detected is normalized to [0, 1], and data enhancement processing is performed;

[0015] The image to be detected is divided into N×N grid units to obtain a plurality of feature images of different sizes.

[0016] Preferably, the perception large kernel convolution comprises a plurality of parallel layers, and (k-1) r+1≤K;

[0017] Wherein, k is the convolution kernel of the parallel layer, r is the inflation rate of the parallel layer, and K is the large kernel.

[0018] Preferably, the large kernel K=13, the number of parallel layers n=5, the convolution kernel of the parallel layer k=(7, 5, 3, 3, 3), and the inflation rate r=(2, 1, 3, 4, 5).

[0019] Preferably, the triple attention mechanism layer comprises:

[0020] The first branch is used to capture the interaction features between the channel C and the spatial dimension W;

[0021] The second branch is used to capture the interaction features between the channel C and the spatial dimension H;

[0022] The third branch is used to construct spatial attention and capture the dependency relationship between the spatial dimensions W and H;

[0023] The outputs of the first branch, the second branch and the third branch are aggregated by averaging to form an attention weighted tensor, and the generated attention weighted tensor is applied to the original tensor of the input image to be detected.

[0024] Preferably, the attention weighted tensor is:

[0025]

[0026] Wherein, σ represents the S-shaped activation function; ψ1, ψ2 and ψ3 represent the standard two-dimensional convolution layer;

[0027] Indicates: the rotated tensor in the first branch;

[0028] Indicates: the rotated tensor in the second branch;

[0029] Indicates: the rotated tensor in the third branch, which is consistent with the original tensor;

[0030] represents: the tensor after the shape is restored by the z-pooling layer in the first branch;

[0031] represents: the tensor after the shape is restored by the z-pooling layer in the second branch.

[0032] Preferably, the optimization of the target detection model further comprises:

[0033] Optimization of the feature extraction network: the feature image is down-sampled using two initial convolutional layers to extract primary pixel features; and then three times of down-sampling are continuously performed through a C2f layer combined with a convolutional layer to obtain three first feature maps of different scales.

[0034] Optimization of the feature fusion network: the multi-scale sparse features obtained after the spatial pyramid pooling layer is pooled are continuously up-sampled twice to obtain two second feature maps of different scales.

[0035] The first feature map and the second feature map of the corresponding scale are fused respectively, and the feature information is further extracted by a C2f layer and then output to the triple attention mechanism layer.

[0036] Preferably, the optimization of the target detection model further comprises: a 4-fold up-sampling layer is added after the two times of up-sampling as a small target detection layer to obtain a third feature map.

[0037] The third feature map and the first feature map are fused, and the fused feature information is output to the target detection layer through a C2f layer.

[0038] Preferably, the detection of the feature information after the adaptive focusing comprises:

[0039] Generating a prediction box: for the feature image of any size containing a target center point in an N×N grid unit, a plurality of bounding boxes and class probabilities are predicted for each grid cell in each feature image of different sizes.

[0040] Prediction box decoding: the offset of the bounding box is predicted by using an Anchor-Free mechanism, and the decoded prediction box is filtered by a confidence threshold and a non-maximum suppression algorithm to remove low-confidence predictions and overlapping prediction boxes.

[0041] Outputting a detection result: the prediction box after post-processing is sorted according to the confidence, and the class, position and confidence of the target in the image to be detected are output.

[0042] According to the specific embodiment disclosed in the present application, the second aspect of the present application discloses a YOLOv8-based infrared dim target detection device, comprising:

[0043] A data acquisition module is configured to pre-process a to-be-detected image to obtain a feature image.

[0044] A detection module is configured to detect the feature image by using a target detection model of the above method.

[0045] Compared with the prior art, the above scheme of the disclosed embodiment has at least the following beneficial effects:

[0046] On the basis of the existing YOLOv8 network model, the sparse features of the target are captured by combining a large-kernel convolution layer in the spatial pyramid pooling; the cross-latitude interaction is realized by using the triple attention mechanism added at the output end of the feature fusion network, and the useful features in the region are dynamically focused. Thus, the detection accuracy of the model for the infrared dim target is improved as a whole. BRIEF DESCRIPTION OF DRAWINGS

[0047] The accompanying drawings, which are incorporated into the specification and constitute a part of the specification, illustrate embodiments consistent with the present application and, together with the specification, serve to explain the principles of the present application. Obviously, the accompanying drawings in the following description are only some embodiments of the present disclosure, and other drawings can be obtained by those skilled in the art without creative labor. In the drawings:

[0048] Figure 1 A flowchart of a dim target detection method based on YOLOv8 of the present application;

[0049] Figure 2 A framework diagram of the target detection model of the present application;

[0050] Figure 3 A schematic diagram of the triple attention mechanism module of the present application;

[0051] Figure 4 A schematic diagram of the dim target detection layer and the dim target detection head of the present application;

[0052] Figure 5 A schematic diagram of the large-kernel convolution UniRepLK Block module of the present application;

[0053] Figure 6 A comparison diagram of the detection results of the present application. DETAILED DESCRIPTION

[0054] In order to make the purposes, technical solutions and advantages of the present disclosure clearer, the following will further describe the present disclosure in detail with reference to the drawings. Obviously, the described embodiments are only some but not all of the embodiments of the present disclosure. Based on the embodiments of the present disclosure, all other embodiments obtained by a person of ordinary skill in the art without creative work fall within the scope of protection of the present disclosure.

[0055] The terms used in the embodiments of the present disclosure are only for the purpose of describing particular embodiments and are not intended to limit the present disclosure. The singular forms "a", "an" and "the" used in the embodiments of the present disclosure and the appended claims are also intended to include the plural forms, unless the context clearly indicates otherwise. "Plural" generally includes at least two.

[0056] It should be understood that although the terms first, second, third, etc. may be used in the embodiments of the present disclosure to describe, these descriptions should not be limited to these terms. These terms are only used to distinguish the description. For example, without departing from the scope of the present disclosure, first can also be referred to as second, and similarly, second can also be referred to as first.

[0057] The optional embodiments of the present disclosure will be described in detail below with reference to the drawings.

[0058] Embodiment 1

[0059] As shown in the figure, the infrared small target detection method based on YOLOv8 of the present application comprises the following steps: Figure 1

[0060] Step S1. Preprocessing the image to be detected to obtain a feature image.

[0061] Step S2. Input the feature image into the optimized target detection model, and use the optimized target detection model to detect small targets.

[0062] The optimization step of the target detection model comprises:

[0063] Fusing a perception large kernel convolution in the spatial pyramid pooling layer of the target detection model to obtain multi-scale sparse features of the target;

[0064] Introducing a triple attention mechanism layer at the output end of the feature fusion network of the target detection model, for adaptively focusing the feature information integrated with the sparse features;

[0065] Constructing a multi-scale target detection network, fusing the adaptively focused feature information and corresponding output.

[0066] ​The application improves the detection precision of the model on the infrared dim small target as a whole by combining a large kernel convolution layer in the spatial pyramid pooling to capture the sparse features of the target, and using a triple attention mechanism added at the output end of the feature fusion network to realize cross-latitude interaction and dynamically focus on useful features in the region.

[0067] Specifically, the step of obtaining a feature image by preprocessing the to-be-detected image in step S1 specifically includes:

[0068] Step S1.1 obtains the to-be-detected image, so that the size of the input to-be-detected image is 640x640.

[0069] The small target in the application can be six types of targets, including kites, high-speed target aircrafts, balloons, civil aviation aircrafts, bird flocks and unmanned aerial vehicles.

[0070] Step S1.2 normalizes the pixel value of the to-be-detected image to [0, 1], and performs mosaic enhancement, mixed enhancement, random disturbance and color disturbance, so as to improve the generalization ability and robustness of the model.

[0071] Step S1.3 divides the to-be-detected image into NxN grid units to obtain a plurality of feature images of different sizes.

[0072] Figure 2 The structure of the target detection model optimized based on the YOLOv8 network model in step S2 is shown, and the optimization method includes the following steps:

[0073] Step S2.1 optimization of the feature extraction network.

[0074] In the transmission direction of the feature extraction network, two initial convolution layers Conv0 and Conv1 are included in sequence, which are used to downsample the feature image by 4 times, extract the primary features of the original pixels, and obtain a sub-feature map of 160x160.

[0075] The channel number is doubled by using the C2f layer, and the sub-feature map is divided into two parts, one part is directly combined with a convolution layer Conv to perform three times of down sampling in succession to obtain a multi-scale first feature map; the other part participates in subsequent feature fusion.

[0076] Specifically, the first feature map can be 20x20, 40x40, 80x80 of three different scales, which are used to predict targets of different sizes.

[0077] Step S2.2 fuses a perception large kernel convolution UniRepLK Block in a spatial pyramid pooling layer SPPF (Spatial Pyramid Pooling Fast), as shown in Figure 5 .

[0078] Specifically, the UniRepLK Block in the embodiment has five parallel layers. For the task target in the embodiment, a large kernel K=13 is selected, and the convolution kernel size and the expansion rate of the five parallel layers are k=(7, 5, 3, 3, 3) and r=(2, 1, 3, 4, 5) respectively.

[0079] The five parallel layers are structurally re-parameterized to obtain a large convolution kernel with an expansion coefficient of 1.

[0080] In other embodiments, different parameters can be defined according to different downstream tasks, but need to satisfy (k-1) r+1≤K;

[0081] wherein k is the convolution kernel of the parallel layer, r is the expansion rate of the parallel layer, and K is the large kernel.

[0082] Since the infrared dim small target usually lacks detailed texture information and occupies few pixels, the traditional convolution form relies on small kernel convolution layers to gradually extract features, which is insufficient to capture the global features of the dim small target when processing infrared images. Although the large kernel convolution is wide but not deep, it has a better shape bias. In the case that the target part is blocked or similar to the background, the shape bias helps the model to identify targets of different scales, so as to more accurately locate and distinguish the boundaries of the target object.

[0083] Therefore, the present application uses a non-expanded convolution layer with a small kernel and multiple expanded convolution layers with small kernels to enhance the non-expanded convolution layer with a large kernel to capture the sparse features of the target.

[0084] Step S2.3. Optimization of the feature fusion network, comprising:

[0085] Step S2.3.1. The sparse features obtained after the spatial pyramid pooling layer is pooled are successively upsampled twice to obtain second feature maps with scales of 40x40 and 80x80 respectively.

[0086] The second feature maps are fused with the first feature maps obtained in step S2.1 according to the corresponding scales at Concat11, Concat14 and Concat26 layers.

[0087] Preferably, it can also include:

[0088] Step S2.3.2. Based on the obtained 80x80 second feature map, a 4-fold upsampling layer is added as a dim small target detection layer, as shown in Figure 4 to obtain a third feature map with a scale of 160x160.

[0089] The third feature map is merged with a 160x160 sub-feature map obtained by 4 times down-sampling in the feature extraction network, and the merged feature map is output after a C2f layer; meanwhile, the obtained 160x160 sub-feature map is continuously down-sampled three times and fused with three second feature maps having different sizes, to further enhance the semantic information of small targets.

[0090] Step S2.4 introduces a triple attention mechanism layer at the output end of the feature fusion network.

[0091] In order to dynamically focus on the key part of the image and strengthen the learning of the region, so as to pay more attention to the useful features of the region, a triple attention mechanism layer is introduced before entering the target detection network, as shown in Figure 3 .

[0092] Suppose the shape of the input tensor is CxHxW, where C is the number of channels, H is the height, and W is the width. Triple attention can capture cross-dimensional interactions in a structured way, rather than dynamically adjusting the weight of each branch during training. The triple attention mechanism consists of three branches, each responsible for capturing different dimensional interaction features. Specifically, each branch is responsible for reviving cross-dimensional interaction features between the spatial dimensions H or W and the channel dimension C of the input:

[0093] The first branch captures the interaction between the channel C and the spatial dimension W;

[0094] The second branch captures the interaction between the channel C and the spatial dimension H;

[0095] The third branch, similar to the structure in the convolution block attention module, is used to build spatial attention and capture the dependency between spatial dimensions H and W.

[0096] Finally, the outputs of the three branches are aggregated by simple averaging to form the final attention weighted tensor, and the generated attention weights are applied to the original tensor, i.e., the feature image before entering the triple attention mechanism, to enhance the features of specific channels and spatial regions.

[0097] In addition, in the first branch and the second branch, the input feature map is rotated by 90 degrees along the H and W axes to build the connection between the channel and the spatial dimension; the result after rotation is maximum-pooled and average-pooled by the Z-pooling layer, which reduces the spatial dimension while preserving important features.

[0098] In this embodiment, the attention weighted tensor is:

[0099]

[0100] where σ is the sigmoid activation function;

[0101] Ψ1, Ψ2 and Ψ3 represent a standard two-dimensional convolution layer;

[0102] represents: the rotated tensor in the first branch

[0103] represents: the rotated tensor in the second branch

[0104] represents: the rotated tensor in the third branch, consistent with the original tensor

[0105] represents: the tensor after the shape is restored by the z pooling layer in the first branch

[0106] represents: the tensor after the shape is restored by the z pooling layer in the second branch

[0107] Step S3. A multi-scale target detection network is constructed, and the target information after adaptive focusing is fused and output.

[0108] Specifically, the 80x80, 40x40 and 20x20 feature images of three different scales output by the triple attention mechanism are respectively sent into three detection heads with sizes of 80x80, 40x40 and 20x20, which are respectively used to predict targets occupying 8, 16 and 32 pixels.

[0109] The 160x160 detection layer output by the weak small target detection layer is sent into the newly added weak small target detection head, so as to avoid the phenomenon of feature loss when facing small targets, and to make it difficult to be recognized by the subsequent detection head.

[0110] In other embodiments of the present application, the target detection model can also be optimized, specifically by adjusting the hyperparameters, model fine-tuning, regularization techniques, and environment hardware optimization detection model.

[0111] In other embodiments of the present application, the step of detecting the feature information after adaptive focusing in the target detection network comprises:

[0112] Step S4.1 generates a prediction box; the image to be detected after the preprocessing in step S1 is divided into a plurality of NxN grids, and the image containing any target center in the grid is predicted by using different size feature maps inside the model to predict different size targets, and each cell of each feature map will predict a plurality of bounding boxes and class probabilities.

[0113] Step S4.2 prediction box decoding: the offset of the bounding box is predicted by using the Anchor-Free mechanism, and the decoded prediction box is filtered by the confidence threshold and the non-maximum suppression algorithm to remove low confidence predictions and overlapping prediction boxes.

[0114] Step S4.3 outputs the detection result: the post-processed prediction box will be sorted according to the confidence, and the class, position information and confidence of the target in the image to be detected are output.

[0115] Figure 6 A comparison chart of infrared monitoring of ground-to-air under unknown air situation by using the target detection model of the application and other models is shown.

[0116] Figure 6 The first column is the original image, and the detected targets from top to bottom are civil aviation aircraft, unmanned aerial vehicle, air ball, high-speed target aircraft and bird group. The second to fifth columns are the detection images of the corresponding targets in the original image using the Resnet50 model, Vgg16 model, Yolov7_tiny model, Yolov8n model and the target detection model of the application, respectively.

[0117] Table 1 is a comparison experiment, and the accuracy of the target detection model of the application is 95.9, which is higher than that of other detection models.

[0118] Table 1 Comparison experiment

[0119]

[0120] In summary, in the construction of the optimized YOLOv8 target detection model, different size pooling operations are performed on the feature map, and a bidirectional fusion strategy is used to integrate the results, so that the model can capture low-level detail information and high-level semantic information at the same time, thereby obtaining a fixed size output. After being adaptively focused by the triple attention mechanism, it is sent to the corresponding detection head for classification and judgment. This method can effectively process targets with various sizes, and the average precision index of infrared weak target recognition is improved to 95.9%, improving the generalization ability and robustness of the neural network.

[0121] Embodiment 2

[0122] The application further discloses a device embodiment for implementing the method of the above embodiment, which is used for implementing the method steps of the above embodiment, has the same technical effects as the above embodiment, and will not be described here.

[0123] The application discloses a kind of infrared weak target detection device based on YOLOv8, device, comprising:

[0124] Data acquisition module, for pre-processing the image to be detected to obtain feature image.

[0125] The detection module comprises a target detection model of an infrared weak small target detection method based on YOLOv8, and the target detection model is used for detecting the infrared weak small target in the to-be-detected image.

[0126] The output module outputs the position information, category and confidence information of all detected weak small targets in the whole to-be-detected image according to the prediction result of the detection module.

[0127] The detection module comprises:

[0128] The image input unit pre-processes the to-be-detected image, so that the size of the input to-be-detected image is 640x640.

[0129] The feature extraction unit extracts features from the pre-processed image, comprising:

[0130] Two continuous initial convolution layers Conv0 and Conv1 perform four times down-sampling on the to-be-detected image to obtain a sub-feature map with a size of 160x160.

[0131] Three down-sampling modules composed of C2f layers and Conv convolution layers are used for three times of continuous down-sampling to obtain first feature maps with scales of 80x80, 40x40 and 20x20.

[0132] The pooling unit is composed of a pyramid pooling layer and a large kernel convolution UniRepLK Block fusion, which is used to capture the sparse features of the weak small target.

[0133] The feature fusion unit is used for multi-scale fusion of shallow features and deep features, comprising a feature fusion layer and a weak small target detection layer.

[0134] The triple attention mechanism unit performs cross-latitude interaction on the fused feature information, enhancing the features of specific channels and spatial regions.

[0135] The detection probe unit comprises a weak small target detection head and three detection heads.

Claims

1. A YOLOv8-based infrared dim small target detection method, characterized in that, The method comprises the following steps: performing preprocessing on a to-be-detected image to obtain a feature image; inputting the feature image into an optimized target detection model, and performing weak and small target detection by using the optimized target detection model; the optimization of the target detection model comprises: fusing a perception large kernel convolution UniRepLK Block in a spatial pyramid pooling layer of a feature extraction network of the target detection model to obtain multi-scale sparse features of a target; introducing a triple attention mechanism layer at an output end of a feature fusion network of the target detection model to adaptively focus feature information integrated with the sparse features; constructing a multi-scale target detection network to detect the adaptively focused feature information; the optimization of the target detection model further comprises: optimization of the feature extraction network: using two initial convolution layers to downsample the feature image to extract original pixel primary features; then continuously performing three times of downsampling by using a C2f layer combined with a convolution layer to obtain three first feature maps of different scales; optimization of the feature fusion network: continuously performing two times of upsampling on the multi-scale sparse features obtained after the spatial pyramid pooling layer is pooled to obtain two second feature maps of different scales; respectively fusing the first feature map and the second feature map of the corresponding scale, and outputting the feature information extracted by using a C2f layer to the triple attention mechanism layer; the optimization of the target detection model further comprises: adding an up-sampling layer with a factor of 4 as a weak and small target detection layer after the two times of upsampling to obtain a third feature map; fusing the third feature map and the first feature map, and outputting the fused feature information to the target detection layer by using a C2f layer.

2. The method of claim 1, wherein, The preprocessing of the to-be-detected image to obtain the feature image comprises: performing adaptive size adjustment on the input to-be-detected image; normalizing pixel values of the to-be-detected image to [0, 1] and performing data enhancement processing; segmenting the to-be-detected image into N×N grid units to obtain a plurality of feature images of different sizes.

3. The method of claim 1, wherein, The perception large kernel convolution comprises: a plurality of parallel layers, and (k-1) r+1≤K; wherein, k is a convolution kernel of the parallel layer, r is an inflation rate of the parallel layer, and K is a large kernel.

4. The method of claim 3, wherein, The large kernel K=13, the number of parallel layers n=5, the convolution kernel of the parallel layer k=(7, 5, 3, 3, 3), and the inflation rates are r=(2, 1, 3, 4, 5).

5. The method of claim 1, wherein, The triple attention mechanism layer comprises: a first branch for capturing interaction features between a channel C and a spatial dimension W; a second branch for capturing interaction features between the channel C and a spatial dimension H; a third branch for constructing spatial attention to capture a dependency relationship between the spatial dimensions W and H; outputs of the first branch, the second branch and the third branch are aggregated by averaging to form an attention weighting tensor, and the generated attention weighting tensor is applied to an original tensor of the input to-be-detected image.

6. The method of claim 5, wherein, The attention weighting tensor is: ； wherein σ represents a sigmoid activation function; 1、 2 and 3 represent a standard two-dimensional convolutional layer; denotes: the rotated tensor in the first branch; denotes: the rotated tensor in the second branch; denotes: the rotated tensor in the third branch, consistent with the original tensor; denotes: the tensor after the first branch is restored in shape by the z-pooling layer; denotes: the tensor after the second branch is reshaped by the z-pooling layer.

7. The method of claim 1, wherein, the detection of the adaptively focused feature information comprises: generating prediction boxes: for each feature image of different sizes containing any target center point within the N×N grid cells, predicting multiple bounding boxes and class probabilities for each grid cell in each feature image of different sizes; prediction box decoding: predicting the offset of the bounding box by the Anchor-Free mechanism, and removing the low-confidence predictions and overlapping prediction boxes by the confidence threshold filtering and non-maximum suppression algorithm after decoding the prediction boxes, outputting the detection results: the post-processed prediction boxes are sorted according to the confidence, and the class, position and confidence of the target in the image to be detected are outputted.

8. An infrared dim small target detection device based on YOLOv8, characterized in that, comprising: a data acquisition module for pre-processing the image to be detected to obtain a feature image; a detection module for detecting the feature image by the target detection model of the method according to any one of claims 1-7.

Images