Diff-Mosaic: Augmenting Realistic Representations in Infrared Small Target Detection via Diffusion Prior

The task of Single-Frame infrared small target (SIRST) detection is to locate abnormal targets in infrared images. Conventional detection methods achieve SIRST detection by comparing the difference between the target and the background. Tophat [10] and New Tophat [11] employ manually designed filters to identify high-response in the target region. The Local Contrast Measure (LCM) [13], WSLCM [14] and TCLCM [15] detect small targets by assessing the contrast between the central pixel and its surrounding neighborhood at various scales. Infrared Patch Image (IPI) model [16] decomposes images into low-rank and sparse signals to achieve SIRST detection. However, due to their inability to capture deep features of images, these methods face challenges in detecting anomalous targets within images characterized by complex backgrounds. To address this issue, researchers have employed Convolutional Neural Networks (CNNs) based method to systematically acquire profound features of infrared images through a data-driven approach.

Unlike traditional methods, CNN-based methods do not rely on any prior knowledge about the target and background. Instead, they learn target features directly from the data through the model. This advantage makes CNN-based methods widely used in downstream tasks of remote sensing, such as classification tasks [26, 27, 28, 29], segmentation tasks [30, 31, 32], and object detection tasks [33]. CNN-based methods can efficiently process large amounts of data and have strong model-fitting capabilities, thus enabling accurate SIRST detection. For example, ALC-Net [19] based on a feature pyramid network proposes an attention local contrast (ALC) network based on the attention mechanism to segment the target region in infrared images. ACM [17] proposes an asymmetric context modulation module and integrates it into U-Net [34] and FPN [35]. DNA-Net  [20] uses a dense nested interaction module. Multiple interactions of deep semantic features and shallow detailed features are performed through this module to maintain deep infrared small targets. UIU-Net[21] adopts a U-Net in U-Net architecture, consisting of two U-Net components. These two U-Net modules are dedicated to feature extraction and feature fusion, utilizing residual connections and skip connections to retain the details of small targets.

Recently, Diffusion Denoising Probabilistic Models (DDPM) [36] have made significant progress in computer vision, particularly in image generation. DDPM is divided into two processes: the diffusion process and the reverse process. During the diffusion process, diffusion models progressively introduce Gaussian noise to systematically undermine the input information. This gradual introduction of noise into the modeling process makes it easier to capture the diversity within images, thereby enhancing the richness and realism of the images generated by DDPM. During the reverse process, DDPM then restores the original information by predicting and gradually removing the increased noise. This step-by-step denoising strategy helps the model more effectively reconstruct the original information, particularly demonstrating impressive performance when confronted with complex noise and image degradation scenarios. Thus DDPM has achieved impressive results in image generation. To further improve the efficiency and quality of image generation, researchers have proposed the Latent Diffusion Model (LDM) [37]. LDM improves the efficiency of model training and inference by using autoencoder [38] to compress the information to be learned into a latent space. The large-scale trained LDM model not only generates satisfactory images, but also incorporates real-world information. Therefore, LDM can serve as a powerful generative prior to image enhancement by incorporating real-world information. This diffusion-based generative prior has found applications in various domains. Generative Diffusion Prior [24] and DiffBIR [39] leverage the Diffusion Prior to achieve high-quality blind image restoration and image enhancement. ControlStyle [25] integrates textual and visual information using Diffusion Prior, accomplishing high-quality text-driven stylized image generation.

As discussed in Section  I, the scale of publicly available datasets [17, 20] for SIRST detection is limited. Furthermore, the data augmentation methods employed to expand datasets suffer from issues of both limited diversity and a lack of realism in SIRST detection. As shown in Fig. 1, we compare the generation quality of our method with the traditional data augmentation method Mosaic. It can be observed that the images generated by Mosaic look very unnatural at the splices. The four sub-images used for the stitching are not coordinated in terms of brightness and contrast, resulting in an overall lack of realism. In addition, these augmented images are derived from existing infrared datasets, which lack diversity. In order to address this issue, we propose a novel data augmentation network named Diff-Mosaic.

In this work, we aim to utilize a powerful diffusion generative prior to generating realistic and diverse augmentation samples. The Diff-Mosaic framework is divided into two aspects: pixel harmonization and detail refinement, the framework is shown in Fig. 3. Diff-Mosaic has two stages: the training stage and the data generation stage. In the training stage, the pipeline of Diff-Mosaic is organized as follows:

• •

  Pixel-prior Machine: We propose an enhancement network that boosts the quality and realism of images generated by Mosaic. First, the input image $I_{input}$ was transformed into $I^{\prime}_{degrade}$ by applying a degradation module. Then, a cut-and-paste operation was performed to generate a mixed image $I^{\prime}_{mix}$. Finally, the mixed image $I^{\prime}_{mix}$ was fed into the enhancement network for training to generate a harmonized image $I^{\prime}_{smooth}$.

• •

  Diff-prior Machine: A powerful generative prior is applied to resample the image $I^{\prime}_{smooth}$ using the diffusion model, thereby incorporating real-world information. The resampled image $I^{\prime}_{realis}$ not only contains richer details, but also incorporates real-world knowledge and information from pre-trained diffusion models. Therefore the resampled image $I^{\prime}_{realis}$ is more realistic and diverse. To make the generated resampled samples more realistic, we fine-tune the diffusion model.

In the data generation stage, we use the Mosaic augmentation to generate $I_{Mosaic}$. Subsequently, image degradation and cut-and-paste operations are performed on $I_{Mosaic}$ to generate the mixed image $I_{mix}$. And $I_{smooth}$ is more diverse and has higher image quality than the Mosaic image $I_{Mosaic}$. Finally, we resample $I_{smooth}$ to get the augmented samples $I_{realis}$ via diffusion model. These images will be input into the network as the augmented samples for training. We will describe each step in detail in the following subsections.

In order to solve the problem of uneven distribution existing in Mosaic images and further improve the quality of Mosaic images. We propose Diff-Mosaic to generate image-coordinated Mosaic images. As shown in Fig. 3, in the training stage, we trained a transformer network to enhance the quality of Mosaic images. Specifically, as shown in Fig. 4, we conducted degradation and cut-and-paste operations on the input image $I_{input}$ to produce the mixed image $I_{mix}$. This mixed image was then used for training in the enhancement network. The degradation process is expressed as follows:

where $\mathcal{D}egrade\left(\cdot\right)$ denotes the degradation module, which employs a variety of degradation sub-modules, including fuzzy, resize, and noise degradation methods. To improve the quality of the reconstructed image, we use the cut-and-paste operation after degradation. Specifically, we cut a part of the region of the original image and paste it into the corresponding region of the degraded image $I^{\prime}_{degrade}$. The operation is as follows:

where $\mathcal{P}(\cdot)$ denotes the cut-and-paste operation, and $M_{select}$ denotes the selection of the region from the original image $I_{input}$. The Cut-and-paste operation can improve the diversity and difficulties of the data, which makes the recovered images of the model of higher quality.

Then the mix image $I^{\prime}_{mix}$ is input into the transformer network for reconstruction. The transformer network consists of shallow feature extraction, deep feature extraction, and image reconstruction. The input image $I^{\prime}_{mix}$ is subjected to shallow convolutional layers to extract shallow features. Then it undergoes multiple Residual Swin Transformer Blocks (RSTB) to extract the deep features [40]. Where the RSTB structure consists of multiple Swin Transformer Layers (STL) which collaborate with each other to capture deeper features of the image. Next, we fuse the shallow and deep features to integrate high-frequency and low-frequency information. Finally, the deep features are reduced to the original image space by an upsampling to get the harmonized image $I^{\prime}_{smooth}$. To optimize the network, we compute the gap $\mathcal{L}_{har}$ between the original image $I_{input}$ and $I^{\prime}_{smooth}$. The loss $\mathcal{L}_{har}$ is expressed as follows:

As shown in Fig. 3, compared to Mosaic image $I^{\prime}_{mosaic}$ , $I^{\prime}_{smooth}$ effectively solves the distribution inconsistency problem and improves the image quality.

In the data generation stage, we have adopted the improved Mosaic process. The images are further enhanced to increase the diversity of the augmented samples. Specifically, first, we used Mosaic to stitch image $I_{input}$ with other images in the dataset to form $I_{Mosaic}$. Next, we perform a degradation operation on the image:

Mixed image $I_{mix}$ can improve the diversity of Mosaic, and the mixed image $I_{mix}$ is inputted into the image reconstruction network to generate the Mosaic image $I_{smooth}$ that gets the coordinated Mosaic image.

Compared to $I_{Mosaic}$, the Pixel-Prior machine-generated image $I_{smooth}$ exhibits a more harmonious and smoother texture. However, there is still a gap between $I_{smooth}$ and the real image. This discrepancy is due to the use of L2 loss in optimizing the enhancement network, which can magnify differences between minor features. Therefore, the Pixel-Prior machine focuses on the uniformity of the overall pixel values while disregarding the details within the image.

To address this problem, we introduce the diffusion model. Unlike traditional methods of data augmentation, the diffusion model employs a powerful generative prior to simulating real-world scenarios based on the distribution of image pixels. This allows diffusion models to incorporate real-world information and generate images with detailed and realistic approximations. In this paper, we use the Latent Diffusion Model(LDM) [37] as a generative prior.

LDM is an advanced diffusion model. It uses pre-trained autoencoders to map images to latent space for learning data distributions. The autoencoder is composed of an encoder $\mathcal{E}$ and a decoder $\mathcal{D}$, where $\mathcal{E}$ encodes input images $I$ into latent code $z$ and $\mathcal{D}$ reconstructs it. In this model, noise is added $T$ times to latent variables $z$, generating Gaussian-distributed codes $z_{T}$. Noise predict network $\epsilon_{\theta}$ is trained to denoise at time step $t$, optimizing with $\mathcal{L}_{ldm}$. $\mathcal{L}_{ldm}$ is calculated as follows:

where $N_{gaussian}$ is the Gaussian noise and $z_{t}$ is the latent noise at time step $t$. The large-scale trained diffusion model [37] demonstrates proficiency in comprehending various attributes within images, such as object shapes and textures, and is thus able to generate visually appealing images. To further generate more high-quality and realistic images, we fine-tuned the diffusion model on the SIRST datasets.

As shown in Fig. 3, given $I^{\prime}_{smooth}$, in the training phase we encode the image $I^{\prime}_{smooth}$ into the latent space coding $z_{0}=\mathcal{E}(I^{\prime}_{smooth})$ via the encoder $\mathcal{E}$. Then, the $z_{0}$ model is subjected to a $T$-step denoising into a noise $z_{T}$ close to a Gaussian distribution, and then $z_{0}^{\prime}$ is generated by $T$-step resampling. Finally, the decoder $\mathcal{D}$ generates the resampled image $I^{\prime}_{realis}$. To bring the results of the diffusion model resampling closer to our dataset, we reduce the gap between $z_{0}^{\prime}$ and $z_{0}$ by $\mathcal{L}_{realis}$. $\mathcal{L}_{realis}$ is expressed as follows:

During the data generation stage, images $I_{realis}$ resampled by the fine-tuned LDM have much finer details and are very closely aligned with the distribution of the original dataset. In addition, augmentation sample $I_{realis}$ is more diverse and realistic due to the implementation of generative prior. Notably, this method can capture latent data features without requiring additional labels as input. Therefore, it conveniently produces more realistic and diverse samples for the dataset reconstruction.

In the data generation stage. By employing diffusion prior, we resample $I_{smooth}$ to generate $I_{realis}$.

As shown in Fig. 1, compared to the Mosaic image, $I_{realis}$ becomes more coherent. In addition, the generated result $I_{realis}$ is very different from other images in the dataset due to the fact that it is extracted information from the real world and generated by the diffusion model, which significantly enhances the diversity of the dataset. The $I_{realis}$ generated in the data generation stage will be involved in the training as the augmented sample.

To demonstrate the advantages of our proposed data augmentation method, we utilize a standard detection network for inference. The structure of the detection network is shown in Fig. 5. Firstly the image $I_{reails}$ is fed into the residual attention block to extract the features. The residual attention block is composed of two convolutional layers, channel attention module, and spatial attention module. These two attention modules are designed to enhance the feature information after passing through two convolutional layers. In order to fuse different features, the features are continuously upsampled and downsampled into different residual attention blocks to form a feature pyramid. Finally, these features containing different scale information are fused to generate a robust feature map for predicting the result $\hat{M}$. Both the real target result $M$ and the predicted result $\hat{M}$ are binary images. In order to get the results close to ground truth $M$, the network reduces the gap between $\hat{M}$ and the ground truth $M$ by loss $\mathcal{L}_{iou}$, which is expressed as follows:

To evaluate the detection performance of the CNN-based approach, we use three different evaluation metrics to assess the network performance.

• •

  Intersection over Union (IoU): IoU is a classical pixel-level semantic segmentation evaluation metric to describe the algorithm’s contour description capability. It is defined as the ratio of intersection and concatenation area between predicted values and labels as follows:

  $$IoU=\dfrac{A_{inter}}{A_{Union}}$$

  (8)

  where $A_{inter}$,$A_{Union}$ denotes the interaction region and union region, respectively.

• •

  Probability of detection($P_{d}$ ): The probability of detection is a target-level evaluation metric. It measures the ratio of the number of correctly predicted targets $T_{correct}$ to the number of all targets $T_{All}$. The definition of $P_{d}$ is as follows:

  $$P_{d}=\dfrac{T_{correct}}{T_{All}}$$

  (9)

• •

  False Alarm Rate($F_{a}$) : False Alarm Rate is another target-level evaluation metric. It is used to measure the ratio of incorrectly predicted pixels $P_{false}$ to all image pixels $P_{ALL}$. The definition of $F_{a}$ is as follows:

  $$F_{a}=\dfrac{P_{false}}{P_{All}}$$

  (10)

We compared state-of-the-art methods with our method on SIRST and NUDT-SIRST. And three metrics, IoU, $P_{d}$, and $F_{a}$, were used for evaluation. As shown in Tab.  I, our method achieved the best results for the IoU, $P_{d}$, $F_{a}$ metrics on SIRST and NUDT-SIRST. To highlight the superior performance of our method, we visualize our method with other detection models. To highlight the superior performance of our method, we use a red dashed box to circle the target area and zoom in to show it. To more clearly compare the difference between the model detection results and ground truth, we use red pixels to indicate the difference between the method and ground truth. In addition, we marked the regions of model error detection with yellow dashed boxes.

As shown in Fig. 6, we visualize the performance comparison between the state-of-the-art method and our method on general and difficult samples. Where difficult samples are uncommon instances in the dataset. They usually have large dimensions, multiple objects or irregular contours. General samples are usually small, singular, and have regular shapes. ACM and ALC-Net had many false alarms. UIU-Net and DNA-Net  [20] show fewer wrongly predicted regions, however, they fail to accurately detect the target in some of the difficult samples (the last two samples). In contrast, detection model trained with the samples generated by Diff-Mosaic performed well in predictions dealing with difficult samples. This suggests that Diff-Mosaic can generate more challenging samples for discriminative training in the detection model.

As shown in Tab. II, we list the computational parameter count and inference time of CNN-based methods. Our approach generates augmented samples using pixel-prior and diffusion-prior mechanisms. This process happens during data generation and produces one sample every 10 seconds without affecting training and inference times. Thus our method is competitive with other methods in terms of the parameter number and inference time.

We analyzed the effect of different numbers of samples generated by Diff-Mosaic as augmentation samples on the performance of the detection model. As shown in Tab.  III, we compare the performance of the models trained with real data alone and when augmented with 125, 250, and 400 additional samples. With an increase in the number of augmentation samples, the performance of the model, as measured by the $F_{a}$, $P_{d}$ and IoU metrics, shows an improvement. This demonstrates the diversity and realism of the augmented samples generated by Diff-Mosaic can improve the robustness of the detection model.

The ablation studies in this section were performed on NUDT-SIRST using the baseline detection Network. We compare the performance of models trained with Mosaic, with Pixel-Prior (‘w/ Pixel-Prior’), and with Diff-Mosaic (‘w/ Pixel-Prior+Diff-Prior’). The comparative results are shown in Tab. IV. Due to the lack of realism and diversity of the samples generated by ‘w/Mosaic’, the improvement of the detection network performance is very limited. And ‘w/ Pixel-Prior’ harmonizes the augmented samples, making them more realistic and therefore improving on all three performance metrics. At last, ’w/ Pixel-Prior+Diff-Prior’ introduces real-world information and generates diverse and realistic augmented samples. Aided by ’w/ Pixel-Prior+Diff-Prior’, the detection network shows significant improvement in all three performance metrics.

To demonstrate the effectiveness of the augmentation samples generated by Diff-Mosaic, we used the generated results of Diff-Mosaic as augmentation samples for training on different detection models. In this section, we compare the effects of training without (‘w/o Diff-Mosaic’) and with Diff-Mosaic’s (‘w/ Diff-Mosaic’) augmented samples using UIU-Net and DNA-Net as detection models. As shown in Tab. V, we compare the performance of the UIU-Net model with and without the augmented samples. To highlight the gaps, we use red bold font to indicate how much ‘w/Diff-Mosaic’ improves on IoU and blue bold font to indicate how much ‘w/ Diff-Mosaic’ decreases on $F_{a}$. It can be seen that the performance of UIUNet improves significantly on the NUDT-SIRST dataset after training with Diff-Mosaic augmented samples. As shown in Fig. 7, it can be seen that ‘w/o Diff-Mosaic’ has a lot of misdirected targets in its detection results and is inaccurate in detecting the contours of small targets. Furthermore, the detection result of ‘w/ Diff-Mosaic’ has no false detection targets and achieves accurate detection of the contours of small irregular targets.

As shown in Tab. VI, we compared the performance of the DNA-Net model with and without the augmented samples generated by Diff-Mosaic. The results indicate a significant improvement in detection performance on the NUDT-SIRST dataset when using augmented samples. As depicted in Fig. 8, it is evident that the ‘w/o Diff-Mosaic” method fails to detect the target in challenging test samples, while the ‘w/ Diff-Mosaic’ method accurately identifies the target with high precision. To show the realism of the images we generate, we input the results of the pixel-prior machine into state-of-the-art image generation techniques, including SwinIR [41], FemaSR [42] and DiffBIR [39], and compare their generated results with our method. As shown in Fig. 9, it can be observed that augmented samples generated by SwinIR, FeMaSR, and DiffBIR fail to integrate the Mosaic strategy well. On the contrary, the augmented samples generated by our diff-mosaic demonstrate a seamless coherence in image distribution and a high level of realism. Additionally, to measure the realism of the generated images, we utilized KID [43] and FID [44] metrics to measure the disparity between them and real infrared images. As shown in tab. VII, our method outperforms others in terms of FID and KID, indicating a significant advantage in enhancing realism.

To demonstrate the effectiveness of diffusion-prior machine, we combined it with CutMix [22]. Subsequently, we generated 400 augmented samples on NUDT-SIRST and input them into the detection network for training. As depicted in Fig. 10, we demonstrate the augmented samples from the original CutMix (‘w/ CutMix’) and the combination with our method (‘w/ CutMix + diffusion prior’). It is evident that the results generated by the original CutMix lack coherence and realism. In contrast, the results generated by ‘w/ CutMix + diffusion prior’ exhibit greater coherence. To validate the effectiveness of our diffusion prior, we compared the performance of the augmented samples generated by ‘w/ CutMix’ and ‘w/ CutMix + our diffusion prior’. As shown in Tab. VIII, it can be observed that ’w/ CutMix’ shows a negligible improvement in the network. In contrast, ’w/ CutMix + diffusion prior’ introduces real-world information, leading to a noticeable improvement in the network’s detection capability.