Infrared small target detection method based on diffusion model

Abstract

The present invention belongs to the field of infrared target detection technology, and specifically relates to an infrared small target detection method based on a diffusion model, and the steps are as follows: S1, constructing a conditional diffusion network model with a dual-branch architecture, the model integrates a conditional network branch for fully extracting prior features from an input infrared image and a denoising network branch for recovering a clean mask from a noise mask; S2, constructing a two-dimensional Transformer module, and significantly enhancing the model's ability to recognize the global structure and detailed features of targets under complex backgrounds through joint analysis of spatial and channel dimensions; S3, designing a cross-attention fusion module, and realizing accurate guidance of the denoising branch by the conditional prior branch information through calculating the cross-attention between the conditional prior branch and the denoising branch. It solves the technical problem of "improving the accuracy and reliability of infrared small target recognition". It transforms the infrared small target detection task into a generative paradigm solution, and significantly improves the accuracy and reliability of small target recognition.

Description

Technical Field

[0001] The invention belongs to the technical field of infrared target detection, and in particular relates to research on a small infrared target detection method based on a diffusion model. Background Art

[0002] Compared with ordinary RGB imaging, infrared imaging technology has detection capabilities beyond the visible spectrum, can penetrate darkness and bad weather, and capture infrared radiation information released by objects. This unique advantage ensures that infrared images can still provide stable target detection and identification services under extreme light shortages or bad weather conditions, and thus become the core technical support for key areas such as natural disaster monitoring and aerial reconnaissance. However, with the advancement of technology, effectively identifying and accurately locating tiny targets from complex infrared images has become the focus of research. Due to the long imaging distance and equipment limitations, these small targets account for a very low proportion in the image, usually not exceeding 0.15%. In addition, the uneven distribution of environmental heat may cause blurred contours and loss of texture information, further exacerbating the problem of target features being submerged by background noise, resulting in a decrease in image resolution. Therefore, considering the current state of technology and actual challenges, how to achieve accurate detection of tiny targets in infrared images under various constraints is still a technical difficulty that needs to be overcome.

[0003] In recent years, infrared small target detection technology has undergone rapid innovation, thanks to the public release of high-quality datasets such as NUDT-SIRST, NUAA-SIRST, and IRSTD-1k, as well as the rapid development of deep learning technology. Driven by this, a series of data-driven deep learning models have emerged, which cleverly use the power of convolutional neural networks to achieve a significant leap in accuracy and stability in infrared image small target detection, successfully surpassing the performance limitations of traditional detection models. Dai et al. innovatively designed an asymmetric context modulation module (ACM), using a novel bottom-up strategy to deeply embed subtle visual features into the network, greatly enhancing the detection ability of the model. On this basis, by introducing a local contrast enhancement mechanism, they launched ALCNet, which once again refreshed the accuracy record of infrared small target detection. Based on the finite difference theory of Taylor expansion, Zhang's team developed ISNet, which effectively aggregates and strengthens edge features at different levels, greatly improving the contrast between the target and the background, and providing a solid theoretical support for the extraction of target boundary information under the guidance of mathematical theory. Li et al. demonstrated the innovative design of DNA-Net, a dense and nested attention network that achieves deep integration and efficient use of small target context information through a cyclic feature fusion and enhancement process, opening up a new path for the advancement of detection technology. In addition, the UIU-Net framework proposed by Wu et al. cleverly achieves dual optimization of global and local context information and multi-scale feature extraction by nesting U-Net inside the classic U-Net architecture. Although the above-mentioned deep learning-based infrared small target detection technology has made remarkable achievements, most of the current methods are still limited by the segmentation paradigm, and the model expression ability faces bottlenecks. In the face of extreme imaging conditions such as strong noise interference and uneven environmental thermal effects, efficient and accurate detection of small targets is still a challenge that needs to be solved. Summary of the invention

[0004] In view of this, the purpose of the present invention is to provide an infrared small target detection field. Most of the existing detection methods are based on the segmentation paradigm, which is limited by the expression ability of the model. When faced with extreme imaging conditions such as strong noise interference and uneven environmental thermal effects, the detection performance is often affected. Therefore, the present invention proposes a method for infrared small target detection based on a diffusion model, defines the detection task as a generation paradigm, and designs a conditional denoising model. In the training stage, Gaussian noise is added to the clean target mask, and the corresponding infrared image is used as a conditional prior training model to reverse this process. In the inference stage, the model gradually refines the initial noise mask through the learned denoising model until they conform to the distribution of the target mask. Its core lies in the following three parts:

[0005] S1. Construct a dual-branch conditional diffusion network model, which integrates the conditional network branch to extract the infrared image from the input. The prior features are fully extracted from the noise mask and the denoising network branch is used to Recover a clean mask .

[0006] S2. Construct a two-dimensional Transformer module to significantly enhance the model's ability to recognize the global structure and detailed features of targets in complex backgrounds through joint analysis of spatial and channel dimensions.

[0007] S3. Design a cross-attention fusion module to calculate the cross-attention between the conditional prior branch and the denoising branch, so as to achieve accurate guidance of the denoising branch by the conditional prior branch information.

[0008] Furthermore, the conditional diffusion network model includes two processes: forward diffusion and reverse denoising. The forward process is based on the Markov chain to the clean target mask Gradually add Gaussian noise: , and finally get an image close to pure Gaussian noise , where the time step is , , represents the maximum time step; is a predefined monotonically increasing time step parameter; is the noise sampled from the standard normal distribution. The backward process gradually removes the noise through the inverse Markov chain, from the Gaussian distribution Step by step, restore to the original clean target mask For each time step , the denoising and restoration process can be expressed as:

[0009]

[0010] in represents the denoising model, Represents its parameters, , The denoising model consists of two parts: the conditional prior branch and the denoising branch.

[0011] The input of the conditional prior branch is the infrared image , firstly, shallow features are extracted from it through a single convolutional layer. Subsequently, three groups of two-dimensional Transformer modules are used to extract effective prior features from the infrared image as guidance information for the denoising branch. In addition, in order to improve the network's ability to perceive multi-scale information, a two-fold downsampling operation is used after the first two groups of two-dimensional Transformer modules. Finally, the output of the last layer of two-dimensional Transformer modules is sent to the residual block, and the result after the upsampling operation is spliced ​​with the output of the second layer of two-dimensional Transformer modules along the channel dimension and sent to the residual block to achieve information fusion. Similarly, the output result of the residual block is doubled upsampled to restore the features to the original resolution of the infrared image and spliced ​​with the output of the first layer of two-dimensional Transformer modules along the channel dimension. The corresponding residual block is used to integrate the multi-scale target information, and the output is used as guidance information for the denoising branch.

[0012] The infrared image and the noisy mask features are concatenated as the input of the denoising branch, and shallow features are extracted through a single convolutional layer. The denoising branch mainly consists of two parts: the encoder and the decoder, which are composed of three sets of two-dimensional Transformer modules connected in series to improve the denoising ability of the model. A cross-attention fusion module is used between the encoder and the decoder to integrate the encoded information with the guidance information extracted by the conditional prior branch, effectively denoising the noisy mask. Finally, a single-layer convolution operation is performed on the output of the decoder to obtain the final clean mask.

[0013] Furthermore, the two-dimensional Transformer module is composed of three parts connected in series: a spatial dimension self-attention calculation module, a channel dimension self-attention calculation module and a feedforward network. At the same time, residual connections are used for these three groups of calculation units to improve the network fitting ability.

[0014] The spatial self-attention calculation module is used to capture the spatial dimension similarity relationship of infrared features. For the input features, in order to reduce the computational complexity, it is first divided into The image block is of size . Then, three sets of linear mapping operations are used to obtain the query, key and value vectors respectively. The key is reshaped and matrix multiplied with the query, and then the spatial dimension attention feature map of the image block is calculated by the softmax function. The attention feature map is matrix multiplied with the value vector and added to the spatial encoding feature representing the spatial position. Finally, the calculation results of all image blocks are summarized and a layer of linear mapping is performed to generate the output result of the spatial dimension self-attention calculation module.

[0015] The channel-dimensional attention calculation module is used to obtain the complex interactive relationship of infrared features at the channel level. The input features are first reshaped to exchange the spatial and channel dimensions of the features. Then, similar to the spatial-dimensional attention calculation module, three sets of linear mapping operations are used to obtain the query, key, and value vectors. The key is reshaped and matrix multiplied with the query, and then the channel-dimensional attention feature map of the image block is calculated through the softmax function. The attention feature map is matrix multiplied with the value vector and added to the channel encoding feature representing the channel position. Finally, a layer of linear mapping generates the output result of the channel-dimensional self-attention calculation module.

[0016] The feedforward network consists of residual blocks, which are implemented through three sets of convolutions and activation functions. It performs nonlinear transformation on the image features after processing by the self-attention layer, significantly enhancing the expressive power of the model and thus learning more complex feature representations.

[0017] Furthermore, the cross-attention fusion module calculates the cross-attention between the conditional prior branch and the denoising branch, thereby achieving efficient guidance of the denoising branch by the conditional prior branch information.

[0018] The input of the cross-attention fusion module is the output of the conditional prior branch, denoted as x, and the input of the encoder in the denoising branch, denoted as y. For these two features, three sets of linear mapping operations are used to obtain the query, key, and value vectors. Then, the query vectors obtained by x and the mapping are exchanged, and the self-attention is calculated separately. The final calculation results are added, and the two features are fused through a linear mapping operation. Finally, the output of the module is obtained through a residual group module.

[0019] The present invention proposes a method for infrared small target detection based on a diffusion model, which transforms the infrared small target detection task into a generative solution, significantly improving the accuracy and reliability of small target recognition. The method designs a conditional denoising network model with a dual-branch architecture, including a conditional prior branch for extracting prior knowledge of infrared images and a denoising branch for removing noise in masks. By deeply mining the intrinsic statistical characteristics of the image, the recognition accuracy of small targets is enhanced. In addition, a two-dimensional Transformer module is designed, which strengthens the model's perception of the global structure and local details of the target in complex scenes by combining spatial and channel information. Finally, the present invention also implements a cross-attention fusion module to achieve efficient information interaction between the conditional prior branch and the denoising branch in the conditional denoising network model. Experimental results show that the present invention has excellent detection performance on infrared small target datasets. BRIEF DESCRIPTION OF THE DRAWINGS

[0020] Figure 1 The figure is a schematic diagram of the invented infrared small target detection model based on the diffusion model;

[0021] Figure 2 Schematic diagram of the spatial self-attention calculation module;

[0022] Figure 3 Schematic diagram of the channel-dimensional self-attention calculation module;

[0023] Figure 4 Schematic diagram of the invented cross attention fusion module;

[0024] Figure 5 The following is a visualization result diagram of a typical scenario. DETAILED DESCRIPTION

[0025] Various exemplary embodiments, features and aspects of the present disclosure will be described in detail below with reference to the accompanying drawings. The same reference numerals in the accompanying drawings represent elements with the same or similar functions. Although various aspects of the embodiments are shown in the accompanying drawings, the drawings are not necessarily drawn to scale unless otherwise specified.

[0026] In addition, in order to better illustrate the present disclosure, numerous specific details are given in the following specific embodiments. It should be understood by those skilled in the art that the present disclosure can also be implemented without certain specific details. In some examples, methods, means, components and circuits well known to those skilled in the art are not described in detail in order to highlight the subject matter of the present disclosure.

[0027] In this embodiment, an infrared small target detection model is provided. The model is constructed based on a diffusion model and includes two processes: forward diffusion and reverse denoising. The forward process is based on a Markov chain to a clean target mask. Gradually add Gaussian noise: , and finally get an image close to pure Gaussian noise , where the time step is , , represents the maximum time step; is a predefined monotonically increasing time step parameter; is the noise sampled from the standard normal distribution. The backward process gradually removes the noise through the inverse Markov chain, from the Gaussian distribution Step by step, restore to the original clean target mask For each time step , the denoising and restoration process can be expressed as:

[0028]

[0029] in represents the denoising model, Represents its parameters, , The denoising model consists of two parts: the conditional prior branch and the denoising branch.

[0030] The conditional prior branch is used to extract the target prior from the infrared image and assist the denoising branch to restore the noisy mask. Specifically, it includes a single convolutional layer, three serially connected two-dimensional Transformer modules, and three groups of residual blocks.

[0031] This single convolution layer is used to extract shallow features from the input infrared image, and its parameters can be set according to the specific application scenario. For example, the convolution kernel size of this single convolution layer can be set to 3×3, the step size can be set to 1, the number of padding layers can be set to 1, and the bias term can be used.

[0032] The two-dimensional Transformer module is used to fully extract the prior knowledge in the infrared image. In order to extract features at different levels, a two-fold downsampling operation is also used after the first two two-dimensional Transformer modules. In order to integrate the information of the features extracted at different levels, the output of the last layer of the two-dimensional Transformer module is first sent to the residual block, and the result after the upsampling operation is spliced ​​with the output of the second layer of the two-dimensional Transformer module along the channel dimension and sent to the residual block to achieve information fusion. Similarly, the output result of the residual block is doubled upsampled to restore the features to the original resolution of the infrared image and spliced ​​with the output of the first layer of the two-dimensional Transformer module along the channel dimension. The corresponding residual block is used to integrate the multi-scale target information, and the output is used as the guidance information for the denoising branch.

[0033] The residual block is implemented by three sets of convolutions and activation functions. The first two convolution kernels are set to 3×3, the step size is set to 1, the number of padded layers is set to 1, and the bias term is used. A residual connection is also set after these two convolutions to add the input of the residual block. The latter set of convolution kernels is set to 1×1, the step size is set to 1, the number of padded layers is set to 0, and the bias term is used. In order to increase the nonlinear fitting ability of the network, a ReLU activation function is set after each convolution layer.

[0034] The input of the denoising branch is the concatenated features of the infrared image and the noisy mask, which is used to denoise the noisy mask according to the guidance of the conditional prior branch. It specifically includes a single convolutional layer, an encoder, a cross self-attention fusion module, a decoder module and a single convolutional layer.

[0035] The first single convolutional layer is used to extract shallow features from the input, and its parameters can be set according to the specific application scenario. For example, the convolution kernel size of the single convolutional layer can be set to 3×3, the stride can be set to 1, the number of padding layers can be set to 1, and the bias term can be used.

[0036] The encoder and decoder are used to learn the denoising process of noisy masks, and both consist of three two-dimensional Transformer modules in series.

[0037] The cross-self-attention fusion module is located between the encoder and the decoder, and is used to integrate the encoded information with the guidance information extracted by the conditional prior branch, effectively guiding the denoising branch to denoise the noisy mask.

[0038] The two-dimensional Transformer module is composed of three parts connected in series: a spatial self-attention calculation module, a channel self-attention calculation module, and a feedforward network. At the same time, residual connections are used for these three groups of calculation units to improve the network fitting ability.

[0039] The spatial self-attention calculation module is used to capture the spatial dimension similarity relationship of infrared features. For the input features, in order to reduce the computational complexity, it is first divided into The image block is of size . Then, three sets of linear mapping operations are used to obtain the query, key and value vectors respectively. The key is reshaped and matrix multiplied with the query, and then the spatial dimension attention feature map of the image block is calculated by the softmax function. The attention feature map is matrix multiplied with the value vector and added to the spatial encoding feature representing the spatial position. Finally, the calculation results of all image blocks are summarized and a layer of linear mapping is performed to generate the output result of the spatial dimension self-attention calculation module.

[0040] The channel-dimensional attention calculation module is used to obtain the complex interactive relationship of infrared features at the channel level. The input features are first reshaped to exchange the spatial and channel dimensions of the features. Then, similar to the spatial-dimensional attention calculation module, three sets of linear mapping operations are used to obtain the query, key, and value vectors. The key is reshaped and matrix multiplied with the query, and then the channel-dimensional attention feature map of the image block is calculated through the softmax function. The attention feature map is matrix multiplied with the value vector and added to the channel encoding feature representing the channel position. Finally, a layer of linear mapping generates the output result of the channel-dimensional self-attention calculation module.

[0041] The feedforward network consists of residual blocks, which are implemented through three sets of convolutions and activation functions. It performs nonlinear transformation on the image features after processing by the self-attention layer, significantly enhancing the expressive power of the model and thus learning more complex feature representations.

[0042] The cross-attention fusion module calculates the cross-attention between the conditional prior branch and the denoising branch, so as to achieve efficient guidance of the denoising branch by the conditional prior branch information.

[0043] The input of the cross-attention fusion module is the output of the conditional prior branch, denoted as x, and the input of the encoder in the denoising branch, denoted as y. For these two features, three sets of linear mapping operations are used to obtain the query, key, and value vectors. Then, the query vectors obtained by x and the mapping are exchanged, and the self-attention is calculated separately. The final calculation results are added, and the two features are fused through a linear mapping operation. Finally, the output of the module is obtained through a residual group module.

[0044] We finally chose to use IPI, NRAM, PSTNN, MDvsFA, ACM and ALCNet methods on the NUAA dataset to verify the advancement and effectiveness of this method. The visualization results in typical scenarios are shown in the figure below. Figure 5 As shown. The present invention adopts several commonly used evaluation indicators in this field, including IoU (Intersection over Union), nIoU (norm Intersection over Union) and Pd, and other infrared small target detection methods for quantitative evaluation and analysis. Among them, IoU represents the intersection over union ratio, and the value range is 0~1. The closer the value is to 1, the better the detection effect. nIoU represents the standardization of the IoU indicator, and the value range is 0~1. The closer the value is to 1, the better the detection effect. Pd represents the detection rate, which means the ratio of correctly detected targets to all targets, and the value range is 0~100%. The closer the value is to 100%, the better the detection effect. The results are shown in Table 1.

[0045] Table 1 Comparison of infrared small target detection results

[0046] method Iou iGO Pd IPI 2.62 4.16 84.40 NRAM 45.68 55.49 85.32 PSTNN 51.95 62.66 82.57 MDvsFA 60.28 58.16 76.15 ACM 71.96 71.05 96.25 ALCNet 73.43 71.44 97.23 The present invention 74.30 72.87 97.32

[0047] From the experimental results, the detection effect of the method of the present invention is obviously better than that of the comparison algorithm, thus verifying the excellence of the present invention.

Claims

1. A method for detecting small infrared targets based on a diffusion model, characterized in that: The following steps are involved: S1. Construct a dual-branch conditional diffusion network model, which integrates the conditional network branch to extract the infrared image from the input. The prior features are fully extracted from the noise mask and the denoising network branch is used to Recover a clean mask ; S2. Construct a two-dimensional Transformer module to significantly enhance the model's ability to recognize the global structure and detailed features of targets in complex backgrounds through joint analysis of spatial and channel dimensions; S3. Design a cross-attention fusion module to calculate the cross-attention between the conditional prior branch and the denoising branch, so that the conditional prior branch information can accurately guide the denoising branch; The step S1 comprises the following steps: The S1.1 conditional diffusion network model includes two processes: forward diffusion and reverse denoising. The forward process is based on the Markov chain to the clean target mask Gradually add Gaussian noise: , and finally get an image close to pure Gaussian noise , where the time step is , , represents the maximum time step; is a predefined monotonically increasing time step parameter; is the noise sampled from the standard normal distribution; the backward process gradually removes the noise through the inverse Markov chain, from the Gaussian distribution Step by step, restore to the original clean target mask ; For each time step , the denoising and restoration process can be expressed as: ； in represents the denoising model, Represents its parameters, , ; The denoising model consists of two parts: the conditional prior branch and the denoising branch; The input of the S1.2 conditional prior branch is an infrared image. First, the shallow features of the image are extracted through a single convolutional layer. Then, three groups of two-dimensional Transformer modules are used to extract effective prior features from the infrared image, and these features will serve as the guiding information of the denoising branch. In order to enhance the network's perception of multi-scale information, a two-fold downsampling operation is introduced after the first two groups of two-dimensional Transformer modules. Finally, the output of the last group of two-dimensional Transformer modules is sent to the residual block. After the upsampling operation, the result is spliced ​​with the output of the second group of two-dimensional Transformer modules in the channel dimension and input into the residual block for information fusion. Similarly, the output of the residual block is upsampled twice to restore the features to the original resolution of the infrared image, and spliced ​​with the output of the first group of two-dimensional Transformer modules in the channel dimension, and then the multi-scale target information is integrated through the corresponding residual block, and finally the output is used as the guiding information of the denoising branch. S1.3 concatenates the infrared image and the noisy mask features as the input of the denoising branch, and extracts shallow features through a single convolutional layer; the denoising branch mainly consists of two parts, the encoder and the decoder, which are composed of three groups of two-dimensional Transformer modules connected in series to improve the denoising ability of the model; a cross-attention fusion module is used between the encoder and the decoder to integrate the encoded information with the guidance information extracted by the conditional prior branch to effectively denoise the noisy mask; finally, a single-layer convolution operation is performed on the output result of the decoder to obtain the final clean mask.

2. The infrared small target detection method based on diffusion model according to claim 1 is characterized in that: The step S2 comprises the following steps: The S2.1 two-dimensional Transformer module is composed of three parts connected in series: a spatial self-attention calculation module, a channel self-attention calculation module, and a feedforward network. At the same time, residual connections are used for these three groups of calculation units to improve the network fitting ability. S2.2 The spatial dimension self-attention calculation module is used to capture the spatial dimension similarity relationship of infrared features. For the input features, in order to reduce the computational complexity, it is first divided into The image block is of size; then three sets of linear mapping operations are used to obtain the query, key and value vectors respectively; the key is reshaped and matrix multiplied with the query, and then the spatial dimension attention feature map of the image block is calculated by the softmax function; the attention feature map is matrix multiplied with the value vector and added with the spatial encoding feature representing the spatial position; finally, the calculation results of all image blocks are summarized and a layer of linear mapping is performed to generate the output result of the spatial dimension self-attention calculation module; S2.3 channel-dimensional attention calculation module is used to obtain the complex interactive relationship of infrared features at the channel level; the input features are first reshaped to exchange the spatial and channel dimensions of the features; then, similar to the spatial-dimensional attention calculation module, three sets of vectors, query, key and value, are obtained through three sets of linear mapping operations respectively; the key is reshaped and matrix multiplied with the query, and then the channel-dimensional attention feature map of the image block is calculated through the softmax function; the attention feature map is matrix multiplied with the value vector and added to the channel coding feature representing the channel position; finally, a layer of linear mapping generates the output result of the channel-dimensional self-attention calculation module; The S2.4 feedforward network consists of residual blocks, which are implemented through three groups of 3×3 convolutions and activation functions. It performs nonlinear transformation on the image features after processing by the self-attention layer, significantly enhancing the expressive power of the model and thus learning more complex feature representations.

3. The infrared small target detection method based on diffusion model according to claim 1 is characterized in that: The step S3 comprises the following steps: The S3.1 cross-attention fusion module calculates the cross-attention between the conditional prior branch and the denoising branch to achieve efficient guidance of the denoising branch by the conditional prior branch information; The input of the S3.2 cross-attention fusion module is the output of the conditional prior branch, denoted as x, and the input of the encoder in the denoising branch, denoted as y; for these two features, three sets of linear mapping operations are used to obtain the query, key and value vectors respectively; then, the query vectors obtained by x and the mapping are exchanged, and the self-attention calculation is performed separately; the final calculation results are added, and the fusion of the two features is realized through a linear mapping operation; finally, the output of the module is obtained through a residual group module.

Images