UIU-Net: U-Net in U-Net for Infrared Small Object Detection

Image segmentation is a pixel-level classification problem that involves dense image prediction. The explosion of convolutional neural networks (CNN) makes the image segmentation mission transform from patch classification [33], fully convolutional networks (FCN) [34], encoder-decoder framework[35] to dilated/atrous framework [36]. Subsequently, new semantic segmentation-based infrared small object detection networks, including ACM[30], ALCNet[31], DNIM [32], have begun arriving. This network’s detection performance has improved tremendously compared to model-based and DL-based detection methods. However, these methods mostly rely on classification backbone networks with a typical downsampling scheme. Even novel down-sampling schemes updated by some researchers [37, 38], the local and global contrast information of the infrared small object has still been ignored.

Different from the above works, we design a multi-scale depth supervision structure to resolve the conflict between feature resolution and network depth, resulting in improved global and local context representation, as well as improved final detection performance.

The attention mechanism was first proposed in the field of computer vision. Deepmind, a division of Google, created an attention mechanism for image classification in $2014$, allowing neural networks to pay more attention to relevant input regions while generating prediction missions. Subsequently, several types of research on attention mechanisms[39, 40, 41, 42, 43] have also been proposed and widely used in many fields, including image classification [41, 40], semantic segmentation [40], object detection [42], and more, with good results. These attention mechanisms embedded in general networks (e.g., classification backbone) are not suited for detecting infrared small objects because small-sized objects are frequently missing in deep semantic feature maps.

Different from previous works, the proposed interactive-cross attention can first encode pixel-based local contextual features that enhance the detail information of infrared small objects, and the encoded features are derived from a resolution-maintaining deep supervision network that can learn multi-scale and global features. The discriminability of infrared small objects and backgrounds can be greatly improved by this integration.

Review of U-Net Framework. To begin, we briefly review the classical U-Net, which was initially developed for semantic segmentation, especially medical image segmentation. U-Net performs multi-scale prediction and provides finer segmentation features. The low-resolution feature map after multiple downsampling in the encoder stage could obtain the global semantic information of the segmented object. Skip connection in the decoder stage ensures the final reconstructed feature map merges more low-level features and also blends multi-scale features. However, multiple downsampling deteriorates feature resolution and object contrast, which yielded a very negative effect on infrared small object detection.

Network Overview. To address the above problem, and ensure the detection performance of the infrared small objects. In this paper, we have proposed an interactive-cross attention-nested U-Net network with training from scratch, called U-Net in U-Net (UIU-Net). UIU-Net is not depending on a classic classification backbone network and is ideal for infrared small object detection. It begins with a resolution-maintenance deep supervision module that learns deep multi-scale features while improving global feature representation. The resolution maintenance here refers specifically to the feature learned throughout each stage of the encoder-decoder network. After that, the learned features from each stage are then fed into the interactive-cross attention module, which encodes objects’ local features to increase the distinguishability ability of infrared small objects. Finally, the multiple intermediate supervision and the last layer are weighted and merged by minimizing a typical cross-entropy loss.

Infrared images typically lack higher contrast due to their own imaging condition, resulting in a very low distinguishability between the object and background. This, coupled with the small size of the objects in the image, makes it difficult to precisely locate infrared small objects with a typical DL-based detection network or backbone-based segmentation network. In this section, we introduce a RM-DS network to overcome the above dilemma. RM-DS network is built on the base of U-Net. Multiple intermediate layers are utilized instead of just the last layer to obtain complete and distinguishable features, yielding better detection results. The output $\mathcal{O}$ of RM-DS network is defined as

where $F_{k}$ denotes the $k$-th layer feature, and $\mathcal{O}$ is the final output of RM-DS network. $\sigma(\cdot)$ denotes $sigmoid$ activation function.

In order to trade off the network depth and the feature resolution, the ReSidual U-blocks (RSU) module [44] is introduced as the backbone network. The detailed structure is illustrated in Fig. 4. RSU module is a type of U-Net. The difference is that: 1) the RSU module takes intermediate feature maps rather than input images as input to learn and encode deep multi-scale features; 2) Dilated convolution is used to improve deep feature resolution in each layer. For shallow layers, only the last convolution is replaced with dilated convolution (see Fig. 4(a)); However, for deep layers, all convolutions are replaced by dilated convolutions with varying dilated rates (see Fig. 4(b)), yielding low memory consumption due to the small size of deep feature maps; 3) Pooling operation has been introduced to the RSU module to reduce computing costs but has not been added to the additional RSU module to reduce feature loss. The feature of RSU backbone $U$ is defined as $F_{U}=U(f(x))+f(x)$, where $f(x)$ is intermediate feature maps.

where $U_{k}(\cdot)(k=1,2,\cdots,K)$ represents $K$ dilated-based convolutions modes, with just the last convolution replaced by dilated convolution and all convolutions replaced by dilated convolutions with different dilated rates. The particular use of which one is dependent on the situation.

Thanks to this design, the increase in network depth inside each stage has no effect on the feature resolution, and the growth in depth of each stage benefits from the growth in overall network depth. Meanwhile, an increase in network depth has enhanced the global context representation. However, only by MS-DS module not being able to locate infrared small object detection more precisely, especially under complex and ever-changing backgrounds.

In this section, an interactive-cross attention module is designed to encode local context representation and significantly improve detection performance. Figure 5 depicts the detailed architecture. This module instead of the original skip connection in U-Net retains more context information from the decoder layer. Specifically, low-level detail and high-level semantic features are the coding objects of the interactive-cross attention module. Cross-channel attention and interactive-cross spatial attention are two suboutputs of the interactive-cross attention module.

Each channel of high-level semantic features can be employed as specific object responses, and they are all associated with different semantic responses. We can focus on discovering the interesting infrared small objects, by making use of their interdependencies.

We define $U=[{{u_{1}},{u_{2}},\cdots,{u_{C}}}]$ as the encoder high-level features after ReSidual U-blocks [44] of deep network layer, where $u_{c}=F_{c}^{h}\in{\mathcal{R}^{W\times H}}$ is the $c$-th channel feature, and $C$ is the total number of channels. For these features, we first adopt AdaptiveAvgPool2d (short for $\mathbf{F}_{ad}$) to traverse features

where $\mathbf{z}_{c},c=1,2,\cdots,C$ stands the $c$-th channel. $H,W$ is the size of the high-level features. ${u_{c}}\left({i,j}\right)$ is the feature of each location in the $c$-th channel.

After that, the Excitation operation (short for $\mathbf{F}_{ex}$) ($C\to{C\mathord{/{\vphantom{Cr}}\kern-1.2pt}r}$, ${C\mathord{/{\vphantom{Cr}}\kern-1.2pt}r}\to C$, r=4) is used to reshape high-level features ($F^{h}$) to minimize network parameters.

where $\delta(\cdot),\mathcal{B}(\cdot)$ are respectively the rectified linear unit (ReLU), and batch normalization (BN). ${\mathbf{W}_{1}}\in{R^{\frac{C}{r}\times C}}$ and ${\mathbf{W}_{2}}\in{R^{C\times\frac{C}{r}}}$ stands for $C\to{C\mathord{/{\vphantom{Cr}}\kern-1.2pt}r}$ and ${C\mathord{/{\vphantom{Cr}}\kern-1.2pt}r}\to C$ Excitation operation, respectively.

Finally cross channel features are a weighted sum of the features with low-level features,

where $\otimes$ denotes the element-wise multiplication, and $F^{l}$ denotes low-level features in the encoder.

Considering infrared small object detection couldn’t neglect the object’s local detailed information to enhance the distinguishability between objects and background, interactive-cross spatial attention has been further designed. It first perform Excitation operation ($\mathbf{F}_{ex}$) ($C\to{C\mathord{/{\vphantom{Cr}}\kern-1.2pt}r}$, r=4) for $F_{C-A}$, and then aggregates it using both average-pooling ($\mathcal{P}_{avg}$) and max-pooling ($\mathcal{P}_{max}$) operations.

where $C^{3\times 3}$ and $C^{1\times 1}$ are convolution operation, respectively. $F_{C-A}^{1\times 1}$ stands for $F_{C-A}$ after $1\times 1$ convolution.

The interactive-cross spatial attention features are a weighted sum of the features with high-level features,

where $\otimes$ denotes the element-wise multiplication, and $F^{h}$ denotes high-level features in the encoder.

Finally, the output of UIU-Net is defined as

In this paper, we model infrared small object detection as a semantic segmentation problem. The detection performance is evaluated by computing three commonly-used indices, i.e., Intersection over union (IoU), normalized intersection over union (nIoU), and the receiver operating characteristic (ROC) curve. In general, the higher values for the three indices means better detection performance in the infrared small object detection task.

1) Intersection over Union (IoU): IoU is widely used for semantic segmentation. It is calculated by the intersection of the real and predicted values of the pixel divided by their union.

where $M$ is the object number of each image/sample, and $N$ is the total sample number of the testing set. ${{A_{{\mathop{\rm int}}{\rm{er}}}}}$ and ${{A_{all}}}$ stands for the summation result for the intersection and union of the all real and predicted objects in the testing set. $T$, $P$, and $TP$ stand for reference label, predict the result, and true positive pixel numbers, respectively. the $i,m$ in $TP_{m}^{i}$, $T_{m}^{i}$, $P_{m}^{i}$, is the $m$-th object in $i$-th sample.

2) normalized Intersection over Union (nIoU): signal-to-noise ratio (SCR) is calculated by the object and its surrounding pixels, which is widely used in the traditional method, is infinite in deep learning-based semantic segmentation. To better balance model-driven and data-driven methodologies, ref [30] introduced the normalized Intersection over Union (nIoU) as a successor for the IoU. It has to be defined as

where $N$ is still the total sample number of the testing sets. $\frac{{TP\left[i\right]}}{{T\left[i\right]+P\left[i\right]-TP[i]}}$ stands for the IoU of each sample. $T$, $P$, $TP$ are same to Equ. 9, and $i$ in $TP[i]$, $T[i]$, $P[i]$, is the $i$-th sample.

3) Receiver Operating Characteristic (ROC) curve: ROC presents the shifting trends between false positive rate (FPR) and true positive rate (TPR). Unlike IoU, which only reflects the segmentation effect under a fixed threshold, it indicates the total effect under a sliding threshold.

where $TP$, $P$ are same to Equ. 9 and Equ. 10. $FP$ is false positive pixel numbers, and $thr$ stands for the threshold variable.

In this section, two ablation experiments for the two infrared small object datasets are listed to assess the effectiveness of the proposed UIU-Net. In detail, we investigated 1) the comparative detection performance of UIU-Net under different backbone networks, and 2) the individual contributions of each module in UIU-Net. For each part of the ablation study, we rigorously retrained the whole UIU-Net with the same parameter settings.

1) the comparative detection performance of UIU-Net under different backbone networks. We compare the IoU and nIoU values of U-Net with a classical residual backbone and multi-scale ReSidual U-block module in Table I. It shows that the IoU value increased from 0.6178 to 0.7825 on the SIRST dataset and from 0.4335 to 0.4773 on the synthetic dataset. For the nIoU value, the SIRST dataset improves by $\sim$ 0.12 to 0.7515, and the synthetic dataset improves by $\sim$ 0.07 to 0.4721. The above analysis clearly demonstrates the efficacy of the multi-scale residual backbone for infrared small object detection, and its ability to learn deep high-resolution features can greatly compensate for tiny object loss or poor feature representation.

2) the individual contributions of each module in UIU-Net. We compare the IoU and nIoU values of two incremental in Table II. For the incremental module (U-Net + RSU), the IoU value is increased from 0.7330 to 0.7825 on the SIRST dataset and from 0.4100 to 0.4773 on the synthetic dataset. For the nIoU value, the SIRST dataset improves by $\sim$ 0.05 to 0.7515, and the synthetic dataset improves by $\sim$ 0.09 to 0.4721. The above analysis clearly demonstrates the efficacy of the incremental module for the proposed UIU-Net. It also indicates that integrating interactive-cross encoder to deep high-resolution and multi-scale features is very necessary.

Table III lists the detection accuracy of nine infrared detection methods, including infrared patch image (IPI) [7], local contrast measure (LCM) [8], reweighted infrared patch-tensor (RIPT) [14], multiscale gray difference weighted image entropy (MGDWIE) [11], nonnegativity-constrained variational mode decomposition (NVMD) [15], partial sum of tensor nuclear norm (PSTNN) [13], novel local contrast descriptor (NLCD) [12], asymmetric contextual modulation (ACM) [30], Miss Detection vs. False Alarm-Conditional generative adversarial network (MDvsFA-cGAN) [45], and lightweight convolutional neural network (TBC-Net) [27]. in terms for two commonly-used indices IoU and nIoU, while Fig. 6 shows four example methods of the corresponding detection results and the reference label.

1) Quantitative Comparison. Overall, benefiting from the feature representation ability, the detection results using a data-driven network are dramatically higher than those model-driven methods, and the proposed data-driven network outperforms other compared methods, which obtains the best IoU and nIoU values in infrared small object detection. For the SIRST dataset with small and real data quantities, the UIU-Net learned deep multi-scale and high-resolution features and focused on the global and local contrast. This, to some extent, yields better detection performance and lower miss detection rates. ACM and TBC-Net have relied on the classification backbone trained by the ImageNet dataset, which makes the network mainly focus on the objects in SIRST with the same distribution as ImageNet rather than on the SIRST data itself, thus limited to the detection performance and also failing to fundamentally improve the detection performance. MDvsFA-cGAN model has balanced the missed detection rate and false alarm rate by adversarial generation network. However, the higher complexity of the MDvsFA-cGAN model cannot be applied to the small data of SIRST, making IoU and nIoU values lower than other data-driven methods.

Different from the IoU and nIoU values with fixed thresholds, we also present the ROC curves for the above methods with dynamic thresholds in Fig. 7. As can be observed, the proposed method still performs the best. More specifically, IPI has displayed in a specific coordinate interval due to the low quality of the data.

2) Visual Comparison. Some visualization examples of the SIRST dataset have shown in Fig. 6, including single objects and multiple objects. Obviously, the model-driven methods with limited feature expressiveness yielded a high rate of missing detection. The visual effect is that detected objects have fewer pixel values than the reference label. The other three data-driven methods with feature auto-learning obtained better detection results, which is consistent with the quantitative results in Table III. Specifically, the proposed method yielded optimal visual results, particularly for tightly connected multiple objects (see the blue circles in the third and fourth row of Fig. 6). But for ACM and MDvsFA-cGAN, their visual results show adhesion or pixel loss (as shown in the second and third-column in Fig. 6). Especially for MDvsFA-cGAN, which is trained by a small dataset (SIRST), the visual results are even worse than the model-based PSTNN method.

Similarly, the quantitative detection performance on the synthetic datasets are listed in Table IV, and the corresponding visualization results are shown in Fig. 8.

In addition, we have added an experiment to verify whether UIU-Net is overly dependent on the size of the dataset. To answer this question, two sub-experiments are designed. Specifically, 1) large-scale infrared training dataset: $6,200$ images in the configuration I am chosen and trained after manually removing the ”poor” images. Here, the ”poor” images are those that have null values or negative values after pre-processing. 2) small-scale infrared training dataset: $500$ images in 1) are randomly selected and reconstructed into a small sample data. similar to the SIRST data, the final detection output is the smallest outer rectangle of the optimal segmentation result, and the quantitative metrics for the comparing methods are IoU, nIoU, and ROC curves.

1) Quantitative Comparison. To begin, we have given the longitudinal comparison between datasets. Compared to the real SIRST dataset as train data, the quantitative results for the synthetic dataset in Table IV has demonstrated the worse IoU and nIoU value. There are two main reasons for this: 1) resolution gap: the synthetic dataset has a resolution of $128\times 128$, which is lower than the SIRST dataset. Feature distinction and context representation ability are limited due to the lower image resolution. 2) reference labeling: the object labeling range in the synthetic data was erroneously inflated, and this labeling way is especially unfriendly to the quantitative evaluation of pixel-based methods. But the referenced labeling in the SIRST data is more closely matched to the real objects.

Secondly, we have shown a data volume comparison. The ratio of large-scale to small-scale samples in the synthetic dataset is $12:1$. Although the IoU and nIoU for the proposed UIU-Nets have been decreased by about $4\%$, they are still optimal. Furthermore, the results from the both two small-scale datasets, including SIRST and small-scale synthetic data, show that UIU-Net is robust for small-scale data. It also indirectly indicates that UIU-Net is extremely friendly to practical applications.

Finally, the ROC curves for the above-comparison approaches are shown in Fig. 10, the proposed method still performs the best. Unfortunately, due to the complicated background and object variability, as well as insufficient image quality, none of the model-driven methods can be exhibited within the coordinate range. The detection performance of the MDvsFA-cGAN with underfitting for a small sample dataset is equally dismal. This emphasizes how important data volume and model fit are in data-driven methods.

2) Visual Comparison. For the visualization results in Fig. 8, the proposed UIU-Net has the highest similarity with the reference label for the object indicated in the blue rectangular box. However, as shown in the fourth line of Fig. 8, some images have some missing detection even though the accuracy is assured. For other data-driven methods, the ACM output objects introduced erroneous bounds in comparison to its reference label. Due to the inaccuracy of the reference labels, MDvsFA-cGAN is unable to adequately balance the false alarm rate and missed alert rate, even with large-scale samples (as seen in the sixth column of Fig. 8). For the model-driven methods, PSTNN achieves a rough object location, and MGDWIE has a higher missed detection, false detection, and lower accuracy.

This section employs the detection model trained with SIRST true infrared small objects to verify the generalization capability of UIU-Net. It primarily consists of two parts of experiments.

1) Synthetic data: As the same with section IV-E, 100 test images have been selected to verify the generalization capability of UIU-Net, as shown in Table V. It shows that the quantitative results of UIU-Net trained by $408$ SIRST data exceed the model trained by 500 synthetic data, although being lower than the results of training by $6,200$ synthetic data. In addition, it exceeds the other compared methods in Table IV, no matter $6,200$ or $500$ trained data.

2) ATR Sequential Data: Due to the shortages of public single-frame infrared small object datasets, we employ ATR sequential data as the test data to verify the generalization capability of UIU-Net. This inspired a question: how does the proposed UIU-Net perform in terms of the video sequence dataset? We sample data2, data6, and data7 segments for visualization validation due to the limits of the center point label, as shown in Fig. 9, where we display a detection result every 50 frames. For data2, UIU-Net has good detection performance. All visualization results except for $\#0001$ are well for cross-flying UAVs in a close and sky background (see rows 1-2 in Fig. 9). Due to the intricacy of the ground background and the failure to consider timeliness, the missed detection is more problematic for the UAV flying near the original ground backdrop in data6 (see the rows 3-4 in Fig. 9). In contrast to the previous two data segments, the ground change of data7 will invariably contain flash elements, interfering with the UAV detection (see rows 5-6 in Fig. 9).

To sum up, the UIU-Net model trained on the SIRST dataset shows superior generalization performance for small object detection in three video data segments. Currently, there are some other works in the field of saliency object detection, e.g., Contrast based filtering [47], Context-aware saliency [48], spectral residual method [49], etc., which can also achieve the high generalization ability in the video object detection.

Currently, detecting unknown scenes or objects is still a challenging and noteworthy research problem. There have been some works, e.g., background subtraction[50], scene independent end-to-end spatiotemporal feature learning framework[51], meta-knowledge learning and domain adaptation[52], achieving the high robustness and generalization ability for unseen scenarios. Actually, the experiments shown in Fig. 9 are a case of object detection in the unseen video dataset (ATR ground/air video dataset). More specifically, the ATR dataset is an infrared detection and tracking dataset with one or more fixed-wing UAV objects that differ in size, interference, storage and mobility, and other properties (cf. the SIRST dataset). This naturally brings more unseen objects to be detected. Moreover, the image’s background is diverse, including the sky, the earth, and so on, leading to more unseen scenes as well (cf. the SIRST dataset).