Single-frame infrared weak and small target detection method based on multi-scale feature enhancement

Abstract

The invention discloses a single-frame infrared weak and small target detection method based on multi-scale feature enhancement, and the method comprises the following steps: 1, obtaining a single-frame infrared image, constructing a single-frame infrared image data set, and dividing the single-frame infrared image data set into a training set and a test set; 2, an infrared weak and small target detection network model is constructed, the infrared weak and small target detection network takes a YOLOv5 network model as a basic network, and the YOLOv5 network model comprises a trunk feature extraction network, a feature fusion network and a detection head; a multi-scale feature enhancement module is introduced into the trunk feature extraction network, and the multi-scale feature enhancement module comprises a context feature fusion module and a feature detail enhancement module which are sequentially connected in series; 3, designing a training strategy, and training the infrared weak and small target detection network model by taking the training set as input according to the training strategy; and step 4, after training is completed, inputting a test set to test and evaluate the infrared weak and small target detection network model.

Description

Technical Field

[0001] The present invention relates to the field of computer vision technology, and mainly to the problem of small target detection and recognition in infrared imaging, and specifically to a single-frame infrared small target detection method based on multi-scale feature enhancement. Background Art

[0002] Infrared imaging technology has shown its importance in many fields with its unique advantages. Compared with visible light imaging systems, infrared imaging systems are small in size, light in weight, and not easily affected by harsh environments such as smoke and lighting conditions. They can work all day and all day. By passively receiving target thermal radiation, they show high concealment and anti-interference capabilities. Therefore, this technology is widely used in civil and military fields, such as autonomous driving, early warning systems, maritime monitoring, precision guidance, etc. However, there are still many problems and difficulties in the detection of infrared weak targets. First, since the target is usually photographed at high altitude or long distance, it is far away from the detector and the image resolution is low, resulting in a lack of detailed information such as texture and shape of the target; secondly, detection under a large field of view will introduce complex background interference and noise, resulting in a low signal-to-noise ratio of the target, which is easily submerged in a complex background, increasing the difficulty of detecting weak targets. Therefore, accurate detection of infrared weak targets remains a challenging problem.

[0003] At present, infrared small target detection methods are divided into traditional image processing methods and deep learning methods. Common traditional methods include filter-based, visual saliency-based, and image data structure-based methods. Traditional image processing detection methods are affected by artificial feature extraction and have low robustness. When faced with complex backgrounds and target shape changes, they have a high false alarm rate. In recent years, with the rapid development of deep learning in many fields, compared with traditional detection methods, it adopts a data-driven end-to-end learning framework with a flexible network structure and can adaptively learn target features without the need for artificial feature extraction. It is a robust and accurate target detection method, and is therefore more suitable for infrared small target detection under complex backgrounds.

[0004] The infrared small target detection method based on deep learning is modeled into two tasks: segmentation and detection. Most of the current research is based on the segmentation task, which provides detailed information at the pixel level and requires a lot of computing resources. This is beneficial in scenarios that require precise distinction. The ultimate goal of infrared small target detection is tracking and positioning. Segmentation is an intermediate representation as the input of tracking and positioning. The integrity of segmentation is only an approximation of detection accuracy and cannot evaluate the specific detection performance. Therefore, in order to reduce the utilization of computing resources, it is of great significance to model infrared small target detection as a target detection task. Summary of the invention

[0005] In view of the shortcomings of the prior art, the present invention proposes a single-frame infrared dim small target detection method based on multi-scale feature enhancement.

[0006] In order to solve the above technical problems, the technical solution of the present invention is:

[0007] A single-frame infrared dim small target detection method based on multi-scale feature enhancement comprises the following steps:

[0008] Step 1: Obtain a single-frame infrared image and construct a single-frame infrared image dataset.

[0009] Step 2: Build an infrared small target detection network model. The infrared small target detection network model is based on the YOLOv5 network model. The YOLOv5 network model includes a backbone feature extraction network, a feature fusion network and a detection head.

[0010] Step 3: A multi-scale feature enhancement module (MFEB, Muti-FeatureEnhance Block) is introduced into the backbone feature extraction network. The module includes a context feature fusion module (CFB, Context Fusion Block) and a feature detail enhancement module (DEB, Detail Enhance Block). The context feature fusion module includes a patch-guided context aware block (PGCAB, Patch-Guide Context Aware Block) and a sequence convolution module; the feature detail enhancement module includes a detail enhancement convolution module DEConv and a CBAM attention mechanism.

[0011] Step 4: Design a training strategy. According to the number of training epochs, adaptively use Focal-NWD and Focal-αECIoU loss functions as the positioning loss for calculation, and achieve coarse-to-fine regression of the bounding box.

[0012] Step 5: Input the training samples obtained in step 1 into the infrared small target detection model constructed in steps 2 to 4 for training.

[0013] Step 6: Input the test set to test and evaluate the model.

[0014] Preferably, the multi-scale feature enhancement module MFEB in step 3 specifically includes:

[0015] Context feature fusion module CFB: The context feature fusion module contains a block-guided context information perception module and a sequence convolution module. Given a feature The average pooling operation, the maximum pooling operation and the Conv operation are performed respectively to extract the average pooling feature F within each block. avg, the maximum pooling feature F within the block max And the local feature F within the block local , and then respectively avg and F max Obtain inter-block global features with long-range interactions through the Nonlocal module and The two global features effectively capture contextual information from different angles, and then the two global features are concatenated to obtain the aggregated global feature F global , F global and F local The corresponding elements are multiplied to obtain the global and local feature fusion feature F fusion Finally, the original input feature F is connected with the fusion feature F through residual connection fusion The final feature F′ is obtained by adding them together. Different scales of features can be obtained by using different k values.

[0016] Preferably, the block-guided context information perception module first blocks the feature F with a window size of k×k and a step size of s, and the feature is represented as F k ′.

[0017] Preferably, the sequence convolution module is a Bottleneck module, which is composed of a 1×1 Conv module, a 3×3 Conv module, and a residual connection. The feature F obtained after passing through the sequence convolution module is F seq Finally, the block-guided context-aware module outputs the feature F k ′ and the sequence convolution module output feature F seq Add to get the output features of the context feature fusion module

[0018] Preferably, the feature detail enhancement module DEB: the feature detail enhancement module includes a detail enhancement convolution module DEConv and a CBAM attention mechanism.

[0019] The detail enhancement convolution module DEConv includes standard convolution and differential convolution. Differential convolution uses a pixel-pair differential calculation strategy to explicitly encode prior knowledge into the model. Differential convolution includes four convolution operations: center differential convolution CDC, angle differential convolution ADC, vertical differential convolution VDC, and horizontal differential convolution HDC. The output of the context feature fusion module As input, it undergoes five convolutions in parallel, and finally the outputs of the five branches are element-wise added to obtain the output of the detail enhancement convolution module.

[0020] As a preference, the CBAM attention mechanism is divided into two parts: spatial attention and channel attention. Given a feature feature First, channel attention is used, and channel attention is respectively global average pooling and global maximum pooling in the spatial dimension to obtain the global feature description and Describing global features and The two channel attention vectors are added together and normalized by the sigmoid function to generate the final channel attention weight. The characteristics and channel attention weight M C The main channel is multiplied to obtain the channel-enhanced feature map Channel enhanced feature map After spatial attention, the feature map of channel enhancement is Perform channel average pooling and channel maximum pooling in the channel dimension to obtain global feature description and Will and After concatenation in the channel dimension, a 7×7 convolution kernel is used for convolution and sigmoid normalization is performed to obtain a spatial attention map. Finally, the channel enhanced feature map With the spatial attention map M S Element-wise multiplication obtains the spatially enhanced feature F″.

[0021] The context feature fusion module and the feature detail enhancement module are connected in series in sequence to obtain the multi-scale feature enhancement module MFEB.

[0022] Preferably, step 4 specifically includes the following steps:

[0023] Define the center coordinates of the predicted box as cx and cy, and the width and height as w and h; the center coordinates of the real box as cx gt andcy gt , with width and height w gt and h gt .

[0024] Construct the Focal-αECIoU loss function, the specific formula is as follows:

[0025]

[0026] L αECIoU =1-αECIoU

[0027] L Focal-αECIoU =IoU γ L αECIoU

[0028] Among them, IoU is the intersection-over-union ratio, which is the ratio of the intersection and union between the model's predicted box and the true box. 2 (b,b gt ), ρ 2 (w,w gt ), ρ 2 (h,h gt ) calculates the center point distance, width distance and length distance of the predicted box and the real box respectively. β is a balance parameter, ν is a parameter used to measure the consistency of aspect ratio, c is the diagonal length of the two bounding boxes, C w and C h are the width and height of the minimum bounding rectangle of the predicted box and the true box, respectively. α is a power parameter used to meet different levels of bounding box regression accuracy and reduce the sensitivity to the data set. γ is a parameter that controls the degree of outlier suppression.

[0029] The specific formula of Focal-NWD loss function is as follows:

[0030]

[0031] L Focal-NWD =NWD γ (1-NWD)

[0032] Where N and N gt Represents the Gaussian distribution of the predicted box and the true box. Indicates N and N gt Wasserstein distance between Gaussian distributions. NWD(N,N g ) represents the exponential normalized form of the Wasserstein distance between two Gaussian distributions, and c is a constant.

[0033] As a preferred strategy, the new training strategy uses the Focal-NWD loss function to perform a rough bounding box regression in the first half of the training rounds according to the set number of epochs, and the Focal-αECIoU loss function is used for a fine bounding box regression in the second half of the rounds. The positioning loss function formula is as follows:

[0034]

[0035] Among them, nepoch is the number of current training iterations, and epoch is the total number of training rounds.

[0036] The present invention has the following characteristics and beneficial effects:

[0037] 1. The context feature fusion module adopts a multi-branch parallel feature extraction strategy. The block-guided context information perception module and the sequence convolution module form multiple branches, thereby extracting different features from multiple scales and levels, obtaining the global and local features of the target, and making the features contain richer semantic information. On the one hand, it can reduce the loss of small target information and improve the detection accuracy of small targets. On the other hand, the background block and the target block may have similar responses in the local features, and the global correlation information can be effectively used to determine the position of the target area, excluding the area similar to the target and the noise interference, thereby eliminating the false alarm of small targets.

[0038] 2. The feature detail enhancement module can enhance the details of the features, making the model pay more attention to the key areas, retaining the key information while suppressing noise and irrelevant information, enhancing the feature representation, and improving the perception ability of the model, thereby improving the detection accuracy of small targets.

[0039] 3. As part of the positioning loss function, NWD reduces the sensitivity of the IoU metric to the position deviation of weak targets during training. αECIoU introduces the length and width loss factor between the predicted box and the real box, which alleviates the problem that the length and width cannot be fitted if the penalty term of CIoU is always 0 when the aspect ratio of the predicted box and the real box is the same. At the same time, a power parameter α is introduced. By adjusting α, we focus on high-quality anchor boxes to obtain higher regression accuracy. In the early stage of training, CIoU leads to accelerated convergence. When the penalty term of CIoU gradually fails, the part of calculating the width and height loss begins to take effect, fine-tuning the regression and effectively improving the regression accuracy.

[0040] 4. The Focal term is introduced into the two loss functions to effectively separate high-quality anchor boxes from low-quality anchor boxes, increase the loss weight of high-quality anchor boxes, and thus improve the regression accuracy.

[0041] 5. Adaptively adjust the timing of using the two loss functions, Focal-NWD and Focal-αECIoU, according to the number of epochs to achieve a training process from coarse to fine regression accuracy of the bounding box, which speeds up the convergence speed and improves the regression accuracy. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative labor.

[0043] Figure 1It is a flowchart of the steps of a single-frame infrared dim small target detection method based on multi-scale feature enhancement in an embodiment;

[0044] Figure 2 It is a framework diagram of the improved YOLOv5 model constructed in the embodiment;

[0045] Figure 3 Schematic diagram of the block-guided contextual information awareness module PGCAB;

[0046] Figure 4 This is a schematic diagram of the Nonlocal module;

[0047] Figure 5 Schematic diagram of feature enhancement convolution DEConv;

[0048] Figure 6 Schematic diagram of the CBAM attention module;

[0049] Figure 7 Schematic diagram of a multi-scale feature enhancement module MFEB in an embodiment;

[0050] Figure 8 It is a comparison chart of the detection results of various schemes in the embodiment. DETAILED DESCRIPTION

[0051] It should be noted that, in the absence of conflict, the embodiments of the present invention and the features in the embodiments may be combined with each other.

[0052] In the description of the present invention, it should be understood that the terms "center", "longitudinal", "lateral", "up", "down", "front", "back", "left", "right", "vertical", "horizontal", "top", "bottom", "inside", "outside" and the like indicate positions or positional relationships based on the positions or positional relationships shown in the accompanying drawings, and are only for the convenience of describing the present invention and simplifying the description, rather than indicating or implying that the device or element referred to must have a specific orientation, be constructed and operated in a specific orientation, and therefore cannot be understood as limiting the present invention. In addition, the terms "first", "second", and the like are only used for descriptive purposes, and cannot be understood as indicating or implying relative importance or implicitly indicating the number of technical features indicated. Thus, features defined as "first", "second", and the like may explicitly or implicitly include one or more of the features. In the description of the present invention, unless otherwise specified, "multiple" means two or more.

[0053] In the description of the present invention, it should be noted that, unless otherwise clearly specified and limited, the terms "installed", "connected", and "connected" should be understood in a broad sense, for example, it can be a fixed connection, a detachable connection, or an integral connection; it can be a mechanical connection or an electrical connection; it can be a direct connection, or it can be indirectly connected through an intermediate medium, or it can be the internal communication of two components. For ordinary technicians in this field, the specific meanings of the above terms in the present invention can be understood by specific circumstances.

[0054] The present invention provides a single-frame infrared dim small target detection method based on multi-scale feature enhancement, such as Figure 1 As shown, the following steps are included:

[0055] Step 1: Obtain a single-frame infrared image to construct a single-frame infrared image dataset, and divide it into a training set and a test set.

[0056] Specifically, in this embodiment, the following sub-steps are included:

[0057] Step 1.1, obtain the single-frame infrared image datasets SIRSTv2 and IRSTD1K and annotate the target real box so that the data format can be correctly read by the network.

[0058] Step 1.2: Construct a single-frame infrared image dataset, a validation set, and a test set. Divide the dataset constructed in step 1.1 into a ratio of 6:2:2. The training set accounts for 60% of the total dataset, and the validation set and test set each account for 20%.

[0059] Step 2: construct an infrared small target detection network model, wherein the infrared small target detection network is based on the YOLOv5 network model, and the YOLOv5 network model includes a backbone feature extraction network, a feature fusion network and a detection head;

[0060] A multi-scale feature enhancement module is introduced into the backbone feature extraction network, and the multi-scale feature enhancement module includes a context feature fusion module and a feature detail enhancement module which are connected in series in sequence.

[0061] It should be noted that, in this embodiment, the input data needs to be preprocessed, including adaptive anchor frame calculation and image scaling.

[0062] Specifically, in this embodiment, the backbone feature extraction network extracts features from the input data, and the backbone feature network includes 5 Conv modules, 2 C3 modules, 2 multi-scale feature enhancement modules MFEB and 1 SPPF module. The Conv module includes convolution, BN normalization function, and SiLU activation function connected in sequence. The C3 is divided into two branches, one of which includes 1 Conv module, and the other includes 1 Conv module and n Bottleneck modules. The bottleneck block of the Bottleneck module consists of a 1×1 Conv module and a 3×3 Conv module and a residual connection. The two branches are spliced ​​and fused on the channel. The C3 module used in the P4 and P5 layers in the original YOLOv5 network is replaced by the multi-scale feature enhancement module MFEB in this embodiment.

[0063] In a further configuration of this embodiment, the context feature fusion module includes a block-guided context information perception module PGCAB and a sequence convolution module. Figure 3 Given a feature The block-guided context information perception module first performs average pooling, maximum pooling and Conv operations on the feature F in blocks of k×k window size with a step size of s to extract the average pooling feature F within each block avg , the maximum pooling feature F within the block max And the local feature F within the block local , and then respectively avg and F max Obtain inter-block global features with long-range interactions through the Nonlocal module and The two global features effectively capture contextual information from different angles, and then the two global features are concatenated to obtain the aggregated global feature F global , F global and F local The corresponding elements are multiplied to obtain the global and local feature fusion feature F fusion Finally, the original input feature F is connected with the fusion feature F through residual connection fusion Add them together to get the final feature F′.

[0064] In this embodiment, the block-guided context information perception module first blocks the feature F with a window size of k×k and a step size of s. Therefore, the output feature of the block-guided context information perception module is recorded as F k ′.

[0065] Understandable, given features It is obtained by inputting single-frame infrared image data into the first few layers of the original YOLOv5 network model.

[0066] Further, such as Figure 4 As shown, the Nonlocal module is shown in Figure 4 As shown, a 1×1Conv module is first used to generate query (Q), key (K) and value (V) matrices respectively. The similarity between features is evaluated by calculating the dot product between Q and K. The value V matrix is ​​modulated with the similarity weight, and the softmax function is used to ensure that the sum of the weights is 1. The output is added to the features obtained by passing V through a 3×3 DWConv module to obtain the final features. The formula is as follows:

[0067]

[0068] where d k Indicates the dimension of the key vector.

[0069] The sequence convolution module is a Bottleneck module. The feature F obtained after passing through the sequence convolution module is F seq Finally, the block-guided context-aware module outputs the feature F k ′ and the sequence convolution module output feature F seq Add to get the output features of the context feature fusion module

[0070] It is understandable that different scales of features can be obtained by using different k values. In this embodiment, k=3 and k=5 are selected to obtain multi-scale features, wherein block features are obtained with a step size of s=1, and padding operations are performed to keep the size of the feature map unchanged.

[0071] In a further configuration of this embodiment, the feature detail enhancement module includes a detail enhancement convolution module DEConv and a CBAM attention mechanism.

[0072] Specifically, Figure 5 As shown in the figure, the DEConv module includes standard convolution and differential convolution, where the standard convolution is a Conv module that can capture local features. Differential convolution explicitly encodes prior knowledge into the model through a pixel-pair differential calculation strategy. Differential convolution includes four convolution operations: center differential convolution CDC, angle differential convolution ADC, vertical differential convolution VDC, and horizontal differential convolution HDC. Different differential convolutions can capture the local information differences in the input data from multiple angles, making the model pay more attention to the high-frequency information of the image, thereby enhancing the details of the features.

[0073] Further, such as Figure 6 As shown in Figure 2, the CBAM attention mechanism is divided into two parts: spatial attention and channel attention. Given a feature feature First, channel attention is used, and channel attention is respectively global average pooling and global maximum pooling in the spatial dimension to obtain the global feature description and Describing global features and The two channel attention vectors are added together and normalized by the sigmoid function to generate the final channel attention weight. Combine the original feature F with the channel attention weight M C The main channel is multiplied to obtain the channel-enhanced feature map Channel enhanced feature map After spatial attention, the feature map of channel enhancement is Perform channel average pooling and channel maximum pooling in the channel dimension to obtain global feature description and Will and After concatenation in the channel dimension, a 7×7 convolution kernel is used for convolution and sigmoid normalization is performed to obtain a spatial attention map. Finally, the channel enhanced feature map With the spatial attention map M S Multiply element by element to get the spatial enhanced feature F″. The specific formula is as follows:

[0074] F′=F⊙M C (F)=F⊙σ(f 1×1 (f 1×1 (AvgPool(F)))+f 1×1 (f 1×1 (MaxPool(F))))

[0075] F″=F′M S (F′)=σ(f 7×7 ([AvgPool(F′);MaxPool(F′)]))

[0076] CBAM(F)=F″

[0077] where f k×k (·) is the convolution operation with a kernel of size k×k.

[0078] In the above technical solution, the context feature fusion module and the feature detail enhancement module are sequentially connected in series to obtain a multi-scale feature enhancement module MFEB, see Figure 7 shown.

[0079] The feature fusion network fuses the features extracted by the backbone feature extraction network from bottom to top and from top to bottom. The feature fusion network includes 4 Conv modules, 4 C3 modules and 2 upsampling modules, and different features are fused using channel splicing operations.

[0080] The detection head detects the three features of different sizes obtained by the feature fusion network and outputs the target category and the bounding box of the target. The detection head includes three detection branches, and the three detection branches respectively include a 1×1 convolution operation.

[0081] Step 3: Design a training strategy and use the training set as input to train the infrared small target detection network model according to the training strategy.

[0082] It is understandable that small objects are more sensitive to the IoU metric, so according to the set number of epochs, the first half of the training rounds use the Focal-NWD loss function to perform a rough bounding box regression, and the second half of the rounds use the Focal-αECIoU loss function to perform a fine bounding box regression. This training strategy can effectively improve the regression accuracy of the bounding box. The formula is as follows:

[0083]

[0084] Among them, nepoch is the number of current training iterations, and epoch is the total number of training rounds.

[0085] Specifically, the following sub-steps are included:

[0086] Step 3.1: Define the center coordinates of the predicted box as cx and cy, with width and height as w and h; the center coordinates of the real box as cx gt andcy gt , with width and height w gt and h gt .

[0087] Step 3.2: Construct the Focal-αECIoU loss function. The specific formula is as follows:

[0088]

[0089]

[0090]

[0091]

[0092] L αECIoU =1-αECIoU

[0093] L Focal-αECIoU =IoUγ L αECIoU

[0094] Where AB is the intersection area of ​​the A bounding box and the B bounding box, AB is the union area, and ρ 2 (b,b gt ), ρ 2 (w,w gt ), ρ 2 (h,h gt ) calculates the center point distance, width distance and length distance of the predicted box and the real box respectively. β is a balance parameter, ν is a parameter used to measure the consistency of aspect ratio, c is the diagonal length of the two bounding boxes, C w and C h are the width and height of the minimum bounding rectangle of the predicted box and the true box, respectively. α is a power parameter used to meet different levels of bounding box regression accuracy and reduce the sensitivity to the data set. γ is a parameter that controls the degree of outlier suppression. In this embodiment, both α and γ are set to 2.

[0095] Step 3.3: Construct the Focal-NWD loss function. The specific derivation process and formula are as follows:

[0096] The bounding box is modeled as a two-dimensional Gaussian distribution to represent the importance of the target pixels within the box. Assume that the center coordinates of the bounding box are (c x ,c y ), with width and height w and h, the formula for the inscribed ellipse of the bounding box is as follows:

[0097]

[0098] The probability density function formula of the two-dimensional Gaussian distribution is as follows:

[0099]

[0100] Where x, μ and Σ represent the coordinates (x, y), mean vector and covariance matrix respectively. When (x-μ) T Σ -1 When (x-μ)=1, the inscribed ellipse of the bounding box can be modeled as a two-dimensional Gaussian distribution N(μ,Σ).

[0101]

[0102] The distance between two Gaussian distributions is used to measure the similarity between the predicted box and the true box. The Wasserstein distance is used to calculate the distance between Gaussian distributions. Assuming that the two Gaussian distributions are μ1=N(m1,Σ1) and μ2=N(m2,Σ2), the Wasserstein distance can be expressed as follows:

[0103]

[0104] Therefore, the two-dimensional Gaussian distribution distance between the predicted box and the true box can be expressed as the following formula:

[0105]

[0106] Normalize its index to the range of 0-1 to obtain the normalized Wasserstein distance NWD, the formula is as follows:

[0107]

[0108] Where N and N gt Represents the Gaussian distribution of the predicted box and the true box, and c is a constant.

[0109] The specific formula of Focal-NWD loss function is as follows:

[0110] L Focal-NWD =NWD γ (1-NWD)

[0111] Wherein γ is a parameter for controlling the degree of suppression of outliers. In this embodiment, γ is set to 2.

[0112] The further setting of this embodiment is to construct the training environment and training parameters of this embodiment during training. Use Ubuntu 18.04 operating system and Pytorch 13.1 deep learning framework, and use NVIDIA GeForce RTX3090 GPU for training. Set the number of training rounds epoch to 200, the batch size Batch-size to 16, use the stochastic gradient descent SGD optimizer, the learning rate lr is set to 0.001, the IoU threshold for NMS is set to 0.5, and the Conf confidence threshold is set to 0.001.

[0113] The training set and validation set in the dataset constructed in step 1 are input into the network. The input image size is set to 640×640. The training set is used to learn the network parameters, and the validation set is used to verify the training effect.

[0114] Step 4: After training is completed, input the test set to test and evaluate the infrared small target detection network model.

[0115] Specifically, the following sub-steps are included:

[0116] Step 4.1: Input the test set constructed in step 1 into the trained model, compare the prediction results and labels output by the model, and test the detection performance of the detection network model.

[0117] Step 4.2: Construct a calculation model performance index to evaluate the detection performance of the model. The index is constructed as follows:

[0118]

[0119]

[0120]

[0121]

[0122] Among them, TP represents the number of actual positive class predictions as positive class; FP represents the number of actual negative class predictions as positive class; FN represents the number of actual positive class predictions as negative class; P represents precision, R represents recall, F1 is the harmonic mean of precision and recall, and P(R) represents the Precision-Recall curve.

[0123] Step 4.3: To verify the effectiveness of each part of the model, this example conducts an ablation experiment on the improved part on the SIRSTv2 dataset, as shown in Table 1.

[0124] Table 1

[0125]

[0126] As can be seen from Table 1, Focal-αECIoU as a loss function can effectively improve the regression accuracy of the target bounding box, thereby improving the precision, but the recall rate is low and there are more missed detections. On this basis, the Focal-NWD loss function is introduced and a new training strategy is used to adaptively adjust the use of the loss function, which effectively improves the recall rate and reduces missed detections. Finally, the multi-scale feature enhancement module MFEB is used to obtain multi-scale features and context information, and perform multi-level enhancement on the features to improve the feature details, effectively improving the detection precision and recall rate, reducing false alarms and missed detections of the target, and greatly improving the detection performance.

[0127] Step 4.4: Since most of the current deep learning-based methods regard infrared small target detection as a pixel-level segmentation task. In order to verify the effectiveness of the improved model, this embodiment simultaneously selects a model based on a pixel-level segmentation task and a model based on a target detection task to compare the improved model, and conducts comparative experiments on the two data sets SIRSTv2 and IRSTD-1K. Among them, the SOTA models AGPCNet, ACM, ISNet, ALCNet, and DNANet in the infrared small target detection task are selected as the models based on the segmentation task; the single-stage target detection algorithms YOLOv5 to YOLOv11, RetinaNet, SSD and the two-stage target detection algorithms Faster-RCNN and Cascade-RCNN are selected as the target detection task models for comparative experiments. The test parameter confidence threshold is set to 0.001, and the IOU threshold is set to 0.5. In order to unify the indicators, this embodiment selects the commonly used evaluation indicators Precision, Recall, and F1 in the infrared weak small target detection task. The results of the comparative experiment are shown in Table 2:

[0128] Table 2

[0129]

[0130]

[0131] The above experimental results show that since the segmentation task model needs to obtain pixel-level segmentation results, slight deficiencies will cause false alarms or missed detections, resulting in poor detection effects at the target level. The model based on the target detection task proposed in the present invention is far superior to various SOTA segmentation task models and the currently commonly used single-stage target detection task models and two-stage target detection task models in the target-level detection tasks. The values ​​of the evaluation indicators Precision, Recall, and F1 of this method on the SIRSTv2 and IRSTD-1K datasets have reached the optimal value, indicating that this method effectively improves the detection performance of infrared weak targets. On the SIRSTv2 dataset, compared with the Baseline YOLOv5, the precision of this method is improved by 7.7%, the recall rate is improved by 6.8%, and the F1 is improved by 7.3%.

[0132] Step 4.4: Visualize the detection results on the test image, such as Figure 8 As shown, the solid circle represents the correctly detected target, the dotted circle represents the missed detection, and the ellipse represents the false alarm.

[0133] The embodiments of the present invention are described in detail above with reference to the accompanying drawings, but the present invention is not limited to the described embodiments. For those skilled in the art, various changes, modifications, substitutions and variations of these embodiments including components are made without departing from the principles and spirit of the present invention, and still fall within the scope of protection of the present invention.

Claims

1. A single-frame infrared dim small target detection method based on multi-scale feature enhancement, characterized in that: The steps include: Step 1: Obtain a single-frame infrared image to construct a single-frame infrared image dataset, and divide it into a training set and a test set; Step 2: construct an infrared small target detection network model, wherein the infrared small target detection network is based on the YOLOv5 network model, and the YOLOv5 network model includes a backbone feature extraction network, a feature fusion network and a detection head; A multi-scale feature enhancement module is introduced into the backbone feature extraction network, and the multi-scale feature enhancement module includes a context feature fusion module and a feature detail enhancement module connected in series in sequence; Step 3: Design a training strategy, and use the training set as input to train the infrared small target detection network model according to the training strategy; Step 4: After training is completed, input the test set to test and evaluate the infrared small target detection network model.

2. According to the single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 1, it is characterized in that: The context feature fusion module includes a block-guided context information perception module and a sequence convolution module.

3. The single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 2 is characterized in that: The feature detail enhancement module includes a detail enhancement convolution module and a CBAM attention mechanism.

4. The single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 3 is characterized in that: In the context feature fusion module, given a feature The average pooling operation, the maximum pooling operation and the Conv operation are performed through the block-guided context information perception module to extract the average pooling feature F within each block. avg , the maximum pooling feature F within the block max And the local feature F within the block local , and then the average pooling features F in the block are avg And the maximum pooling feature F within the block max Obtain inter-block global features with long-range interactions through the Nonlocal module and The global features between two blocks and Perform channel splicing to obtain the aggregated global feature F global , aggregate global features F global and the local feature F within the block local The corresponding elements in are multiplied to obtain the global and local feature fusion feature F fusion Finally, the original input feature F is connected with the fusion feature F through residual connection fusion Add them together to get the final feature F′.

5. The single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 4 is characterized in that: When the block-guided context information perception module extracts features, the window size is set to k×k, the step size is s, and the feature representation is F k ′.

6. The single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 4 is characterized in that: The sequence convolution module is a Bottleneck module, which is composed of a 1×1 Conv module, a 3×3 Conv module and a residual connection. The feature F obtained after passing through the sequence convolution module is F seq ; Finally, the block-guided context-aware module outputs the feature F k ′ and the sequence convolution module output feature F seq Add to get the output features of the context feature fusion module 7. The single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 4 is characterized in that: The detail enhancement convolution module includes standard convolution and differential convolution, wherein the differential convolution includes center differential convolution CDC, angle differential convolution ADC, vertical differential convolution VDC and horizontal differential convolution HDC, and the output of the context feature fusion module As input, it undergoes five convolutions in parallel, and finally the outputs of the five branches are element-wise added to obtain the output of the detail enhancement convolution module.

8. The single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 7 is characterized in that: The CBAM attention mechanism includes spatial attention and channel attention. Given a feature The features The output of the detail enhancement module is firstly passed through the channel attention, and the channel attention is respectively subjected to global average pooling and global maximum pooling in the spatial dimension to obtain the global feature description and Describing global features and The two channel attention vectors are added together and then normalized by the sigmoid function to generate the final channel attention weight. The input features and channel attention weight M C The main channel is multiplied to obtain the channel-enhanced feature map Channel enhanced feature map After spatial attention, the feature map of channel enhancement is Perform channel average pooling and channel maximum pooling in the channel dimension to obtain global feature description and Describing global features and After concatenation in the channel dimension, a 7×7 convolution kernel is used for convolution and sigmoid normalization is performed to obtain a spatial attention map. Finally, the channel enhanced feature map With the spatial attention map M S Element-wise multiplication yields the spatially enhanced feature F″′.

9. The single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 1 is characterized in that: In step 3, the training strategy of the infrared small target detection network model is: According to the set number of epochs, the first half of the training rounds use the Focal-NWD loss function to perform a rough bounding box regression, and the second half of the rounds use the Focal-αECIoU loss function to perform a fine bounding box regression. The positioning loss function formula is as follows: Among them, nepoch is the number of current training iterations, and epoch is the total number of training rounds.

10. The single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 9, characterized in that: In step 3, when the infrared small target detection network model is trained, the center coordinate values ​​of the prediction box are defined as cx and cy, and the width and height are defined as w and h; the center coordinate value of the real box is cx gt andcy gt , with width and height w gt and h gt ; Construct the Focal-αECIoU loss function, the specific formula is as follows: L αECIoU =1-αECIoU L Focal-αECIoU =IoU γ L αECIoU Among them, IoU is the intersection-over-union ratio, which is the ratio of the intersection and union between the model's predicted box and the true box; ρ 2 (b,b gt ), ρ 2 (w,w gt ), ρ 2 (h,h gt ) are the center point distance, width distance and length distance of the predicted box and the real box respectively; β is a balance parameter, ν is a parameter used to measure the consistency of aspect ratio, c is the diagonal length of the two bounding boxes, C w and C h are the width and height of the minimum bounding rectangle of the predicted box and the true box respectively; α is a Power parameter, and γ is a parameter that controls the degree of outlier suppression.

11. The single-frame infrared dim small target detection method based on multi-scale feature enhancement according to claim 9, characterized in that: The specific formula of the Focal-NWD loss function is as follows: L Focal-NWD =NWD γ (1-NWD) Where N and N gt Represents the Gaussian distribution of the predicted box and the true box; Indicates N and N gt Wasserstein distance between Gaussian distributions; NWD(N,N g ) represents the exponential normalized form of the Wasserstein distance between two Gaussian distributions, and c is a constant.

Images