Infrared dim eye segmentation method based on improved U-Net

Abstract

The invention aims to provide an infrared dim eye segmentation method based on improved U-Net, which comprises the following steps of: A, constructing an improved U-Net network which comprises a coding network and a decoding network; the coding network comprises five groups of multi-scale residual error modules MSRB which are connected in sequence, and a down-sampling module HWD is arranged between two adjacent groups of multi-scale residual error modules MSRB; b, training the improved U-Net network model to obtain a trained network model; c, inputting an infrared image into a coding network based on the trained deep neural network, and sequentially processing the infrared image through each group of multi-scale residual error modules MSRB and a down-sampling module HWD; and D, respectively inputting the processing result of each multi-scale residual module MSRB into a decoding network, and carrying out feature recovery by the decoding network to obtain a prediction result. According to the method, the feature expression capability can be enhanced, and the weak and small target segmentation capability is improved.

Description

Technical Field

[0001] The present invention relates to the technical field of infrared dim small target detection using artificial intelligence, and in particular to an infrared dim small target segmentation method based on an improved U-Net. Background Art

[0002] With the development of infrared imaging technology, infrared detectors have been widely used in military, medical, security and other fields due to their advantages of good concealment, all-weather operation, clear imaging, strong anti-electromagnetic interference ability and simple structure. Due to the small size of targets and weak signals in infrared images, as well as the complex background environment and low signal-to-noise ratio, detection is difficult. Most of the existing methods are difficult to achieve ideal results and there are many problems. In practical applications, a large number of false alarms are often generated, and the weakness of the target also makes it easy to miss detection. In recent years, with the complexity of the application scenarios of weak target detection, the requirements for weak target detection algorithms have become higher and higher. Therefore, the study of infrared weak target detection algorithms with robustness and robustness has become an urgent problem to be solved. With the development of computer hardware and artificial intelligence related theories, the use of deep learning methods to realize infrared weak target detection has become a development trend. It is of great practical significance to study weak target detection algorithms that meet the needs of actual engineering. Summary of the invention

[0003] The present invention aims to provide an infrared dim small target segmentation method based on improved U-Net, which can better focus on multi-dimensional useful information, enhance the expression ability of features, and improve the ability to segment dim small targets.

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

[0005] The infrared weak small eye segmentation method based on improved U-Net comprises the following steps:

[0006] A. Build an improved U-Net network, which includes an encoding network and a decoding network.

[0007] The encoding network includes five groups of multi-scale residual modules MSRB connected in sequence, and a downsampling module HWD is provided between two adjacent groups of multi-scale residual modules MSRB; the decoding network includes four upsamplings, and each upsampling is processed in sequence by a triple attention module TRA and a multi-scale residual module MSRB;

[0008] B. Train the improved U-Net network model to obtain a trained improved U-Net network model;

[0009] C. Based on the trained deep neural network, the infrared image is input into the encoding network and processed by each group of multi-scale residual modules MSRB and downsampling modules HWD in turn. The processing process in the multi-scale residual module MSRB is as follows:

[0010] After the input result is first subjected to 3×3 convolution, the obtained feature map is divided into four parts according to the channel. The first part of the features is not processed and is used as a feature reuse layer. The features of the second, third, and fourth parts are processed by deep convolution with dilation rates of 1, 3, and 5 respectively. The three dilated convolution results are concatenated with the features of the first part, and then the concatenated feature maps are fused between channels using a 1×1 convolution unit, and residually connected with the input feature map to obtain the output result.

[0011] D. The processing results of each multi-scale residual module MSRB are respectively input into the decoding network, which performs feature recovery to obtain the final prediction result.

[0012] The processing process in the downsampling module HWD is as follows:

[0013] The input result is divided into two paths. The first path is processed by the low-pass filter H0 and then downsampled twice to obtain the first downsampled result. The second path is processed by the high-pass filter H1 and then downsampled twice to obtain the second downsampled result.

[0014] The first downsampling result is divided into two branches. The first branch is processed by a low-pass filter H0 to obtain a low-frequency component A, and the second branch is processed by a high-pass filter H1 to obtain a horizontal component H.

[0015] The second downsampling result is divided into two branches. The first branch is processed by a low-pass filter H0 to obtain a vertical component V, and the second branch is processed by a high-pass filter H1 to obtain a diagonal component D.

[0016] The low-frequency component A, horizontal component H, vertical component V, and diagonal component D are concatenated, and the concatenated result is processed by a 1×1 convolution layer, batch normalization, and ReLu activation function to adjust the number of channels to obtain the output result.

[0017] The processing process in the decoding network is as follows:

[0018] The output result of the fifth group of multi-scale residual modules MSRB in the encoding network is upsampled twice and concatenated with the output result of the fourth group of multi-scale residual modules MSRB in the encoding network through the cat function, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in turn to obtain the first TRA+MSRB processing result;

[0019] The first TRA+MSRB processing result is upsampled twice and then concatenated with the output result of the third group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in sequence to obtain the second TRA+MSRB processing result;

[0020] The second TRA+MSRB processing result is upsampled twice and then concatenated with the output result of the second group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in sequence to obtain the third TRA+MSRB processing result;

[0021] The third TRA+MSRB processing result is upsampled twice and then concatenated with the output result of the first group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in sequence to obtain the fourth TRA+MSRB processing result;

[0022] The fourth TRA+MSRB processing result is fused between channels using a 1×1 convolution unit to obtain the final prediction result.

[0023] The processing process in the triple attention module TRA is as follows: the input feature tensor is divided into three branches for processing. The first branch rotates the tensor of size C×H×W by 90° counterclockwise along the H axis to obtain a tensor of size W×H×C. The rotation result is sequentially subjected to Z-pool calculation, 7×7 convolution, batch normalization, and Sigmoid activation function calculation. After the attention weight is generated, it is multiplied with the input result, and then the tensor is rotated 90° clockwise along the H axis to restore it to its initial shape to obtain the first branch result;

[0024] After the second branch rotates 90° counterclockwise along the W axis, the resulting tensor size is H×C×W. The rotation result is sequentially subjected to Z-pool calculation, 7×7 convolution, batch normalization, and Sigmoid activation function calculation. After the attention weight is generated, it is multiplied with the input result, and then the tensor is rotated 90° clockwise along the W axis to restore it to its initial shape to obtain the result of the second branch.

[0025] The third branch performs Z-pool calculation, 7×7 convolution, batch normalization, and Sigmoid activation function calculation in sequence. After generating the attention weight, it is multiplied by the input result to obtain the output result, thus obtaining the result of the third branch.

[0026] Finally, the first branch result, the second branch result, and the third branch result are added together and the average is taken to obtain the output result.

[0027] The method of the present invention adopts a down-sampling method based on Haar wavelet transform in the U-Net encoder, which can retain more detail information and reduce information loss during the feature extraction process.

[0028] The method of the present invention designs a multi-scale residual unit as the basic component unit of U-Net, which is used to more efficiently obtain multi-scale context information and improve the efficiency of feature extraction. Finally, a triple attention mechanism is added to the decoder structure to make the network focus on multi-dimensional useful information and enhance the expressiveness of features.

[0029] The present invention adopts the NUDT-SIRST infrared small target dataset for model training, and then evaluates and predicts the trained model on the test set. Experiments show that the algorithm invented in this paper can effectively improve the ability to segment small targets. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] Figure 1 This is a flow chart of the improved U-Net algorithm of Example 1;

[0031] Figure 2 Schematic diagram of the structure of the multi-scale residual module MSRB of Example 1;

[0032] Figure 3 Schematic diagram of the structure of the down-sampling module HWD based on Haar wavelet transform in Example 1;

[0033] Figure 4 Schematic diagram of the structure of the triple attention module TRA of Example 1;

[0034] Figure 5 The segmentation results of the method in Example 1 and other U-Net segmentation methods in the prior art are compared on the NUDT-SIRST dataset. DETAILED DESCRIPTION

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

[0036] Example 1

[0037] The infrared weak small eye segmentation method based on improved U-Net in this embodiment includes the following steps:

[0038] A. Figure 1 As shown, an improved U-Net network is constructed, and the improved U-Net network includes: an encoding network and a decoding network;

[0039] The coding network comprises five groups of multi-scale residual modules MSRB connected in sequence, and a downsampling module HWD is arranged between two adjacent groups of multi-scale residual modules MSRB; the module structure diagram of the multi-scale residual module MSRB is shown in Figure 2; The HWD structure diagram of the downsampling module is shown in Figure 3 ; The decoding network includes four upsamplings, and each upsampling is processed by the triple attention module TRA and the multi-scale residual module MSRB in turn. The structure diagram of the triple attention module TRA is shown in Figure 4 ;

[0040] B. Train the improved U-Net network model, randomly select 60% of the samples for training the model, and 40% of the samples for testing, where the sample resolution is 256. Use the trained network model to predict the test samples to achieve automatic segmentation of infrared dim targets;

[0041] C. Based on the trained deep neural network, the 256×256 infrared image is input into the encoding network and processed by each group of multi-scale residual modules MSRB and downsampling modules HWD in turn. The processing process in the multi-scale residual module MSRB is as follows:

[0042] The input result is first subjected to 3×3 convolution, and the obtained feature map is divided into four parts according to the channel. The first part of the features is not processed and is used as a feature reuse layer. The features of the second, third, and fourth parts are processed by deep convolution with dilation rates of 1, 3, and 5 respectively. The three dilated convolution results are concatenated with the features of the first part, and then the concatenated feature maps are fused between channels using a 1×1 convolution unit, and residually connected with the input feature map to obtain the output result.

[0043] The processing process in the downsampling module HWD is as follows:

[0044] The input result is divided into two paths. The first path is processed by the low-pass filter H0 and then downsampled twice to obtain the first downsampled result. The second path is processed by the high-pass filter H1 and then downsampled twice to obtain the second downsampled result.

[0045] The first downsampling result is divided into two branches. The first branch is processed by a low-pass filter H0 to obtain a low-frequency component A, and the second branch is processed by a high-pass filter H1 to obtain a horizontal component H.

[0046] The second downsampling result is divided into two branches. The first branch is processed by a low-pass filter H0 to obtain a vertical component V, and the second branch is processed by a high-pass filter H1 to obtain a diagonal component D.

[0047] The low-frequency component A, horizontal component H, vertical component V, and diagonal component D are all They are concatenated, and the concatenated results are processed in turn through a 1×1 convolutional layer, batch normalization, and ReLu activation function to adjust the number of channels to obtain the output result.

[0048] D. The processing results of each multi-scale residual module MSRB are respectively input into the decoding network, and the final prediction result is obtained after decoding by the decoding network.

[0049] The processing process in the decoding network is as follows:

[0050] The output result of the fifth group of multi-scale residual modules MSRB in the encoding network is upsampled twice and then concatenated with the output result of the fourth group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in turn to obtain the first TRA+MSRB processing result;

[0051] The first TRA+MSRB processing result is upsampled twice and then concatenated with the output result of the third group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in sequence to obtain the second TRA+MSRB processing result;

[0052] The second TRA+MSRB processing result is upsampled twice and then concatenated with the output result of the second group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in sequence to obtain the third TRA+MSRB processing result;

[0053] The third TRA+MSRB processing result is upsampled twice and then concatenated with the output result of the first group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in sequence to obtain the fourth TRA+MSRB processing result;

[0054] The fourth TRA+MSRB processing result is fused between channels using a 1×1 convolution unit to obtain the final prediction result.

[0055] like Figure 4As shown, the processing process in the triple attention module TRA is as follows: the input feature map needs to pass through three branch structures to output the result. For the input tensor X with a size of C×H×W, the interaction between the channel dimension C and the height dimension H is established in the first branch. First, the input tensor X is rotated 90° counterclockwise along the H axis to obtain a tensor with a shape size of W×H×C. After Z-pool calculation, a 2×H×C tensor is obtained. Then, convolution, batch normalization, and Sigmoid activation function calculation are performed to generate attention weights, which are multiplied with the input. Finally, the tensor is rotated 90° clockwise along the H axis to restore it to its initial shape. In the second branch, the interaction between the channel dimension C and the width dimension W is established. The input tensor is rotated 90° counterclockwise along the W axis to obtain a tensor with a shape size of H×C×W. After Z-pool calculation, a 2×C×W tensor is obtained. Then convolution, batch normalization, and Sigmoid activation function calculation are performed. After the attention weight is generated, it is multiplied with, and finally the tensor is rotated 90° clockwise along the W axis to restore it to its initial shape. The third branch performs Z-pool calculation on the input tensor X to obtain a tensor with a shape size of 2×H×W. Then convolution, batch normalization, and Sigmoid activation function calculation are performed. After the attention weight is generated, it is multiplied with the input to obtain the output result. Finally, the output results of the three branches are added and averaged to obtain the output features.

[0056] Z-pool calculation is to perform maximum pooling and average pooling calculations on the 0th dimension of the tensor, and concatenate the calculation results to reduce its dimension to 2D. This operation can enable the tensor to retain rich feature information and reduce the amount of calculation, which can be expressed as follows: where 0d represents the 0th dimension of the tensor.

[0057] Z-pool(x)=[Maxpool 0d (x),Avgpool 0d (x)] (1)

[0058] Example 2

[0059] Using the method of Example 1, the present invention and various semantic segmentation methods based on improved U-Net in the prior art achieve weak small target segmentation on the public dataset NUDT-SIRST. The specific evaluation data is shown in Figure 5 .

[0060] like Figure 5 As shown in Figure 1, ten representative backgrounds in the NUST-SIRST dataset were selected for testing. Figure 5It can be seen that background 1 is a building, ResUnet missed detection, and other algorithms can correctly segment the target point. Background 2 is the ground, U-Net has a small number of false detection targets, ResUnet and ResUnet++ have a large number of false detections, UCTransNet and U-Net3+ have one false detection target, DeepLab3+ did not detect the target, and the method of the present invention correctly segmented the target. Background 3 is the sky, and there are shallow clouds. U-Net has false detections, ResUnet++ only detects one target point, and the pixel description is too different from the real target. ResUnet, UCTransNet, U-Net3+, DeepLab3+ and the method of the present invention all correctly detect the target. Background 4 is a grassland in the wild. Due to the low contrast between the target and the background, the performance of each algorithm is not good. U-Net only detects one target and has two false detections. ResUnet, ResUnet++, and DeepLab3+ did not detect the target. UCTransNet and U-Net3+ only detected one target. The segmentation effect of the method of the present invention is the best, and both targets are accurately segmented. Background 5 is the water surface and the sky, without too many clutter signals, so all algorithms correctly segment the target. The method of the present invention achieves the best effect in describing the contour of the target.

Claims

1. An infrared weak small eye segmentation method based on improved U-Net, characterized in that: The following steps are involved: A. Build an improved U-Net network, which includes an encoding network and a decoding network. The encoding network includes five groups of multi-scale residual modules MSRB connected in sequence, and a downsampling module HWD is provided between two adjacent groups of multi-scale residual modules MSRB; the decoding network includes four upsamplings, and each upsampling is processed in sequence by a triple attention module TRA and a multi-scale residual module MSRB; B. Train the improved U-Net network model to obtain a trained improved U-Net network model; C. Based on the trained deep neural network, the infrared image is input into the encoding network and processed by each group of multi-scale residual modules MSRB and downsampling modules HWD in turn. The processing process in the multi-scale residual module MSRB is as follows: The input result is first subjected to 3×3 convolution, and the obtained feature map is divided into four parts according to the channel. The first part of the features is not processed and is used as a feature reuse layer. The features of the second, third, and fourth parts are processed by deep convolution with dilation rates of 1, 3, and 5 respectively. The three dilated convolution results are concatenated with the features of the first part, and then the concatenated feature maps are fused between channels using a 1×1 convolution unit, and residually connected with the input feature map to obtain the output result. D. The processing results of each multi-scale residual module MSRB are respectively input into the decoding network, and the decoding network restores the features to obtain the final prediction result.

2. The infrared weak small eye segmentation method based on improved U-Net as claimed in claim 1, characterized in that: The processing process in the downsampling module HWD is as follows: The input result is divided into two paths. The first path is processed by the low-pass filter H0 and then downsampled twice to obtain the first downsampled result. The second path is processed by the high-pass filter H1 and then downsampled twice to obtain the second downsampled result. The first downsampling result is divided into two branches. The first branch is processed by a low-pass filter H0 to obtain a low-frequency component A, and the second branch is processed by a high-pass filter H1 to obtain a horizontal component H. The second downsampling result is divided into two branches. The first branch is processed by a low-pass filter H0 to obtain a vertical component V, and the second branch is processed by a high-pass filter H1 to obtain a diagonal component D. The low-frequency component A, horizontal component H, vertical component V, and diagonal component D are concatenated, and the concatenated result is processed by a 1×1 convolution layer, batch normalization, and ReLu activation function to adjust the number of channels to obtain the output result.

3. The infrared weak small eye segmentation method based on improved U-Net as claimed in claim 1, characterized in that: The processing process in the decoding network is as follows: The output result of the fifth group of multi-scale residual modules MSRB in the encoding network is upsampled twice and then concatenated with the output result of the fourth group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in turn to obtain the first TRA+MSRB processing result; The first TRA+MSRB processing result is upsampled twice and then concatenated with the output result of the third group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in sequence to obtain the second TRA+MSRB processing result; The second TRA+MSRB processing result is upsampled twice and then concatenated with the output result of the second group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in sequence to obtain the third TRA+MSRB processing result; The third TRA+MSRB processing result is upsampled twice and then concatenated with the output result of the first group of multi-scale residual modules MSRB in the encoding network, and then processed by the triple attention module TRA and the multi-scale residual module MSRB in sequence to obtain the fourth TRA+MSRB processing result; The fourth TRA+MSRB processing result is fused between channels using a 1×1 convolution unit to obtain the final prediction result.

4. The infrared weak small eye segmentation method based on improved U-Net as claimed in claim 1, characterized in that: The processing process in the triple attention module TRA is as follows: the input feature tensor is divided into three branches for processing. The first branch rotates the tensor of size C×H×W by 90° counterclockwise along the H axis to obtain a tensor of size W×H×C. Z-pool calculation, 7×7 convolution, batch normalization, and Sigmoid activation function calculation are performed in sequence. After the attention weight is generated, it is multiplied with the input result, and then the tensor is rotated 90° clockwise along the H axis to restore it to its initial shape to obtain the first branch result; After the second branch rotates 90° counterclockwise along the W axis, the resulting tensor size is H×C×W. Z-pool calculation, 7×7 convolution, batch normalization, and Sigmoid activation function calculation are performed in sequence. After the attention weight is generated, it is multiplied with the input result, and then the tensor is rotated 90° clockwise along the W axis to restore it to its initial shape to obtain the result of the second branch; The third branch performs Z-pool calculation, 7×7 convolution, batch normalization, and Sigmoid activation function calculation in sequence. After generating the attention weight, it is multiplied by the input result to obtain the output result, thus obtaining the result of the third branch. Finally, the first branch result, the second branch result, and the third branch result are added together and the average is taken to obtain the output result.

Images