Infrared dim small target detection method based on improved YOLOv5

Abstract

The present application aims to provide an infrared dim small target detection method based on improved YOLOv5, comprising the following steps: A, constructing a deep neural network structure, the deep neural network structure is as follows: Backbone network, Neck network, YOLO detection head; B, training the deep neural network to obtain a trained deep neural network; C, based on the trained deep neural network, inputting an infrared image into the Backbone network, and sequentially processing through each module in the network, wherein the processing results of the CSP1_2 module, the CSP1_3 module and the SPPF module are input into the Neck network; D, obtaining three feature maps with different sizes by processing the Neck network, and inputting the feature maps into the YOLO detection head, and obtaining the final result after the prediction processing of the YOLO detection head. The method can effectively improve the recognition accuracy and speed of the infrared dim small target.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of infrared weak small target detection, and particularly relates to an infrared weak small target detection method based on an improved YOLOv5s. BACKGROUND

[0002] Infrared weak small target detection is a key technology in the field of computer vision, and is widely used in military and civilian fields, such as aerospace, precision guidance, infrared early warning, and unmanned aerial vehicle detection. At present, there are still many difficulties in the technology of infrared weak small target detection. First, because the target is far away from the infrared detector, the size of the target on the image is small, usually a few pixels, the signal-to-noise ratio of the target is low, the signal is relatively weak, and there is a lack of rich detail information. In addition, long-distance imaging makes the scene span very large, the background is complex, and it is easily affected by external factors such as weather and noise, which makes the target submerged in the background, which increases the difficulty of weak small target detection. Due to the above difficulties, the existing weak small target detection algorithm has limitations and cannot meet the actual needs. Therefore, it is of great significance to study an accurate and fast infrared weak small target detection algorithm. SUMMARY

[0003] The present application aims to provide an infrared weak small target detection method based on an improved YOLOv5s, which can effectively improve the recognition accuracy and speed of infrared weak small targets.

[0004] The technical solution of the present application is as follows:

[0005] The infrared weak small target detection method based on the improved YOLOv5s includes the following steps:

[0006] A. Construct a deep neural network structure, which is specifically as follows:

[0007] Backbone network, Neck network, YOLO detection head;

[0008] Among them, the Backbone network includes CBS module, CBS module, CSP1_1 module, CBS module, CSP1_2 module, CBS module, CSP1_3 module, CBS module, CSP_BT module, SPPF module connected in turn;

[0009] B. Train the deep neural network to obtain a trained deep neural network;

[0010] C. Based on the trained deep neural network, the infrared image is input into the Backbone network and processed sequentially through each module in the network. The processing results of the CSP1_2 module, CSP1_3 module, and SPPF module are respectively input into the Neck network for processing.

[0011] D. The Neck network processes the data to obtain three feature maps of different sizes, which are then input into the YOLO detection head. After prediction processing by the YOLO detection head, the final result is obtained.

[0012] In step C, the processing procedure of the neck network is as follows:

[0013] The processing result of the SPPF module in the Backbone network is processed by the Coord Conv convolution, then upsampled to the same size as the processing result of the CSP1_3 module in the Backbone network. These results are then concatenated together in the channel dimension using Concat. Next, the result passes through the CSP2_1 module and the Coord Conv convolution to obtain processing result a. Result a is upsampled to the same size as the processing result of the CSP1_2 module in the Backbone network, then concatenated together in the channel dimension using Concat. It is then processed by the CSP2_1 module to obtain processing result b. Result b is split into two paths: one path is processed by the Coord Conv convolution and input into the YOLO detection head module; the other path is processed by the CBS module and concatenated together with result a in the channel dimension using Concat. Finally, it is processed by the CSP2_1 module to obtain processing result c. Result c is split into two paths: one path is processed by the Coord Conv convolution and input into the YOLO detection head module; the other path is processed by the CBS module and then concatenated together in the channel dimension with result a. Finally, it is processed by the CSP2_1 module to obtain processing result c. Result c is split into two paths: one path is processed by the Coord Conv convolution module and input into the YOLO detection head module. After Conv convolution, the result is input into the yoyo detection head module. The result of the other path is processed by the CBS module and then processed by the SPPF module. The result of the Coord Conv convolution is then concatenated in the channel dimension by Concat, and then processed by the CSP2_1 module to obtain the processed result d. The processed result d is then processed by the Coord Conv convolution and input into the yoyo detection head module.

[0014] The processing procedure in the CSP2_1 module is as follows: the input result is divided into two paths. One path is processed by two CBS modules in sequence and then input into the Concat function; the other path is processed by the CBS module and then input into the Concat function; the Concat module concatenates the channel dimensions and then processes the results in the CBS module to obtain the output result.

[0015] The general expression for modules CSP1_1, CSP1_2, and CSP1_3 is CSP1_n, where n represents the number of times the internal module is repeated.

[0016] The processing in CSP1_n is as follows: The input result is divided into two paths. One path is processed by the CBS module, then by the res_n module, and then input into the Concat function. The other path is processed by the CBS module and then input into the Concat function. Concat concatenates the results in the channel dimension and then processes them by the CBS module to obtain the output result.

[0017] The processing procedure in the res_n module is as follows: the input result is processed by two CBS modules in sequence, and then added to the input result to obtain the output result.

[0018] The processing procedure in the CSP_BT module is as follows: the input result is divided into two paths. One path is processed by the CBS module and the BoTR module in sequence and then input into the Concat function. The other path is processed by the CBS module and then input into the Concat function. The Concat module concatenates the channel dimensions and then processes the results by the CBS module to obtain the output result.

[0019] The processing procedure in the BoTR module is as follows: the input result is processed sequentially by the CBS module and the MHSA module, and then added to the input result to obtain the output result.

[0020] The CBS module includes a convolutional function, a BN normalization function, and a SiLU activation function connected in sequence.

[0021] Step B, the training process of the deep neural network, is as follows:

[0022] 70% of the samples were randomly selected for training the model, 10% for testing, and 20% for validation. The sample resolution was 256. The NWD loss function was used as the localization loss to train the model. The trained network model was then used to predict the test samples, achieving automatic detection of infrared-weak targets.

[0023] The NWD loss function is as follows:

[0024]

[0025] Loss NWD =1-NWD(N a N b )

[0026] Where, N a N represents the Gaussian distribution of the predicted bounding box. b W2 represents the Gaussian distribution of the ground truth bounding boxes. 2 (N a N b ) represents N a and N bWasserstein distance between the Gaussian distributions of the two boxes; NWD(N a N b The expression denoted by represents the exponential form of the normalized Wasserstein distance between two Gaussian distributions, used to represent the similarity of the Gaussian distributions of the bounding boxes; C represents a constant; loss NWD The formula serves as the localization loss function.

[0027] The derivation of the NWD loss function is as follows:

[0028] The bounding box is modeled as a two-dimensional Gaussian distribution to represent the importance of the target pixels within the box. Assume the center point of the bounding box is at coordinates (c...). x c y Given a bounding box with width w and height h, the equation of the inscribed ellipse can be expressed as Equation 3. The probability density function of a two-dimensional Gaussian distribution can be expressed as Equation 4, where X represents the coordinates (x, y), μ represents the mean vector, and ∑ represents the covariance matrix. When (X-μ) T ∑ -1 When (X-μ)=1, the inscribed ellipse of the bounding box is a two-dimensional Gaussian distribution established by N(μ,∑), as shown in Equation 5.

[0029]

[0030]

[0031]

[0032] The similarity between predicted and ground truth bounding boxes is represented by the distance between two Gaussian distributions. The distance between the two Gaussian distributions is calculated using the Wasserstein distance, which is defined as shown in Equation 6. For two Gaussian distributions of bounding boxes, the Wasserstein distance can be expressed as Equation 7. However, a simple distance metric cannot be used to represent similarity. Therefore, after normalization, the similarity metric is expressed in exponential form, resulting in Equations 1 and 2.

[0033]

[0034]

[0035] The method of this invention incorporates a multi-head self-attention mechanism into the CSP structure to enhance the expressive power of global information;

[0036] The method of this invention uses CoordConv instead of 1×1 convolution, enabling the convolution to perceive positional information and optimize network parameters;

[0037] The method of this invention uses NWD instead of CIOU as the localization loss during training, which reduces the sensitivity to positional deviations of weak targets during training.

[0038] The method of this invention has been further optimized based on the YOYOv5s structure, which can effectively improve the accuracy and speed of identifying small infrared targets and has good application prospects. Attached Figure Description

[0039] Figure 1 This is a schematic diagram of the deep neural network structure in Embodiment 1 of the present invention;

[0040] Figure 2 This is a schematic diagram of the MHSA module in Embodiment 1 of the present invention;

[0041] Figure 3 This is a schematic diagram of the coordinate convolution CoordConv structure in Embodiment 1 of the present invention;

[0042] Figure 4 This is a comparison chart of the detection results of the method of Embodiment 1 of the present invention and other YOLO algorithms in the same scene. Detailed Implementation

[0043] The present invention will now be described in detail with reference to the accompanying drawings and embodiments.

[0044] Example 1

[0045] The infrared weak target detection method based on the improved YOLOv5s in this embodiment includes the following steps:

[0046] The infrared weak target detection method based on the improved YOLOv5s includes the following steps:

[0047] A. Construct a deep neural network structure, the specific structure of which is as follows:

[0048] Backbone network, Neck network, YOLO detection head;

[0049] The Backbone network includes the following modules connected in sequence: CBS module, CBS module, CSP1_1 module, CBS module, CSP1_2 module, CBS module, CSP1_3 module, CBS module, CSP_BT module, and SPPF module.

[0050] B. Train the deep neural network to obtain a trained deep neural network;

[0051] The training process of a deep neural network is as follows:

[0052] 70% of the samples were randomly selected for training the model, 10% for testing, and 20% for validation. The sample resolution was 256. The NWD loss function was used as the localization loss to train the model. The trained network model was then used to predict the test samples, achieving automatic detection of infrared-weak targets.

[0053] The NWD loss function is as follows:

[0054]

[0055] Loss NWD =1-NWD(N a N b (2)

[0056] Where, N a N represents the Gaussian distribution of the predicted bounding box. b W2 represents the Gaussian distribution of the ground truth bounding boxes. 2 (N a N b ) represents N a and N b Wasserstein distance between the Gaussian distributions of the two boxes; NWD(N a N b The expression ) represents the exponential form of the normalized Wasserstein distance between two Gaussian distributions, used to represent the similarity of the Gaussian distributions of the bounding boxes; C represents a constant; Loss NWD The formula serves as the localization loss function.

[0057] C. Based on the trained deep neural network, the infrared image is input into the Backbone network and processed sequentially through each module in the network. The processing results of the CSP1_2 module, CSP1_3 module, and SPPF module are respectively input into the Neck network for processing.

[0058] The processing procedure for the Neck network is as follows:

[0059] The processing result of the SPPF module in the Backbone network is processed by the Coord Conv convolution, then upsampled to the same size as the processing result of the CSP1_3 module in the Backbone network. These results are then concatenated together in the channel dimension using Concat. Next, the result passes through the CSP2_1 module and the Coord Conv convolution to obtain processing result a. Result a is upsampled to the same size as the processing result of the CSP1_2 module in the Backbone network, then concatenated together in the channel dimension using Concat. It is then processed by the CSP2_1 module to obtain processing result b. Result b is split into two paths: one path is processed by the Coord Conv convolution and input into the YOLO detection head module; the other path is processed by the CBS module and concatenated together with result a in the channel dimension using Concat. Finally, it is processed by the CSP2_1 module to obtain processing result c. Result c is split into two paths: one path is processed by the Coord Conv convolution and input into the YOLO detection head module; the other path is processed by the CBS module and then concatenated together in the channel dimension with result a. Finally, it is processed by the CSP2_1 module to obtain processing result c. Result c is split into two paths: one path is processed by the Coord Conv convolution module and input into the YOLO detection head module. After Conv convolution, the result is input into the YOLO detection head module. The result of the other path is processed by the CBS module and then processed by the SPPF module. After the result of the Coord Conv convolution is processed, it is concatenated in the channel dimension by Concat. Then it is processed by the CSP2_1 module to obtain the processing result d. The processing result d is processed by the Coord Conv convolution and then input into the YOLO detection head module.

[0060] D. The Neck network processes the data to obtain three feature maps of different sizes, which are then input into the YOLO detection head. After prediction processing by the YOLO detection head, the final result is obtained.

[0061] The processing procedure in the CSP2_1 module is as follows: the input result is divided into two paths. One path is processed by two CBS modules in sequence and then input into the Concat function; the other path is processed by the CBS module and then input into the Concat function; the Concat module concatenates the channel dimensions and then processes the results in the CBS module to obtain the output result.

[0062] The general expression for modules CSP1_1, CSP1_2, and CSP1_3 is CSP1_n, where n represents the number of repetitions of the internal module.

[0063] The processing in CSP1_n is as follows: The input result is divided into two paths. One path is processed by the CBS module, then by the res_n module, and then input into the Concat function. The other path is processed by the CBS module and then input into the Concat function. Concat concatenates the results in the channel dimension and then processes them by the CBS module to obtain the output result.

[0064] The processing procedure in the res_n module is as follows: the input result is processed by two CBS modules in sequence, and then added to the input result to obtain the output result.

[0065] The processing procedure in the CSP_BT module is as follows: the input result is divided into two paths. One path is processed by the CBS module and the BoTR module in sequence and then input into the Concat function. The other path is processed by the CBS module and then input into the Concat function. The Concat module concatenates the channel dimensions and then processes the results by the CBS module to obtain the output result.

[0066] The processing procedure in the BoTR module is as follows: the input result is processed sequentially by the CBS module and the MHSA module, and then added to the input result to obtain the output result.

[0067] The CBS module includes a convolutional function, a BN normalization function, and a SiLU activation function connected in sequence.

[0068] Example 2

[0069] The specific structure and process of the MHSA module in Example 1 are as follows: Figure 2 As shown:

[0070] The input feature matrix, after being convolved by 1×1, is used as W. Q W K W V The dimensions are H×W×d, where d represents the dimension of a single token. The height and width are initialized simultaneously using encoding R. h and R w The position code r is obtained by adding the parameters q, k, v, and r through a broadcast mechanism. The dimensions of the four parameters q, k, v, and r are H×W×d. The content-position output matrix qr is obtained by multiplying q and r through matrix multiplication. T Multiplying q and k by matrices yields the content-content output matrix qk. T ,qr T With qk T The elements are summed as matrices and then normalized using Softmax to obtain an attention matrix of size HW×HW. This attention matrix is ​​then multiplied by the value projection matrix v to obtain the output value Z, which aggregates global information. The computation of the multi-head self-attention layer is performed in parallel. We used four heads, assigning different weights to the feature maps through a self-attention mechanism. This allows the network to focus more on the target information in the feature maps, significantly improving the detection performance of the CNN model.

[0071] like Figure 3As shown, Example 1 employs coordinate convolution (CoordConv). Compared to traditional convolution, CoordConv adds two position channels to the input: one representing the x-coordinate and the other representing the y-coordinate. These two coordinate channels are concatenated with the input channels before the convolution operation, enabling the convolution process to be aware of positional information.

[0072] Example 3

[0073] The detection results of YOYOv3, YOYOv4, YOYOv5, YOYOv7 and the algorithm of Example 1 (Ours Method) were compared in the same scene. The specific comparison results are as follows: Figure 4 ,like Figure 4 visible:

[0074] The detection performance of the improved YOLOv5s algorithm was evaluated using image data from six different scenarios: sky, mountains, water, buildings, roads, and woodlands. The same parameters were used for testing, and the results are shown in the figure. YOLOv4-tiny failed to detect targets in the sky and mountain scenarios and exhibited false detections against a water background. YOLOv7-tiny also showed false detections against a water background, but successfully detected targets in other scenarios. YOLOv3-tiny and YOLOv5s detected targets normally in all six scenarios without false detections or missed detections. However, the target recognition rate in each scenario was lower than that of the algorithm in Example 1. The algorithm in Example 1 showed a significant improvement in the recognition rate for small targets. Overall, the improved YOLOv5s algorithm of this invention significantly outperforms other algorithms in terms of detection and recognition rate.

Claims

1. A method for detecting weak infrared targets based on an improved YOLOv5s, characterized in that, Includes the following steps: A. Construct a deep neural network structure, the specific structure of which is as follows: Backbone network, Neck network, YOLO detection head; The Backbone network includes the following modules connected in sequence: CBS module, CBS module, CSP1_1 module, CBS module, CSP1_2 module, CBS module, CSP1_3 module, CBS module, CSP_BT module, and SPPF module. B. Train the deep neural network to obtain a trained deep neural network; C. Based on the trained deep neural network, the infrared image is input into the Backbone network and processed sequentially through each module in the network. The processing results of the CSP1_2 module, CSP1_3 module, and SPPF module are respectively input into the Neck network for processing. D. The Neck network processes the data to obtain three feature maps of different sizes, which are then input into the YOLO detection head. After prediction processing by the YOLO detection head, the final result is obtained. In step C, the processing procedure of the neck network is as follows: The processing result of the SPPF module in the Backbone network is processed by the Coord Conv convolution, then upsampled to the same size as the processing result of the CSP1_3 module in the Backbone network. These results are then concatenated together in the channel dimension using Concat. Next, the result passes through the CSP2_1 module and the Coord Conv convolution to obtain processing result a. Result a is upsampled to the same size as the processing result of the CSP1_2 module in the Backbone network, then concatenated together in the channel dimension using Concat. It is then processed by the CSP2_1 module to obtain processing result b. Result b is split into two paths: one path is processed by the Coord Conv convolution and input into the YOLO detection head module; the other path is processed by the CBS module and concatenated together with result a in the channel dimension using Concat. Finally, it is processed by the CSP2_1 module to obtain processing result c. Result c is split into two paths: one path is processed by the Coord Conv convolution and input into the YOLO detection head module; the other path is processed by the CBS module and then concatenated together in the channel dimension with result a. Finally, it is processed by the CSP2_1 module to obtain processing result c. Result c is split into two paths: one path is processed by the Coord Conv convolution module and input into the YOLO detection head module. After Conv convolution, the result is input into the YOLO detection head module. The result of the other path is processed by the CBS module and then processed by the SPPF module. The result of the Coord Conv convolution is then concatenated in the channel dimension by Concat, and then processed by the CSP2_1 module to obtain the processed result d. The processed result d is then processed by the Coord Conv convolution and input into the YOLO detection head module.

2. The infrared weak target detection method based on improved YOLOv5s as described in claim 1, characterized in that: The processing procedure in the CSP2_1 module is as follows: the input result is divided into two paths. One path is processed by two CBS modules in sequence and then input into the Concat function; the other path is processed by the CBS module and then input into the Concat function; the Concat module concatenates the channel dimensions and then processes the results in the CBS module to obtain the output result.

3. The infrared weak target detection method based on improved YOLOv5s as described in claim 1, characterized in that: The general expression for modules CSP1_1, CSP1_2, and CSP1_3 is CSP1_n, where n represents the number of repetitions of the internal module. The processing in CSP1_n is as follows: The input result is divided into two paths. One path is processed by the CBS module, then by the res_n module, and then input into the Concat function. The other path is processed by the CBS module and then input into the Concat function. Concat concatenates the results in the channel dimension and then processes them by the CBS module to obtain the output result.

4. The infrared weak target detection method based on improved YOLOv5s as described in claim 3, characterized in that: The processing procedure in the res_n module is as follows: the input result is processed by two CBS modules in sequence, and then added to the input result to obtain the output result.

5. The infrared weak target detection method based on improved YOLOv5s as described in claim 1, characterized in that: The processing procedure in the CSP_BT module is as follows: the input result is divided into two paths. One path is processed by the CBS module and the BoTR module in sequence and then input into the Concat function. The other path is processed by the CBS module and then input into the Concat function. The Concat module concatenates the channel dimensions and then processes the results by the CBS module to obtain the output result.

6. The infrared weak target detection method based on improved YOLOv5s as described in claim 5, characterized in that: The processing procedure in the BoTR module is as follows: the input result is processed sequentially by the CBS module and the MHSA module, and then added to the input result to obtain the output result.

7. The infrared weak target detection method based on improved YOLOv5s as described in any one of claims 1-6, characterized in that: The CBS module includes a convolutional function, a BN normalization function, and a SiLU activation function connected in sequence.

8. The infrared weak target detection method based on improved YOLOv5s as described in claim 1, characterized in that: Step B, the training process of the deep neural network, is as follows: 70% of the samples were randomly selected for training the model, 10% for testing, and 20% for validation. The sample resolution was 256. The NWD loss function was used as the localization loss to train the model. The trained network model was used to predict the test samples to achieve automatic detection of infrared weak targets.

9. The infrared weak target detection method based on the improved YOLOv5s as described in claim 8, characterized in that: The NWD loss function is as follows: ⑴ ⑵ in, The Gaussian distribution represents the predicted bounding box. The Gaussian distribution represents the true bounding box; express and Wasserstein distance between the Gaussian distributions of the two boxes; This represents the exponential form of the normalized Wasserstein distance between two Gaussian distributions, used to represent the similarity of the Gaussian distributions of the bounding boxes; C represents a constant. The formula serves as the localization loss function.

Images