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

Recently, researchers have proposed various deep learning methods to accurately detect infrared targets with the characteristics of indistinct shape and texture. Due to the limited variety of infrared datasets, training deep learning models with good generalization poses a challenge. To augment the infrared dataset, researchers employ data augmentation techniques, which often involve generating new images by combining images from different datasets. However, these methods are lacking in two respects. In terms of realism, the images generated by mixup-based methods lack realism and are difficult to effectively simulate complex real-world scenarios. In terms of diversity, compared with real-world scenes, borrowing knowledge from another dataset inherently has a limited diversity. Currently, the diffusion model stands out as an innovative generative approach. Large-scale trained diffusion models have a strong generative prior that enables real-world modeling of images to generate diverse and realistic images. In this article, we propose Diff-Mosaic, a data augmentation method based on the diffusion model. This model effectively alleviates the challenge of diversity and realism of data augmentation methods via diffusion prior. Specifically, our method consists of two stages. First, we introduce an enhancement network called Pixel-Prior, which generates highly coordinated and realistic Mosaic images by harmonizing pixels. In the second stage, we propose an image enhancement strategy named Diff-Prior. This strategy utilizes diffusion priors to model images in the real-world scene, further enhancing the diversity and realism of the images. Extensive experiments have demonstrated that our approach significantly improves the performance of the detection network. The code is available at https://github.com/YupeiLin2388/Diff-Mosaic.

I. Introduction

Single-frame infrared small target (SIRST) detection has a wide range of applications in several fields, such as video surveillance [2], [3], [4], early warning systems [5], raining scenes [6], [7], [8], and military surveillance [9]. Accurately detecting small infrared targets is crucial for ensuring safety, proper navigation, and successful mission execution. However, infrared small target detection faces multiple challenges. First, due to the limited number of pixels occupied by small targets in infrared images, they are easily overwhelmed by complex backgrounds and noise, making them difficult to detect. Second, infrared small targets often lack distinct shapes and textures, making their detection more challenging. Additionally, these small targets exhibit weak features in the images, requiring overcoming the impact of lighting variations and other environmental factors for accurate detection. Therefore, achieving accurate detection of infrared small targets in the face of these issues remains a challenging problem.