Infrared small target detection system and method based on improved YOLOv5s and lightweight method

Abstract

The invention discloses an infrared small target detection system and method based on improved YOLOv5s and a lightweight method, and relates to the technical field of infrared target detection. The invention provides a light-weight infrared small target detection algorithm based on improved YOLOv5s, namely, YOLOv5s-SIT (Small Infruled Target), and the light-weight infrared small target detection algorithm based on the improved YOLOv5s is a light-weight infrared small target detection algorithm based on improved YOLOv5s-SIT (SIT, Small Infruled Target). The main improvements in the designed YOLOv5s-SIT algorithm network comprise an ECBAM attention module for enhancing important features and inhibiting non-important features, an sppx fusion module for enriching the expression ability of a feature map, and a loss function ComLoss for detecting the resetting of an anchor frame and reconstructing the algorithm. In order to improve the precision of a YOLOv5s-SIT algorithm on an infrared small target detection task, a model lightweight strategy is designed for the structural characteristics of the algorithm, corresponding pruning and compression are performed on a backbone network of the algorithm, and finally a lightweight infrared small target detection algorithm is obtained.

Description

Technical Field

[0001] The present invention relates to the technical field of infrared target detection, and in particular to an infrared small target detection system, method and lightweight method based on improved YOLOv5s. Background Art

[0002] At present, with the increasing demand for precise strike applications of infrared guided weapon systems in the military field and the increasing demand for applications in forest fire prevention, wildlife observation and rescue in the civilian field, the performance requirements for infrared small target detection have been correspondingly increased. Since infrared small targets have the characteristics of low signal-to-noise ratio, poor contrast, small target area, weak grayscale, and the feature extraction of traditional infrared small target detection methods relies on manual experience, its generalization ability is poor and it cannot achieve effective migration of algorithms between multiple platforms. Therefore, infrared small target detection algorithms based on deep learning have become a research hotspot. However, due to the particularity of infrared small target detection tasks, there are currently almost no public data sets that meet the conditions, and target detection algorithms based on deep learning are very dependent on data sets, which to a certain extent has led to less public research on the application of deep learning theory to infrared small target detection tasks at home and abroad.

[0003] At the same time, the platform on which the infrared guidance system / infrared small target detection system is based has its own particularity, and it is impossible to increase its computing power and data storage space indefinitely. volume (that is, the storage space occupied) is large and it is difficult to meet the actual deployment requirements. Therefore, the lightweight research of algorithm models is crucial. There are still some problems in its application to actual scenarios. Among them, the most prominent problem is how to migrate it to an available hardware platform, because the algorithm model has the disadvantages of huge number of parameters, large amount of calculation, and high power consumption. The idea of ​​lightweight design of algorithm models is to reduce network parameters without losing detection accuracy, so as to achieve the effect of slimming down the algorithm model and no longer be restricted by the conditions of the hardware platform. Therefore, in order to meet the three requirements of faster, more accurate, and smaller, it is inevitable to study the lightweight of infrared small target detection algorithm.

[0004] There are two research directions for infrared small target detection technology: one is to automatically identify targets by extracting target features, including classic pattern recognition algorithms, knowledge-based algorithms, and artificial neural network technology; the other is to automatically identify targets by template matching. Since the second method is relatively backward and time-consuming, it is not suitable for the application scenarios mentioned in this patent. Here we only introduce the first infrared small target detection algorithm based on deep learning.

[0005] In 2016, Redmon et al. proposed the YOLO algorithm, which was the first to realize the application of a single-stage process in target detection. Its advantages are fast speed and small algorithm size, and its main disadvantage is low accuracy. In 2017, Lin et al. proposed RetinaNet, which uses Focal Loss to express cross entropy and improves the sample weights in the image set, so that the algorithm can still achieve good results in samples with higher detection difficulty. In 2019, Zhao et al. proposed TBC-net, which uses strategies such as semantic constraint modules to solve the problem caused by extremely unbalanced targets and backgrounds when the network learns small target features. In 2021, Hou et al. proposed RISTD-net, which combines traditional feature extraction methods with deep learning feature extraction network frameworks to obtain an infrared small target detection algorithm with good robustness.

[0006] Existing research on lightweight algorithm models is mainly divided into improvements to the algorithm model itself and improvements to the algorithm model deployment environment. The improvement method for the algorithm model itself is mainly achieved by reducing network parameters and reducing model complexity, which is also the method used in this patent. Commonly used algorithm model compression technologies include network pruning, network decomposition, knowledge distillation, parameter sharing, quantization, and compact network design. Improvements to the algorithm model deployment environment include hardware acceleration and software acceleration.

[0007] Some current target detection algorithms such as RCNN, FasterR-CNN, FPN, etc. have limited detection speed and are prone to misjudgment. The emergence of the YOLO series of algorithms has gradually solved the problems of speed and misjudgment. However, the obvious disadvantages of the previous generations of algorithms are insufficient accuracy and high missed detection rates. With the continuous optimization and upgrading of the YOLO series, such as the improvement of RetinaNet, TBCnet, and RISTDnet network structures, this problem has been greatly improved, but it has not yet been completely and properly solved.

[0008] Due to the rapid development of deep learning theory, the updated target detection algorithms are continuously improving their ability to detect small targets, and the algorithm model's ability to extract features is also constantly improving, resulting in the complexity and size of the algorithm model gradually increasing. Therefore, in order to transplant the algorithm from the laboratory stage to the available mobile platform, the algorithm model must be lightweight to facilitate better adaptation and operation of the algorithm. Existing research on lightweight algorithm models is mainly divided into improvements to the algorithm model itself, reducing network parameters, and reducing the complexity of the algorithm model; as well as improvements to the algorithm model deployment environment, hardware and software acceleration. Commonly used algorithm model lightweight technologies include network pruning, network decomposition, knowledge distillation, parameter sharing, quantization, and compact network design.

[0009] It can be seen that the existing technology has the following shortcomings that have not yet been resolved:

[0010] Infrared small targets lack detailed texture structures. As they are too small to be effectively captured and identified in the field of view, the detection accuracy of existing algorithms is not high. In addition, the detection speed of existing infrared small target detection algorithms is not fast enough to cope with the fast movement of the platform on which they are mounted. Summary of the invention

[0011] In view of this, the present invention provides an infrared small target detection system, method and lightweight method based on improved YOLOv5s, which can improve the accuracy of the YOLOv5s-SIT algorithm for infrared small target detection tasks and increase the detection speed.

[0012] To achieve the above-mentioned purpose, the technical solution of the present invention is: a lightweight infrared small target detection system based on improved YOLOv5s, the YOLOv5s network includes a backbone network BackBone, an enhanced backbone feature extraction network Neck and a target detection head network Head, and the Spp module in the backbone network Backbone part of the YOLOv5s network is replaced by a Spp_x module; the Backbone part receives the input image, and the Spp-x module outputs the solved image data stream to the first CSP2-1.

[0013] The enhanced backbone feature extraction network Neck includes a first CSP2-1, a first CBL, a first Upsample, a first Contcat, a second CSP2-1, a first ECBAM, a second CBL, a second Upsample, a second Contcat, a third CSP2-1, a third Contcat, a third CBL, a second ECBAM, a fourth CSP2-1, a fourth Contcat, a fourth CBL, a third ECBAM and a fifth CSP2-1 connected in sequence; the first CBL is connected to the fourth Contcat, and the second CBL is connected to the third Contcat.

[0014] The target detection head network Head includes three Conv blocks, wherein the processing results of the third to fifth CSP2-1 in Neck are respectively connected to the three Conv blocks, and the three Conv blocks respectively output images of three detected target sizes.

[0015] The Spp_x module includes an eighth CBL block and three Max pool 2Ds connected in sequence, the outputs of the CBL block and the three Maxpool 2Ds are all connected to a Concat, the output of the Concat is connected to a ninth CBL block, and the output of the ninth CBL block is used as the output of the Spp_x module.

[0016] The structures of the first to third ECBAMs are the same, specifically including an improved channel attention module and a spatial attention module; the feature map F is input into the improved channel attention module to extract the channel attention M c (F), and then fused with the input feature map F to obtain the feature map F'. The input feature map enters the spatial attention module to extract the spatial attention M s (F), and then fused with the feature map F' to obtain the output feature map F".

[0017] Furthermore, the backbone network Backbone part includes an input module, a Focus, a fifth CBL, a CSP1-1, a sixth CBL, CSP1-2, a seventh CBL, CSP1-3, an eighth CBL and an Spp-x module connected in sequence.

[0018] The processing result of the CSP1-2 module is connected to the second Concat in the Neck part.

[0019] The processing results of the CSP1-3 modules are connected to the first Concat in the Neck part.

[0020] Furthermore, the operation process of the improved channel attention module is as follows:

[0021] The input convolution feature map is F C×H×W , then the channel attention M c (F) can be expressed as formula (3):

[0022] M c (F) = σ(Conv k (Avgpool(F))+Conv k (Maxpool(F))) (3)

[0023] Among them, σ represents the Sigmoid activation function, Conv k It represents a one-dimensional convolution kernel with a length of k, Avgpool(F) is average pooling, and Maxpool(F) is maximum pooling.

[0024] The formula (4) for calculating the convolution kernel size k based on the number of feature map channels c is as follows:

[0025]

[0026] Among them, odd represents the odd number closest to c, γ is 2, and b is 1. Through the above formula (4), the length of the one-dimensional convolution kernel is automatically specified for the channel attention mechanism to avoid introducing additional hyperparameters.

[0027] The operation process of the spatial attention module is:

[0028] Spatial AttentionM s The calculation process of (F) is expressed as formula (5):

[0029] M s (F) = σ(f 7×7 ([Avgpool(F);Maxpool(F)])) (5)

[0030] Among them, F is the input feature map, σ represents the Sigmoid activation function, and f 7×7 It represents the convolution calculation with a convolution kernel size of 7×7, Avgpool(F) is the average pooling, and Maxpool(F) is the maximum pooling.

[0031] Furthermore, the output of the target detection head network Head contains three scale information, corresponding to three target sizes 19×19, 38×38 and 76×76.

[0032] Another embodiment of the present invention further provides a lightweight infrared small target detection method based on improved YOLOv5s. The target size detected by the above system is calibrated in the following manner:

[0033] The output of the target detection head network Head contains three scale information, each scale includes three anchor frames, that is, nine scales of anchor frames are required. Then the detection method clusters all the true value frames into nine categories.

[0034] Set the width and height of the target truth box to w and h respectively, and then use the K-means algorithm to And w, h cluster the training data sets respectively.

[0035] The specific process of this detection method is as follows:

[0036] Step 1: First select 9 target truth boxes as the initial category centroids.

[0037] Step 2: Calculate the distance between each target ground truth box and the centroids of the nine categories, and then assign each ground truth box to the category with the closest distance.

[0038] Step 3: Calculate the centroid of each category based on the latest category assignment results.

[0039] Step 4: If the termination condition is met, the algorithm is completed, otherwise it returns to the second step to continue execution.

[0040] Termination conditions: the number of reallocated truth boxes is less than the minimum value; the offset of the K-means center is less than the minimum value; the sum of squared errors between the truth box and the category centroid reaches a local minimum.

[0041] Among them, the centroid distance calculation method in the clustering process adopts L1 The distance measurement principle, expression (1) is as follows:

[0042]

[0043] Among them, α i represents the i-th element, is the centroid element in the current iteration process, and k is the number of iterations.

[0044] The clustering process will produce 3 The centroid of and 3 centroids of w / h are used to generate 9 anchor boxes of different scales and aspect ratios according to the clustering results as the anchor boxes of this method; among them, represents the basic scale of the anchor box, and w / h represents the aspect ratio of the anchor box.

[0045] When calculating the specific anchor frame of a certain input feature map, let ah ij and aw ij are the height and width of the jth anchor box of the ith feature map, respectively. Then ah ij and aw ij It is expressed by formula (2):

[0046]

[0047] Among them, ρ j Indicates the jth w / h; basesize i Indicates the i-th

[0048] Furthermore, the YOLOv5s-SIT algorithm adopts the composite loss function ComLoss strategy, in which the confidence loss is calculated using the binary cross entropy loss function; the category loss is calculated using the multi-classification cross entropy loss function; and the position loss is calculated using the CIoU loss function.

[0049] Furthermore, the mathematical principle of the binary cross entropy loss function is as follows (6):

[0050]

[0051] Among them, L(x,c) is the binary cross entropy function, x is the label of the category, is the true probability of the category label in the i-th row and j-th column, the positive category takes the value of 1, the negative category takes the value of 0, and c is the predicted probability of the corresponding category. is the predicted expectation of the current i-th category.

[0052] The mathematical principle of the multi-classification cross entropy loss function is as follows (7):

[0053]

[0054] Among them, L(x,c) is the multi-classification cross entropy function, x is the label of the category, is the true probability of the m category in the i-th row and j-th column, the positive category takes the value of 1, the negative category takes the value of 0, and c is the predicted probability of the corresponding category. is the predicted expectation of the current i-th category.

[0055] The CIoU loss function is used to calculate the anchor box position loss, and its mathematical principle is as follows (8).

[0056]

[0057] Among them, To predict the distance between the anchor box and the centroid of the true value box, is the minimum diagonal distance between the predicted anchor box and the true value box, IoU is the intersection over union of the bounding boxes of the predicted anchor box and the true value box, α is the adjustment coefficient, v is the aspect ratio of the bounding box, and as mentioned above, the width and height of the target true value box are w and h respectively.

[0058] Another embodiment of the present invention further provides a lightweight method for an infrared small target detection system based on an improved YOLOv5s. For the above system, the following steps are used for pruning during the training process to achieve lightweighting:

[0059] First, the YOLOv5s-SIT algorithm is sparsely trained by introducing the parameter γ in the loss function of sparse training. The loss function in sparse training is shown in formula (9):

[0060] L sp =∑ (x,y) l(f(x,W),y)+λ∑ γ∈Γ M(γ) (9)

[0061] where ∑ (x,y) l(f(x,W),y) is the original composite loss function of the algorithm, L sp is the loss function of the algorithm's sparse training; f(x,W) is the output of the algorithm's convolutional layer; λ is the sparse parameter, which represents the intensity of the sparse training; Σ γ∈Γ M(γ) represents the L of all BN layer parameters γ 1 norm, where M(γ) = |γ|, that is, taking L 1 The regularization forms the sparsification parameter γ.

[0062] The goal of sparse training of the system as claimed in claim 1 is to make the loss function L sp When the parameter γ approaches 0, the training stops.

[0063] Afterwards, a pruning strategy is adopted for the number of convolution kernels in the convolution layer of Neck, the feature extraction part of the backbone network.

[0064] Beneficial effects:

[0065] 1: The present invention provides an infrared small target detection system and method based on improved YOLOv5s, and proposes a lightweight infrared small target detection algorithm based on improved YOLOv5s, namely YOLOv5s-SIT (SIT, Small Infrared Target). Among them, the main improvements in the designed YOLOv5s-SIT algorithm network include the ECBAM attention module that enhances important features and suppresses non-important features, the spp_x fusion module that enriches the expression ability of feature maps, and the resetting of the detection anchor frame and the loss function ComLoss reconstructed for the algorithm. In order to improve the accuracy of the YOLOv5s-SIT algorithm for infrared small target detection tasks, a model lightweight strategy is designed based on the structural characteristics of the algorithm, and its backbone network is pruned and compressed accordingly, and finally a lightweight infrared small target detection algorithm is obtained.

[0066] 2: The present invention provides a lightweight method for an infrared small target detection system based on the improved YOLOv5s. The lightweight strategy in the algorithm and the designed model lightweight strategy can effectively reduce the number of network parameters and the size of the algorithm model, while significantly improving the network's running detection speed, solving the problem that the existing infrared small target detection algorithm occupies too much storage space and is difficult to be well deployed on mobile platforms such as aircraft. BRIEF DESCRIPTION OF THE DRAWINGS

[0067] Figure 1 This is the network structure diagram of the YOLOv5s-SIT algorithm;

[0068] Figure 2 A block diagram of the basic components of the YOLOv5s-SIT algorithm;

[0069] Figure 3 The block diagram of the ECBAM attention module;

[0070] Figure 4 A block diagram for improving the channel attention module;

[0071] Figure 5 Build a block diagram for the spatial attention module;

[0072] Figure 6 It is the structure diagram of Spp_x;

[0073] Figure 7 Schematic diagram of pruning strategy;

[0074] Figure 8The figure is a schematic diagram of the algorithm flow composed of embodiments 1, 2 and 3 of the present invention. DETAILED DESCRIPTION

[0075] The present invention is described in detail below with reference to the accompanying drawings and embodiments.

[0076] This patent conducts research on the infrared small target detection task in the infrared guidance system, and proposes a lightweight infrared small target detection algorithm based on the improved YOLOv5s, namely YOLOv5s-SIT (SIT, Small Infrared Target).

[0077] Among them, the main improvements in the designed YOLOv5s-SIT algorithm network include the ECBAM attention module that enhances important features and suppresses unimportant features, the spp_x fusion module that enriches the expression ability of feature maps, and the resetting of the detection anchor box and the loss function ComLoss reconstructed for the algorithm. In order to improve the accuracy of the YOLOv5s-SIT algorithm for infrared small target detection tasks, a model lightweight strategy was designed based on the structural characteristics of the algorithm, and its backbone network was pruned and compressed accordingly, and finally a lightweight infrared small target detection algorithm was obtained.

[0078] Embodiment 1:

[0079] A lightweight infrared small target detection system based on improved YOLOv5s. The existing YOLOv5s network includes the backbone network BackBone, the enhanced backbone feature extraction network Neck, and the target detection head network Head.

[0080] The important component structure of the feature extraction backbone network of YOLOv5s-SIT is CSPDarknet-53, which eliminates the limitations of the YOLOv5s version network. The CSPDarknet-53 structure can enhance the learning ability of convolutional neural networks, remove computing bottlenecks, reduce video memory usage, speed up network reasoning, and ensure high accuracy while reducing the amount of calculation. The main feature of CSPDarknet-53 is the addition of an improved residual structure block ResUnit to the network to prevent the model from degrading when the depth of the neural network model is too deep. YOLOv5s-SIT applies the CSP1_X structure and CSP2_X structure of CSPDarknet-53 to the backbone network BackBone and the enhanced backbone feature extraction network Neck respectively.

[0081] In the present invention, the Spp module in the backbone network Backbone part of the YOLOv5s network is replaced by the Spp_x module; the Backbone part receives the input image, and the Spp-x module outputs the solved image data stream to the first CSP2-1;

[0082] The enhanced backbone feature extraction network Neck includes a first CSP2-1, a first CBL, a first Upsample, a first Contcat, a second CSP2-1, a first ECBAM, a second CBL, a second Upsample, a second Contcat, a third CSP2-1, a third Contcat, a third CBL, a second ECBAM, a fourth CSP2-1, a fourth Contcat, a fourth CBL, a third ECBAM and a fifth CSP2-1 connected in sequence; the first CBL is connected to the fourth Contcat, and the second CBL is connected to the third Contcat;

[0083] The target detection head network Head includes three Conv blocks, wherein the processing results of the third to fifth CSP2-1 in Neck are respectively connected to the three Conv blocks, and the three Conv blocks respectively output images of three detected target sizes;

[0084] The Spp_x module includes an eighth CBL block and three Max pool 2Ds connected in sequence, the outputs of the CBL block and the three Maxpool 2Ds are connected to a Concat, the output of the Concat is connected to a ninth CBL block, and the output of the ninth CBL block is used as the output of the Spp_x module;

[0085] The structures of the first to third ECBAMs are the same, specifically including an improved channel attention module and a spatial attention module; the feature map F is input into the improved channel attention module to extract the channel attention M c (F), and then fused with the input feature map F to obtain the feature map F'. The input feature map enters the spatial attention module to extract the spatial attention M s (F), and then fused with the feature map F' to obtain the output feature map F".

[0086] The backbone network Backbone part includes an input module, a Focus, a fifth CBL, a CSP1-1, a sixth CBL, a CSP1-2, a seventh CBL, a CSP1-3, an eighth CBL and an Spp-x module connected in sequence.

[0087] The processing result of the CSP1-2 module is connected to the second Concat in the Neck part; the processing result of the CSP1-3 module is connected to the first Concat in the Neck part.

[0088] In the embodiment of the present invention, the ResUnit module (residual structure block) of CSP1_X (Cross Stage Paritial) is composed of two CBLs and a tensor add operation, which can increase the depth of network construction and improve the accuracy of small target detection. The composition structure is as follows: Figure 1 . Figure 1 In the example, CSP1_1 uses 1 residual component; CSP1_2 uses 2 residual components; CSP1_3 uses 3 residual components; and CSP2_1 uses 2*1=2 groups of convolutions. The advantages of using CSP1_X and CSP2_X structures are: enhancing the learning ability of the network, making the network lightweight while maintaining accuracy; reducing computing bottlenecks and improving the algorithm reasoning speed; and reducing the memory cost of the computing platform on which the algorithm is carried.

[0089] The CBL convolution block is the smallest component in the network, consisting of a Conv convolution operation, a BN (BatchNormalization) layer, and a Leaky-ReLU activation function.

[0090] Mish is a new type of activation function, and its mathematical principle is Mish = x*tanh(ln(1+e^x)). Mish activation function is an important component of CSP2_X structure.

[0091] The Focus structure consists of slices, concatenation, and convolution blocks. After the image passes through the Focus structure, the most intuitive effect is that it is downsampled and ensures that no information is lost.

[0092] The continuous deepening of the network structure in YOLOv5s-SIT also continuously increases the network's feature extraction and feature fusion capabilities. The YOLOv5s-SIT algorithm network achieves detection of three target sizes, including 19×19, 38×38, and 76×76, by layering and cascading the above basic components. The data flow of the image to be inspected starts from the Input at the bottom left of the BackBone (backbone network), enters the Focus and other modules in turn, and then begins to layer and cascade with the key modules in the Neck (enhanced backbone feature extraction network), and finally flows to the Head (target detection head network) to start the small target detection function.

[0093] The difference between the YOLOv5s-SIT algorithm network and the original algorithm YOLOv5s is that the Backbone part uses the Spp_x module with stronger receptive field capability to focus on and extract more features; the Neck part uses three ECBAM attention modules, which can effectively extract features of different dimensions from the input feature map and focus on specific small targets of different sizes, thereby achieving the purpose of the algorithm network to detect small targets. In addition, this patent redesigned the anchor frame that conforms to the algorithm network based on the task characteristics of infrared small target detection; improved and upgraded its loss function, and finally the algorithm network uses the composite loss function ComLoss.

[0094] The network structure and basic components of the YOLOv5s-SIT (SIT, Small Infrared Target) algorithm are as follows: Figure 1 and 2 shown. Figure 1 It is the network structure of YOLOv5s-SIT algorithm; Figure 2 This is a diagram of the basic components of the YOLOv5s-SIT algorithm.

[0095] ECBAM Attention Module

[0096] Based on the network design concept of ECANet, this patent improves the ECBAM attention module, which performs attention calculations on the input feature map at the spatial scale and channel scale, thereby realizing feature extraction of different dimensions of the input feature map. Applying the ECBAM attention module to the YOLOv5s target detection head, important features will be enhanced, irrelevant features will be suppressed, and the feature fusion network can eventually output fused features, which is beneficial for the subsequent convolutional layer to encode the target information at the corresponding position. Its structure is as follows Figure 3 As shown in the figure. is the matrix tensor product operator symbol.

[0097] The usage process of the ECBAM attention module: ① Input the feature map F into the improved channel attention module to extract the channel attention M c (F), and then fused with the input feature map F to obtain the feature map F'. ② The input feature map enters the spatial attention module to extract the spatial attention M s (F), and then fused with the feature map F' to obtain the output feature map F".

[0098] in, Figure 4 To improve the schematic diagram of the channel attention module, the operation process of the improved channel attention module is as follows:

[0099] Set the convolution feature map of a certain input to be F∈R C×H×W , then the channel attention M c(F) can be expressed as formula (3):

[0100] M c (F) = σ(Conv k (Avgpool(F))+Conv k (Maxpool(F)))(3)

[0101] In the formula, σ represents the Sigmoid activation function, Conv k Represents a one-dimensional convolution kernel of length k. Avgpool(F) is average pooling. Maxpool(F) is maximum pooling.

[0102] The formula (4) for calculating the convolution kernel size k according to the number of feature map channels c is as follows.

[0103]

[0104] In the formula, odd represents the odd number closest to c. In addition, γ is taken as 2 and b is taken as 1. Through the above formula, the algorithm can automatically specify the length of the one-dimensional convolution kernel for the channel attention mechanism to avoid introducing additional hyperparameters.

[0105] Figure 5 This is a schematic diagram of the spatial attention module. The operation process of the spatial attention module is as follows Figure 5 As shown. Spatial Attention M s The calculation process of (F) can be expressed as formula (5):

[0106] M s (F) = σ(f 7×7 ([Avgpool(F);Maxpool(F)]))(5)

[0107] Among them, F is the input feature map, σ represents the Sigmoid activation function, and f 7×7 Indicates convolution calculation with a convolution kernel size of 7×7. Avgpool(F) is average pooling. Maxpool(F) is maximum pooling.

[0108] The difference between ECBAM and CBAM lies in the feature extraction and operation method of the channel attention module: in the improved channel attention module of ECBAM, the input feature map F is subjected to average pooling and maximum pooling extraction feature operations, followed by convolution fusion, and then enters the Sigmoid activation function, and finally outputs the channel attention Mc. CBAM, on the other hand, sends the two feature maps obtained by maximum pooling and average pooling into the same MLP (multi-layer perceptron), then fuses the features output by the MLP and activates them through Sigmoid to obtain the final channel attention features. The advantage of the ECBAM attention module of this patent lies in the feature extraction operation in the channel attention module, which can realize deep feature extraction of feature maps, and its operation method is simpler and easier to implement than the MLP method.

[0109] Improved loss function: The YOLOv5s-SIT algorithm adopts the composite loss function ComLoss strategy, which can improve the performance and generalization ability of the model and avoid the shortcomings of poor learning ability and poor expression ability of the model when using a single loss function. Among them, the confidence loss is calculated using the binary cross entropy loss function; the category loss is calculated using the multi-classification cross entropy loss function; and the position loss is calculated using the CIoU loss function.

[0110] Among them, the mathematical principle of the binary cross entropy loss function is as follows (6).

[0111]

[0112] In the above formula, L(x,c) is the binary cross entropy function, x is the label of the category, is the true probability of the category label in the i-th row and j-th column, the positive category takes the value of 1, the negative category takes the value of 0, and c is the predicted probability of the corresponding category. is the predicted expectation of the current i-th category.

[0113] The mathematical principle of the multi-classification cross entropy loss function is as follows (7).

[0114]

[0115] In the above formula, L(x,c) is the multi-classification cross entropy function, x is the label of the category, is the true probability of the m category in the i-th row and j-th column, the positive category takes the value of 1, the negative category takes the value of 0, and c is the predicted probability of the corresponding category. is the predicted expectation of the current i-th category.

[0116] The CIoU loss function has higher efficiency and better convergence when used to calculate the anchor box position loss. Its mathematical principle is as follows:

[0117]

[0118] In the formula, To predict the distance between the anchor box and the centroid of the true value box, is the minimum diagonal distance between the predicted anchor box and the true value box, IoU is the intersection over union of the bounding boxes of the predicted anchor box and the true value box, α is the adjustment coefficient, v is the aspect ratio of the bounding box, and as mentioned above, the width and height of the target true value box are w and h respectively.

[0119] Hyperspace Pyramid Pooling Module

[0120] The detection performance of the YOLOv5s algorithm for targets of extreme sizes (too small and too large) is relatively poor. The reason is that the algorithm has a relatively fixed perception of the feature layer of the image to be detected. Therefore, it is proposed to upgrade the spp module in BackBone (backbone network) to Spp_x (hyperspace pyramid pooling module), fuse local features with global features, and increase the receptive field size of the corresponding feature layer in the algorithm, so that the YOLOv5s-SIT algorithm can take into account small targets at different network layers. In view of the fact that the present invention is mainly aimed at the detection of small infrared targets, the Spp_x structure is set to a multi-layer Max pool cascade, and its structure is as follows Figure 6 shown.

[0121] The difference between the Spp_x structure and the SPP structure is that the Spp_x structure performs maximum pooling on pooling kernels of different sizes to better extract deep features. The main upgrade is to replace the Max pool of the original SPP structure with Max pool 2D, and cascade the three Max pool 2Ds with each other, which can ultimately improve the receptive field of the algorithm network and focus on and extract more features.

[0122] Example 2

[0123] The detection anchor frame is used to frame the detected small infrared target. The parameter setting principle of the detection anchor frame is as follows:

[0124] The present invention uses the K-means clustering method to calculate the optimal anchor frame suitable for the infrared small target detection task. From a popular understanding, we know that the length and width of the truth value frame that has been annotated in all training sets, and the anchor frame is a potential candidate frame for predicting these truth value frames, so the length and width of the anchor frame should be as close to the truth value frame as possible. (The truth value frame in the training data set is already annotated. The algorithm first trains the training data set multiple times to achieve the trained anchor frame and the truth value frame are close, which means that the algorithm model has reached the level of being able to detect infrared small targets. Afterwards, the algorithm model is applied to the test data set without annotated truth value frames to detect infrared small targets, and the detected infrared small targets are framed with the trained anchor frame.) Since the prediction layer of the YOLO network contains 3 scale information (corresponding to 3 receptive fields, that is, the above-mentioned three target sizes of 19×19, 38×38 and 76×76), and there are three anchor frames of each scale, we need 9 scale anchor frames, that is, we need to cluster all truth value frames into 9 categories. First, the width and height of the target truth box are set to w and h respectively, and then the K-means algorithm is used to calculate And w, h cluster the training data set respectively. The specific process is as follows:

[0125] Step 1: First select 9 target truth boxes as the initial category centroids;

[0126] Step 2: Calculate the distance between each target truth box and the centroid of the nine categories, and then assign each truth box to the category with the closest distance;

[0127] Step 3: Calculate the centroid of each category based on the latest category assignment results;

[0128] Step 4: If the termination condition is met, the algorithm is completed, otherwise it returns to the second step to continue execution.

[0129] Termination conditions: the number of reallocated truth boxes is less than the minimum value; the offset of the K-means center is less than the minimum value; the sum of squared errors between the truth box and the category centroid reaches a local minimum.

[0130] Among them, the centroid distance calculation method in the clustering process adopts L 1 The distance measurement principle, expression (1) is as follows:

[0131]

[0132] In the formula, α i represents the i-th element, is the centroid element in the current iteration process, and k is the number of iterations.

[0133] The clustering process will produce 3 The centroid of and 3 centroids of w / h, therefore, according to this clustering result, 9 anchor frames of different scales and aspect ratios can be generated as the anchor frames of the present invention. Among them, represents the basic scale of the anchor box, and w / h represents the aspect ratio of the anchor box. When calculating the specific anchor box of a certain input feature map, let ah ij and aw ij are the height and width of the jth anchor box of the ith feature map, respectively. Then ah ij and aw ij It can be expressed by formula (2):

[0134]

[0135] In the formula, ρ j Indicates the jth w / h; basesize i Indicates the i-th

[0136] Example 3

[0137] The present invention designs a lightweight method for an infrared small target detection system based on an improved YOLOv5s. Under the premise of substantially not reducing the detection accuracy and speed, the redundant part of the algorithm network is eliminated by appropriately reducing the number of convolution kernels in the convolution layer of the feature extraction part in the YOLOv5s-SIT algorithm, thereby accelerating the inference speed of the algorithm network.

[0138] Since the feature extraction layer of the YOLOv5s-SIT algorithm detection network designed by the present invention is a feature pyramid structure, the structure extracts the features of the image to be inspected through an upsampling operation. In order to ensure the integrity of the feature extraction, the upsampling layer is not pruned.

[0139] The process of the lightweight algorithm is as follows:

[0140] First, the present invention introduces the parameter γ into the loss function of sparse training to perform sparse training on the YOLOv5s-SIT algorithm. The loss function in sparse training is shown in formula (9).

[0141] L sp =∑ (x,y) l(f(x,W),y)+λ∑ γ∈Γ M(γ) (9)

[0142] where Σ (x,y) l(f(x,W),y) is the original composite loss function of the algorithm, L sp is the loss function of the algorithm's sparse training; f(x,W) is the output of the algorithm's convolutional layer; λ is the sparse parameter, which represents the intensity of the sparse training; Σ γ∈ΓM(γ) represents the L of all BN layer parameters γ 1 norm, where M(γ) = |γ|, that is, taking L 1 Regularization can form the sparsification parameter γ.

[0143] The goal of the sparse training YOLOv5s-SIT algorithm is to make the loss function L sp As the training progresses, the value of the BN layer parameter γ gradually approaches 0, and its sparsification strength increases; when most of the parameters γ approach 0, the training is stopped. At this time, most of the redundant BN layer parameters γ are near 0, which means that the corresponding convolution channels cannot extract important features. At this time, the YOLOv5s-SIT algorithm can be significantly pruned according to the trained γ value to achieve the purpose of lightweight algorithm. Pruning strategy is as follows Figure 6 shown.

[0144] Figure 7 Diagram of pruning strategy

[0145] Afterwards, the present invention adopts a pruning strategy for the number of convolution kernels in the convolution layer of the feature extraction part Neck of the backbone network. In the actual channel pruning operation, binary coding is used to represent the pruning operation, and the union of the coding set is taken as the pruned channel. The specific example is shown in the figure above. The final pruning result needs to ensure that the dimensions of the output features of CBL_0, CBL_2 and CBL_4 remain consistent. The CBL in Neck includes CBL_0, CBL_2, CBL_3 and CBL_4, Add1 and Add2

[0146] The steps are as follows:

[0147] (1) The quantization code of a CBL module can be expressed as {0, 0, 1, 0, 1, 1}, where 1 indicates the channel to be retained according to the γ parameter obtained by BN sparsification training, and 0 indicates the channel to be pruned during the pruning process.

[0148] (2) In actual operation, the binary code of each CBL is first determined according to the γ parameter obtained by BN sparsification training. For example, the codes of CBL_0, CBL_2 and CBL_4 are {1,0,1,1,0,0}, {1,0,0,1,0,1} and {0,0,1,1,0,1} respectively.

[0149] (3) By performing an intersection operation on the codes in the above example, we can obtain the actual pruned code as {1,0,1,1,0,1}.

[0150] (4) According to the pruning code calculated in (3), the second and fifth convolution kernels in CBL_0, CBL_2 and CBL_4 are pruned and deleted.

[0151] The above pruning strategy can ensure the model compression rate while greatly maintaining the detection accuracy of the algorithm model.

[0152] The algorithm principles of embodiments 1, 2 and 3 of the present invention are as follows Figure 8 shown.

[0153] Embodiment 4:

[0154] The experimental platform of the present invention runs on the Windows 10 operating system, uses two GTX TITAN XGPU graphics cards for algorithm model training, CUDA version 10.1, CPU model i7-13700F, and the program is written based on the PyTorch1.4 deep learning framework. The batch size is set to 16 to maximize the use of GPU video memory. The Adam optimizer is used for training optimization, and the weight decay coefficient is 0.0005. The initial learning rate is set to 0.001, and then automatically adjusted according to the gradient. When the epoch increases, the learning rate is appropriately reduced to make the algorithm model converge accurately. A total of 200 epochs are trained.

[0155] Ablation experiment verification

[0156] ECBAM Attention Module Verification

[0157] The ECBAM attention module is added to YOLOv5s to verify its improvement on the detection effect of the YOLOv5s algorithm. The experimental results are shown in Table 1.

[0158] Table 1 Comparative experimental results of ECBAM attention module

[0159] algorithm mAP0.5 mAP0.5:0.95 YOLOv5s 0.8857 0.6835 YOLOv5s-ECBAM 0.9086 0.6987

[0160] As can be seen from Table 1, the use of the ECBAM attention module effectively improves the detection accuracy compared to the original YOLOv5s. The evaluation indicators mAP0.5 and mAP0.5:0.95 are improved by 0.0229 and 0.0152 respectively.

[0161] Hyperspace pyramid pooling module spp_x verification

[0162] A hyperspatial pyramid pooling module Spp_x is added separately to YOLOv5s to verify its improvement on the detection effect of the YOLOv5s algorithm.

[0163] Table 2 Comparative experimental results of hyperspace pyramid pooling module Spp_x

[0164] algorithm mAP0.5 mAP0.5:0.95 YOLOv5s 0.8857 0.6835 YOLOv5s-Spp_x 0.9091 0.7012

[0165] As can be seen from Table 2, the improvement of the original YOLOv5s using the hyperspatial pyramid pooling module Spp_x effectively improves the detection accuracy. The evaluation indicators mAP0.5 and mAPO.5:0.95 are improved by 0.0240 and 0.0177 respectively.

[0166] Composite loss function ComLoss comparison and verification

[0167] The improved loss function ComLoss is added to YOLOv5s to verify its improvement on the detection effect of the YOLOv5s algorithm. The experimental results are shown in Table 3.

[0168] Table 3. Comparative experimental results of composite loss function

[0169] algorithm mAP0.5 mAP0.5:0.95 YOLOv5s 0.8857 0.6835 YOLOv5s-ComLoss 0.8986 0.6903

[0170] As can be seen from Table 3, the use of the improved composite loss function effectively improves the detection accuracy of the original YOLOv5s. The evaluation indicators mAP0.5 and mAPO.5:0.95 are improved by 0.0129 and 0.0068 respectively.

[0171] Model lightweight experiment

[0172] During the experiment, the algorithm model was first trained to sparse the BN layer, then the channels of the algorithm model were pruned according to the pruning rule, and finally the lightweight algorithm model was trained with the infrared data set. The overall pruning rate was set, and then the adaptive pruning rate was optimized to screen each convolutional layer.

[0173] Model lightweight experiment

[0174] First, all the γ parameters in the algorithm model are obtained and sorted. The principle is as follows:

[0175] T=sort([γ 1 ,γ 2 ,γ 3 ,…,γ N ])(10)

[0176] will be less than γ 0.8N The parameters with a value greater than 0.8 are selected and their weights are set to 0, and the parameters with a value greater than 0.8 are selected and their weights are set to 1. The comparison of the pruning experiment results is shown in Table 4.

[0177] Table 4 Comparison of pruning experiment results

[0178] Pruning rate Parameter quantity (millions) Model size (MB) FPS 0 67.5 28.7 42.93 20％ 51.4 23.1 50.97 40％ 35.2 17.4 59.68 80％ 4.3 6.8 71.17

[0179] From the analysis of Table 4, we can see that pruning the model can effectively reduce the number of parameters and size of the algorithm model, and improve the running speed. Compared with the original algorithm model, when the pruning rate reaches 80%, the number of parameters of the lightweight algorithm model is reduced from 67.5 million to 4.3 million, the size is reduced from 28.7M to 6.8M, and the running speed is increased from 42.93 frames per second to 71.17 frames per second, achieving the design goal of the lightweight algorithm model.

[0180] Analysis of measurement accuracy changes

[0181] The accuracy change of the lightweight algorithm model was analyzed using mAP0.5 and mAPO.5:0.95 as indicators. The experimental results are shown in Table 5.

[0182] Table 5 Comparison of detection accuracy before and after pruning

[0183] Pruning rate mAP0.5 mAP0.5:0.95 0 0.9234 0.7124 20％ 0.9172 0.7101 40％ 0.9103 0.7059 80％ 0.8971 0.6922

[0184] The experimental results in Table 7 show that with the continuous increase of the pruning rate, the pruned algorithm model has a small decrease in the two indicators of mAPO.5 and mAPO.5:0.95, but the detection accuracy of the algorithm model can still be guaranteed in practical applications, which effectively proves the feasibility of lightweighting the algorithm model. In actual use, the algorithm model can be selectively lightweighted according to the task attributes and the performance of the platform on which the algorithm model is carried, so as to better achieve overall consideration and balance.

[0185] The experimental results generally show that the designed lightweight infrared small target detection network and method based on improved YOLOv5s has the advantages of small size, accurate detection and fast speed, and can be better adapted to low computing power, low memory and low-cost hardware computing platforms.

[0186] In summary, the above are only preferred embodiments of the present invention and are not intended to limit the protection scope of the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principles of the present invention should be included in the protection scope of the present invention.

Claims

1. A lightweight infrared small target detection system based on improved YOLOv5s. The YOLOv5s network includes a backbone network BackBone, an enhanced backbone feature extraction network Neck, and a target detection head network Head. The characteristics are: The Spp module in the backbone network Backbone part of the YOLOv5s network is replaced with the Spp_x module; the Backbone part receives the input image, and the Spp-x module outputs the solved image data stream to the first CSP2-1; The enhanced backbone feature extraction network Neck includes a first CSP2-1, a first CBL, a first Upsample, a first Contcat, a second CSP2-1, a first ECBAM, a second CBL, a second Upsample, a second Contcat, a third CSP2-1, a third Contcat, a third CBL, a second ECBAM, a fourth CSP2-1, a fourth Contcat, a fourth CBL, a third ECBAM and a fifth CSP2-1 connected in sequence; the first CBL is connected to the fourth Contcat, and the second CBL is connected to the third Contcat; The target detection head network Head includes three Conv blocks, wherein the processing results of the third to fifth CSP2-1 in Neck are respectively connected to the three Conv blocks, and the three Conv blocks respectively output images of three detected target sizes; The Spp_x module includes an eighth CBL block and three Max pool 2Ds connected in sequence, the outputs of the CBL block and the three Max pool 2Ds are all connected to a Concat, the output of the Concat is connected to the ninth CBL block, and the output of the ninth CBL block is used as the output of the Spp_x module; The structures of the first to third ECBAMs are the same, specifically including an improved channel attention module and a spatial attention module; Input feature map F into the improved channel attention module to extract channel attention M c (F), and then fused with the input feature map F to obtain the feature map F'. The input feature map enters the spatial attention module to extract the spatial attention M s (F), and then fused with the feature map F' to obtain the output feature map F".

2. The lightweight infrared small target detection system based on improved YOLOv5s as claimed in claim 1, characterized in that: The backbone network Backbone part includes an input module, a Focus, a fifth CBL, a CSP1-1, a sixth CBL, a CSP1-2, a seventh CBL, a CSP1-3, an eighth CBL and an Spp-x module connected in sequence; The processing result of CSP1-2 module is connected to the second Concat in Neck part; The processing results of the CSP1-3 modules are connected to the first Concat in the Neck part.

3. The lightweight infrared small target detection system based on improved YOLOv5s as claimed in claim 1, characterized in that: The operation process of the improved channel attention module is specifically as follows: The input convolution feature map is F C×H×W , then the channel attention M c (F) can be expressed as formula (3): M c (F)=σ(Conv k (Avgpool(F))+Conv k (Maxpool(F))) (3) Among them, σ represents the Sigmoid activation function, Conv k represents a one-dimensional convolution kernel of length k, Avgpool(F) is average pooling, and Maxpool(F) is maximum pooling; The formula (4) for calculating the convolution kernel size k based on the number of feature map channels c is as follows: Among them, odd represents the odd number closest to c, γ is 2, and b is 1. Through the above formula (4), the length of the one-dimensional convolution kernel is automatically specified for the channel attention mechanism to avoid introducing additional hyperparameters. The operation process of the spatial attention module is: Spatial AttentionM s The calculation process of (F) is expressed as formula (5): M s (F)=σ(f 7×7 ([Avgpool(F);Maxpool(F)])) (5) Among them, F is the input feature map, σ represents the Sigmoid activation function, and f 7×7 It represents the convolution calculation with a convolution kernel size of 7×7, Avgpool(F) is the average pooling, and Maxpool(F) is the maximum pooling.

4. The lightweight infrared small target detection system based on improved YOLOv5s as claimed in claim 1, characterized in that: The output of the target detection head network Head contains three scale information, corresponding to three target sizes 19×19, 38×38 and 76×76.

5. A lightweight infrared small target detection method based on improved YOLOv5s, characterized in that: According to the target size detected by the system described in claims 1 to 3, the target calibration is performed in the following manner: The output of the target detection head network Head contains three scale information, each scale includes three anchor frames, that is, nine scale anchor frames are required. Then the detection method clusters all the true value frames into nine categories. Set the width and height of the target truth box to w and h respectively, and then use the K-means algorithm to and w and h to cluster the training data sets respectively; The specific process of this detection method is as follows: Step 1: First select 9 target truth boxes as the initial category centroids; Step 2: Calculate the distance between each target truth box and the centroid of the nine categories, and then assign each truth box to the category with the closest distance; Step 3: Calculate the centroid of each category based on the latest category assignment results; Step 4: If the termination condition is met, the algorithm is completed, otherwise it returns to step 2 to continue execution; Termination conditions: the number of real-value boxes that are reallocated is less than the minimum value; the offset of the K-means center is less than the minimum value; The sum of squared errors between the true value box and the category centroid reaches a local minimum; The centroid distance calculation method in the clustering process adopts the L1 distance measurement principle, and its expression (1) is as follows: Among them, α i represents the i-th element, is the centroid element in the current iteration process, and k is the number of iterations; The clustering process will produce 3 The centroid of and 3 centroids of w / h are used to generate 9 anchor boxes of different scales and aspect ratios according to the clustering results as the anchor boxes of this method; among them, represents the basic scale of the anchor box, and w / h represents the aspect ratio of the anchor box; When calculating the specific anchor frame of a certain input feature map, let ah ij and aw ij are the height and width of the jth anchor box of the ith feature map, respectively. Then ah ij and aw ij It is expressed by formula (2): Among them, ρ j Indicates the jth w / h; basesize i Indicates the i-th 6. The lightweight infrared small target detection method based on improved YOLOv5s as claimed in claim 5, characterized in that: The YOLOv5s-SIT algorithm adopts the composite loss function ComLoss strategy, in which the confidence loss is calculated using the binary cross entropy loss function; the category loss is calculated using the multi-classification cross entropy loss function; and the position loss is calculated using the CIoU loss function.

7. The lightweight infrared small target detection method based on improved YOLOv5s as claimed in claim 6, characterized in that: The mathematical principle of the binary cross entropy loss function is as follows (6): Among them, L(x,c) is the binary cross entropy function, x is the label of the category, is the true probability of the category label in the i-th row and j-th column, the positive category takes the value of 1, the negative category takes the value of 0, and c is the predicted probability of the corresponding category. is the predicted expectation of the current i-th category; The mathematical principle of the multi-classification cross entropy loss function is as follows (7): Among them, L(x,c) is the multi-classification cross entropy function, x is the label of the category, is the true probability of the m category in the i-th row and j-th column, the positive category takes the value of 1, the negative category takes the value of 0, and c is the predicted probability of the corresponding category. is the predicted expectation of the current i-th category; The CIoU loss function is used to calculate the anchor box position loss, and its mathematical principle is as follows (8). Among them, To predict the distance between the anchor box and the centroid of the true value box, is the minimum diagonal distance between the predicted anchor box and the true value box, IoU is the intersection over union of the bounding boxes of the predicted anchor box and the true value box, α is the adjustment coefficient, v is the aspect ratio of the bounding box, and as mentioned above, the width and height of the target true value box are w and h respectively.

8. A lightweight method for infrared small target detection system based on improved YOLOv5s, characterized in that: For the system as claimed in claim 1, the following steps are used to perform pruning during the training process to achieve lightweight: First, the YOLOv5s-SIT algorithm is sparsely trained by introducing the parameter γ in the loss function of sparse training. The loss function in sparse training is shown in formula (9): L sp =∑ (x,y) l(f(x,W),y)+λ∑ γ∈Γ M(γ) (9) where ∑ (x,y) l(f(x,W),y) is the original composite loss function of the algorithm, L sp is the loss function of the algorithm's sparse training; f(x,W) is the output of the algorithm's convolutional layer; λ is the sparse parameter, which represents the intensity of the sparse training; ∑ γ∈Γ M(γ) represents the L1 norm of all BN layer parameters γ, where M(γ) = |γ|, that is, L1 regularization is used to form the sparse parameter γ; The goal of sparse training of the system as claimed in claim 1 is to make the loss function L sp Minimum, when the parameter γ approaches 0, stop training; Afterwards, a pruning strategy is adopted for the number of convolution kernels in the convolution layer of Neck, the feature extraction part of the backbone network.

Images