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

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

Diff-Mosaic：通过扩散先验增强红外小目标检测中的真实表征

摘要

近年来，研究人员提出了多种深度学习方法来精确检测具有形状和纹理不清晰特征的红外目标。由于红外数据集的多样性有限，训练具有良好泛化能力的深度学习模型面临挑战。为了扩充红外数据集，研究人员采用数据增强技术，这些技术通常涉及通过组合不同数据集的图像来生成新图像。然而，这些方法在两个方面存在不足。在真实性方面，基于混合的方法生成的图像缺乏真实性，难以有效模拟复杂的现实场景。在多样性方面，与现实世界场景相比，从另一个数据集借鉴知识本质上多样性有限。目前，扩散模型作为一种创新的生成方法脱颖而出。大规模训练的扩散模型具有强大的生成先验，能够对图像进行真实世界建模，生成多样且真实的图像。在本文中，我们提出了 Diff-Mosaic，一种基于扩散模型的数据增强方法。该模型通过扩散先验有效缓解了数据增强方法在多样性和真实性方面的挑战。具体而言，我们的方法包括两个阶段。首先，我们引入了一个名为 Pixel-Prior 的增强网络，通过协调像素生成高度协调和真实的马赛克图像。在第二阶段，我们提出了一种名为 Diff-Prior 的图像增强策略。该策略利用扩散先验对现实世界场景中的图像进行建模，进一步增强图像的多样性和真实性。大量实验表明，我们的方法显著提高了检测网络的性能。代码可在https://github.com/YupeiLin2388/Diff-Mosaic获取。

关键词：红外小目标检测，数据增强，马赛克增强，扩散模型。

一、引言

单帧红外小目标（SIRST）检测在多个领域有着广泛的应用，如视频监控 [2]-[4]、预警系统 [5]、雨天场景 [6]-[8] 和军事监控 [9]。准确检测红外小目标对于确保安全、正确导航和成功执行任务至关重要。然而，红外小目标检测面临多个挑战。首先，由于红外图像中的小目标所占据的像素数量有限，它们容易被复杂的背景和噪声淹没，难以检测。其次，红外小目标通常缺乏明显的形状和纹理，使其检测更具挑战性。此外，这些小目标在图像中表现出较弱的特征，需要克服光照变化和其他环境因素的影响才能进行准确检测。因此，在面对这些问题时实现红外小目标的准确检测仍然是一个具有挑战性的问题。

图 1：我们比较了传统马赛克 [1] 与我们方法的效果。为了强调我们方法生成的样本的多样性，在组合其余四张图像时，我们固定了马赛克的右上角图像。马赛克生成的样本具有碎片化的质量，无法像一个完整的图像。相比之下，Diff-Mosaic 具有均匀的分布和协调的灰度。特别是在红色圆圈标记的红外小目标方面，Diff-Mosaic 的结果能够更好地增强多样性和真实性。

为了实现红外小目标检测，研究人员提出了许多传统方法。这些传统方法包括滤波、局部对比度和基于低秩的方法。Tophat [10] 和 New Tophat [11] 使用人工设计的滤波器从红外图像中选择性地提取视觉上突出的目标。二维最小二乘滤波器 [12] 通过估计周围像素来预测背景，并通过比较预测背景与红外图像之间的差异来检测目标。受人类眼睛启发的局部对比度测量（LCM）[13]、WSLCM [14] 和 TCLCM [15] 利用局部对比度来捕获目标特征。红外斑块图像（IPI）模型 [16] 将小目标检测建模为涉及低秩（背景）和稀疏信号（目标）的分解问题。然而，这些方法需要选择适当的特征和精心调整的超参数，难以适应现实世界场景中的各种噪声。因此，这些方法会产生许多误报。为了处理现实世界场景中的各种数据集，研究人员因此转向使用基于 CNN 的网络。

与传统方法不同，基于 CNN 的方法可以从数据集中自动提取和学习红外小目标的特征，更适合复杂多变的红外小目标检测任务。非对称上下文模块（ACM）[17] 通过整合浅层和深层特征信息来准确检测红外小目标。MDvsFA-GAN [18] 采用对抗性神经网络方法，通过生成器学习数据分布来实现红外目标识别。注意力局部对比度网络（ALC-Net）[19] 通过将低级细节信息嵌入到高级中实现了优异的检测性能。DNA-Net [20] 使用密集嵌套交互模块实现准确的目标检测。该模块使高级和低级特征之间能够进行渐进式交互，从而在深层中保留小目标的信息。UIU-Net [21] 使用由两个 U-Net 组成的 “U-Net 中的 U-Net” 结构。这两个 U-Net 分别用于特征提取和特征融合，并使用残差连接和跳跃连接来保留小目标的细节。

这些方法通过改进模型来提高检测性能，但没有考虑提出新的数据增强方法这一方面。在 SIRST 检测中，研究人员通常采用典型的数据增强技术，如马赛克 [1]、Cut-Mix [22] 和 Mixup [23]。这些技术涉及混合来自不同真实图像的信息以生成新图像。然而，这些数据方法并非专门为红外数据集量身定制，因此在 SIRST 检测中带来两个挑战：

• 多样性：这些主流的数据增强方法，如马赛克和 Cut-Mix，通过混合现有数据生成新图像。然而，受限于样本数量有限，这些数据增强方法生成的图像多样性有限。

• 真实性：如图 1 所示，由于红外图像之间的对比度变化显著，马赛克生成的样本的像素分布不一致，导致缺乏真实性。

这些数据增强方法的局限性显著影响了模型的性能。因此，提高增强样本的多样性和真实性是一项迫切需要关注的关键任务。近年来，去噪扩散概率模型（DDPM）由于其强大的生成先验，在低级领域引起了相当大的兴趣。扩散模型使用强大的生成先验，基于图像像素的分布对现实世界场景进行建模。扩散模型生成的图像在细节和真实性方面有显著提高。生成扩散先验 [24] 基于其扩散生成先验，在去模糊和图像着色方面表现出色。ControlStyle [25] 通过扩散先验整合文本和视觉信息，实现高质量的文本驱动风格化图像生成。

受这些方法的启发，在本文中，我们提出了 Diff-Mosaic，一种基于扩散模型的新型数据增强方法。我们的方法采用扩散生成先验将现实世界信息整合到图像中，从而增强生成样本的多样性和真实性。具体而言，我们在马赛克数据增强的基础上提出了一个增强网络 Pixel-Prior，它有助于图像像素的协调。该网络可以在没有额外标签的情况下协调图像，生成高质量、真实的马赛克图像。此外，我们引入了扩散模型 Diff-Prior，它通过使用扩散生成先验对 Pixel-Prior 的结果进行重采样来整合现实世界信息。Diff-Prior 在重采样中融入了现实世界知识，这使得生成的样本更加真实和多样。最后，我们通过将我们的方法应用于最先进的方法并与其他检测基线进行比较，验证了我们方法的有效性。如图 2 所示，我们比较了我们的方法与其他方法的可视化结果。我们的方法实现了准确的目标检测，没有误报。我们方法的贡献如下：

图 2：我们展示了我们的方法与不同的 SIRST 检测方法在两个 SIRST 数据集上的视觉结果。目标区域、放大的目标区域、错误预测区域和与地面实况（GT）不同的预测分别用红色虚线圆圈、红色框、黄色虚线圆圈和红色像素标注。与其他存在误报和与地面实况不一致的基线不同，我们的方法实现了准确的目标检测，没有误报。

• 在本文中，我们引入了一个增强网络，在像素级别协调图像。它在没有额外标签的情况下生成大量有效的增强样本，克服了传统数据增强方法固有的真实性不足的问题。

• 我们引入了一种基于扩散模型先验的图像重采样策略。该方法使用扩散模型先验对图像的各个方面进行真实建模，包括物体形状、纹理和光照，从而确保生成结果的真实性和多样性。此外，这种方法确保我们的方法是第一个将扩散先验引入 SIRST 的方法。

• 我们将我们提出的数据增强方法应用于最先进的基线。通过三个评估指标验证了我们方法的有效性。消融研究表明，所提出方法的每个部分都能提高检测性能。

本文的组织结构如下。第二部分简要回顾相关工作。第三部分描述所提出模型的流程。第四部分进行了大量实验和消融实验，以证明我们方法的可靠性。第五部分总结了本文。

三、方法

A. 观察

如第一部分所讨论，用于 SIRST 检测的公开可用数据集 [17]、[20] 的规模有限。此外，用于扩充数据集的数据增强方法在 SIRST 检测中存在多样性有限和缺乏真实性的问题。如图 1 所示，我们比较了我们的方法与传统数据增强方法马赛克的生成质量。可以观察到，马赛克生成的图像在拼接处看起来非常不自然。用于拼接的四个子图像在亮度和对比度方面不协调，导致整体缺乏真实性。此外，这些增强图像来自现有的红外数据集，缺乏多样性。为了解决这个问题，我们提出了一种新型的数据增强网络，名为 Diff-Mosaic。

在这项工作中，我们旨在利用强大的扩散生成先验来生成真实且多样的增强样本。Diff-Mosaic 框架分为两个方面：像素协调和细节 refinement，该框架如图 3 所示。Diff-Mosaic 有两个阶段：训练阶段和数据生成阶段。在训练阶段，Diff-Mosaic 的流程如下：

Pixel-prior 机器：我们提出了一个增强网络，用于提高马赛克生成的图像的质量和真实性。首先，输入图像\(I_{input}\)通过应用退化模块转换为\(I_{degrade}'\)。然后，执行剪切和粘贴操作以生成混合图像\(I_{mix}'\)。最后，将混合图像\(I_{mix}'\)输入到增强网络中进行训练，以生成协调的图像\(I_{smooth}'\)。

Diff-prior 机器：应用强大的生成先验，使用扩散模型对图像\(I_{smooth}'\)进行重采样，从而整合现实世界信息。重采样的图像\(I_{realis}'\)不仅包含更丰富的细节，还整合了来自预训练扩散模型的现实世界知识和信息。因此，重采样的图像\(I_{realis}'\)更加真实和多样。为了使生成的重采样样本更真实，我们对扩散模型进行了微调。

在数据生成阶段，我们使用马赛克增强来生成\(I_{Mosaic}\)。随后，对\(I_{Mosaic}\)执行图像退化和剪切粘贴操作以生成混合图像\(I_{mix}\)。并且\(I_{smooth}\)比马赛克图像\(I_{Mosaic}\)更具多样性且具有更高的图像质量。最后，我们通过扩散模型对\(I_{smooth}\)进行重采样，以获得增强样本\(I_{realis}\)。这些图像将作为增强样本输入到网络中进行训练。我们将在以下小节中详细描述每个步骤。

图 3：Diff-Mosaic 的框架概述。我们展示了 Diff-Mosaic 的训练阶段和数据增强阶段之间的工作流程。在训练期间，对图像\(input\)执行剪切和粘贴操作以获得\(I_{mix}'\)。随后，将\(I_{mix}'\)输入到增强网络中以生成协调的图像\(I_{smooth}'\)。最后，将\(I_{smooth}'\)输入到扩散模型中进行训练，以获得细节丰富的图像\(I_{realis}'\)。在数据生成阶段，对图像\(I_{input}\)应用马赛克操作以产生马赛克图像\(I_{Mosaic}\)。随后，将\(I_{Mosaic}\)输入到 Pixel-Prior 机器中以生成\(I_smooth\)。最后，通过采用扩散先验，将真实但更丰富的表征整合到图像\(I_{smooth}\)的信息中，以生成视觉上更多样化和有纹理的图像\(I_{realis}\)。

B. 带有 Pixel-Prior 机器的马赛克

为了解决马赛克图像中存在的分布不均匀问题并进一步提高马赛克图像的质量。我们提出了 Diff-Mosaic 来生成图像协调的马赛克图像。如图 3 所示，在训练阶段，我们训练了一个 Transformer 网络来增强马赛克图像的质量。具体而言，如图 4 所示，我们对输入图像\(I_{input}\)进行退化和剪切粘贴操作，以产生混合图像\(I_{mix}\)。然后，该混合图像用于增强网络的训练。退化过程表示如下：

\(I_{degrade}' = \mathcal{D}egrade\left(I_{input}\right) \quad (1)\)

其中 Degrade (・) 表示退化模块，它采用多种退化子模块，包括模糊、 resize 和噪声退化方法。为了提高重建图像的质量，我们在退化后使用剪切和粘贴操作。具体而言，我们从原始图像中剪切一部分区域，并将其粘贴到退化图像\(I_{degrade}'\)的相应区域。该操作如下：

\(\begin{aligned} I_{mix}' &=\mathcal{P}\left(I_{input}, I_{degrade}', M_{select}\right) \\&= \left(1-M_{select}\right) \cdot I_{input}+M_{select} \cdot I_{degrad}' \end{aligned}\)

其中\(P(\cdot)\)表示剪切和粘贴操作，\(M_{select}\)表示从原始图像\(I_{input}\)中选择的区域。剪切和粘贴操作可以提高数据的多样性和难度，这使得模型的恢复图像质量更高。

然后，将混合图像\(I_{mix}'\)输入到 Transformer 网络中进行重建。Transformer 网络由浅层特征提取、深层特征提取和图像重建组成。输入图像\(I_{mix}'\)经过浅层卷积层以提取浅层特征。然后，它经过多个残差 Swin Transformer 块（RSTB）以提取深层特征 [40]。其中 RSTB 结构由多个 Swin Transformer 层（STL）组成，这些层相互协作以捕获图像的更深层特征。接下来，我们融合浅层和深层特征，以整合高频和低频信息。最后，通过上采样将深层特征缩减到原始图像空间，以获得协调的图像\(I_{smooth}'\)。为了优化网络，我们计算原始图像\(I_{input}\)和\(I_{smooth}'\)之间的差距\(L_{har}\)。损失\(L_{har}\)表示如下：

\(\mathcal{L}_{har} = \left\| I_{smooth}'-I_{input}\right\| _{2}^{2} \quad (3)\)

如图 3 所示，与马赛克图像\(I_{mosaic}'\)相比，\(I_{smooth}'\)有效解决了分布不一致问题并提高了图像质量。

在数据生成阶段，我们采用了改进的马赛克过程。对图像进行进一步增强，以增加增强样本的多样性。具体而言，首先，我们使用马赛克将图像\(I_{input}\)与数据集中的其他图像拼接形成\(I_{Mosaic}\)。接下来，我们对图像执行退化操作：

\(I_{mix} = \mathcal{P}\left(I_{Mosaic}, \mathcal{D}egrade\left(I_{Mosaic}\right), M_{select}\right)\)

混合图像\(I_{mix}\)可以提高马赛克的多样性，并且将混合图像\(I_{mix}\)输入到图像重建网络中，生成协调的马赛克图像\(I_{smooth}\)。

图 4：退化和剪切粘贴过程。将图像\(I_{input}\)输入到退化模块中以获得\(I_{degrade}'\)。从\(I_{input}\)中选择一个随机区域 M，将其剪切出来，并粘贴到退化图像\(I_{degrade}'\)的相应区域，以生成混合图像\(I_{mix}'\)。

C. 带有 Diffusion-Prior 机器的马赛克

与\(I_{Mosaic}\)相比，Pixel-Prior 机器生成的图像\(I_{smooth}\)表现出更和谐和更平滑的纹理。然而，\(I_{smooth}\)与真实图像之间仍然存在差距。这种差异是由于在优化增强网络时使用了 L2 损失，这会放大次要特征之间的差异。因此，Pixel-Prior 机器关注整体像素值的均匀性，而忽略了图像内的细节。

为了解决这个问题，我们引入了扩散模型。与传统的数据增强方法不同，扩散模型采用强大的生成先验，基于图像像素的分布来模拟现实世界场景。这使得扩散模型能够整合现实世界信息，并生成具有详细和真实近似的图像。在本文中，我们使用潜在扩散模型（LDM）[37] 作为生成先验。

LDM 是一种先进的扩散模型。它使用预训练的自动编码器将图像映射到潜在空间以学习数据分布。自动编码器由编码器 ε 和解码器 D 组成，其中 ε 将输入图像 I 编码为潜在代码 z，D 对其进行重建。在该模型中，向潜在变量 z 添加 T 次噪声，生成高斯分布的代码\(z_{T}\)。训练噪声预测网络\(\epsilon_{\theta}\)在时间步 t 进行去噪，并使用\(L_{ldm}\)进行优化。\(L_{ldm}\)计算如下：

\(\mathcal{L}_{ldm} = \mathbb{E}_{z, t \sim \mathcal{U}(0,1)}\left\| N_{gaussian}-\epsilon_{\theta}\left(z_{t}, t, c\right)\right\| _{2}^{2}\)

其中\(N_{gaussian}\)是高斯噪声，\(z_{t}\)是时间步 t 的潜在噪声。大规模训练的扩散模型 [37] 能够理解图像中的各种属性，如物体形状和纹理，因此能够生成视觉上吸引人的图像。为了进一步生成更高质量和更真实的图像，我们在 SIRST 数据集上对扩散模型进行了微调。

如图 3 所示，给定\(I_{smooth}'\)，在训练阶段，我们通过编码器 ε 将图像\(I_{smooth}'\)编码为潜在空间编码\(z_{0} = E(I_{smooth}')\)。然后，对\(z_{0}\)模型进行 T 步去噪，得到接近高斯分布的噪声\(z_{T}\)，然后通过 T 步重采样生成\(z_{0}'\)。最后，解码器 D 生成重采样图像\(I_{realis}'\)。为了使扩散模型重采样的结果更接近我们的数据集，我们通过\(L_{realis}\)减小\(z_{0}'\)和\(z_{0}\)之间的差距。\(L_{realis}\)表示如下：

\(\mathcal{L}_{realis} = \left\| z_{0}-z_{0}'\right\| _{2}^{2}\)

在数据生成阶段，通过微调的 LDM 重采样的图像\(I_{realis}\)具有更精细的细节，并且与原始数据集的分布非常接近。此外，由于生成先验的实现，增强样本\(I_{realis}\)更加多样和真实。值得注意的是，这种方法可以捕获潜在的数据特征，而不需要额外的标签作为输入。因此，它方便地为数据集重建生成更真实和多样的样本。

在数据生成阶段。通过采用扩散先验，我们对\(I_{smooth}\)进行重采样以生成\(I_{realis}\)。

如图 1 所示，与马赛克图像相比，\(I_{realis}\)变得更加连贯。此外，由于生成结果\(I_{realis}\)是从现实世界中提取信息并由扩散模型生成的，因此它与数据集中的其他图像有很大不同，这显著增强了数据集的多样性。在数据生成阶段生成的\(I_{realis}\)将作为增强样本参与训练。

D. 红外小目标检测

为了证明我们提出的数据增强方法的优势，我们利用标准检测网络进行推理。检测网络的结构如图 5 所示。首先，将图像\(I_{realis}\)输入到残差注意力块中以提取特征。残差注意力块由两个卷积层、通道注意力模块和空间注意力模块组成。这两个注意力模块旨在增强经过两个卷积层后的特征信息。为了融合不同的特征，将特征不断上采样和下采样到不同的残差注意力块中，形成特征金字塔。最后，融合这些包含不同尺度信息的特征，生成用于预测结果\(\hat{M}\)的鲁棒特征图。真实目标结果 M 和预测结果\(\hat{M}\)都是二值图像。为了使结果接近地面实况 M，网络通过损失\(L_{iou}\)减小\(\hat{M}\)与地面实况 M 之间的差距，其表示如下：

\(\mathcal{L}_{iou} = 1-\frac{\hat{M} \cdot M+\alpha}{\hat{M}+M-\hat{M} \cdot M+\alpha}\)

图 5：检测框架。我们将 Diff-Mosaic 生成的图像\(I_{realis}\)作为增强样本输入到检测网络中进行训练。骨干网络由多组注意力模块组成，每个嵌套的注意力模块由两个卷积层、一个通道注意力模块和一个空间注意力模块组成。输入到密集嵌套注意力模块的图像生成不同尺度的特征。最后，融合这些特征以产生预测结果\(\hat{M}\)。红色圆圈表示目标所在的区域，黄色圆圈表示错误预测的区域。模型通过\(L_{iou}\)缩小\(\hat{M}\)与地面实况 M 之间的差距。

四、实验

A. 实现细节

1. 数据集：为了证明我们方法的有效性，我们选择了两个数据集进行比较：

单帧红外小目标检测（SIRST）是一个开源的单帧红外小目标检测数据集，它从序列中选择图像，该数据集在 2020 年被亚利桑那大学指定为公共数据集。它包含来自数百个不同场景的真实世界视频的 427 个不同场景的代表性图像。这些图像是在短波、中波和 950nm 波长下拍摄的。它们以五种不同的形式进行标记，以支持用于检测任务和分割任务的模型。数据集的总量为 462，我们使用 70% 的 SIRST 数据集进行训练，30% 用于测试。

NUDT-SIRST：该数据集是一个手动合成的数据集，包含城市、田野、高光、海洋和云五个主要背景场景。每个图像都是从真实背景合成的，具有不同的目标（例如，点目标、圆点目标和扩展目标），具有不同的 SCR 和丰富的姿态。数据总量为 1362，我们将数据集分为 50% 用于训练，50% 用于测试。

1. 实现细节：为了实现从真实场景中提取知识的数据增强，我们训练了 Diff-Mosaic。对于 Pixel-Prior 机器中的退化部分，我们使用二阶退化。对输入图像重复两次模糊 - resize - 噪声过程 [39]。对于重建网络，我们将批处理大小设置为 4，训练 epoch 设置为 100，学习率设置为 0.001。对于 Diff-Prior 机器，我们微调扩散模型 100 个 epoch，并使用 Adam 优化器，将学习率设置为\(10^{-4}\)。在本实验中，我们通过 Diff-Mosaic 在 NUDT-SIRST 数据集上生成了 400 个增强样本，在 SIRST 数据集上生成了 100 个增强样本。检测网络的训练批处理大小为 40， epoch 数为 3000。

2. 比较方法：为了证明我们方法的有效性，我们将我们的方法与最先进的方法进行比较，包括传统方法（New Tophat [11]、WSLCM [14]、TLLCM [15]、IPI [16]）和基于 CNN 的方法（ACM [17]、ALC-Net [19]、UIU-Net [21] 和 DNA-Net [20]）。为了公平比较，所有基于 CNN 的方法都训练了 3000 个 epoch，其他配置与原始论文一致。ACM 和 ALC-Net 使用 Mxnet 实现，DNANet 和 UIUNet 使用 PyTorch 1.10 实现，并在 NVIDIA Geforce RTX 3090 上运行。

B. 评估指标

为了评估基于 CNN 的方法的检测性能，我们使用三种不同的评估指标来评估网络性能。

• 交并比（IoU）：IoU 是一种经典的像素级语义分割评估指标，用于描述算法的轮廓描述能力。它定义为预测值和标签之间的交集和并集面积的比率，如下所示：

\(IoU = \frac{A_{inter}}{A_{Union}} \quad (8)\)

其中\(A_{inter}\)、\(A_{Union}\)分别表示交互区域和联合区域。

• 检测概率（\(P_d\)）：检测概率是一种目标级评估指标。它衡量正确预测的目标数量\(T_{correct}\)与所有目标数量\(T_{All}\)的比率。\(P_d\)的定义如下：

\(P_d = \frac{T_{correct}}{T_{All}} \quad (9)\)

• 误报率（\(F_a\)）：误报率是另一种目标级评估指标。它用于衡量错误预测的像素\(P_{false}\)与所有图像像素\(P_{All}\)的比率。\(F_a\)的定义如下：

\(F_a = \frac{P_{false}}{P_{All}} \quad (10)\)

C. 实验结果

我们在 SIRST 和 NUDT-SIRST 上比较了最先进的方法与我们的方法。并使用 IoU、\(P_d\)和\(F_a\)三个指标进行评估。如表 I 所示，我们的方法在 SIRST 和 NUDT-SIRST 上的 IoU、\(P_d\)、\(F_a\)指标上取得了最佳结果。为了突出我们方法的优越性能，我们将我们的方法与其他检测模型进行可视化。为了突出我们方法的优越性能，我们使用红色虚线框圈出目标区域并放大显示。为了更清楚地比较模型检测结果与地面实况之间的差异，我们使用红色像素表示该方法与地面实况之间的差异。此外，我们用黄色虚线框标记了模型错误检测的区域。

如图 6 所示，我们可视化了最先进的方法与我们的方法在一般样本和困难样本上的性能比较。其中困难样本是数据集中不常见的实例。它们通常具有较大的尺寸、多个对象或不规则的轮廓。一般样本通常较小、单一且具有规则的形状。ACM 和 ALC-Net 有许多误报。UIU-Net 和 DNA-Net [20] 显示出较少的错误预测区域，然而，它们在一些困难样本（最后两个样本）中未能准确检测目标。相比之下，使用 Diff-Mosaic 生成的样本训练的检测模型在处理困难样本的预测中表现良好。这表明 Diff-Mosaic 可以生成更具挑战性的样本，用于检测模型中的判别训练。

表 I：不同最先进方法与我们的方法在 SIRST 和 NUDT-SIRST 数据集上的 IoU、\(P_d\)和\(F_a\)值比较，其中 IoU 和\(P_d\)的值越大表示性能越高，\(F_a\)的值越小表示性能越高。我们用红色和蓝色标记最佳和次佳结果，可以观察到我们的方法取得了最佳结果。

方法描述	SIRST(Tr=70%)			NUDT-SIRST(Tr=50%)
	IoU (×10⁻²) ￪	P_d (×10⁻²) ￪	F_a (×10⁻⁶) ￬	IoU (×10⁻²) ￪	P_d (×10⁻²) ￪	F_a (×10⁻⁶) ￬
基于滤波	New Tophat [11]	34.05	66.91	743	14.51	54.4
基于局部对比度	WSLCM [14]	20.64	60.86	1342	0.85	74.6
	TLLCM [15]	9.56	56.12	3046	7.06	62.01
基于低秩	IPI [16]	58.2	67.1	279	17.76	74.49
基于 CNN	ALC-Net [19]	70	71.71	23.91	72.13	78.81
	ACM [17]	76.17	86.31	16.07	71.11	85.13
	UIU-Net [21]	72.03	98.1	26.15	89	98.73
	DNA-Net [20]	76.97	95.41	3.54	88.38	97.99
我们的方法	79.44	99.99	3.19	91.18	99.47	1.91

如表 II 所示，我们列出了基于 CNN 的方法的计算参数数量和推理时间。我们的方法使用 pixel-prior 和 diffusion-prior 机制生成增强样本。这个过程发生在数据生成期间，每 10 秒生成一个样本，不影响训练和推理时间。因此，我们的方法在参数数量和推理时间方面与其他方法相比具有竞争力。

表 II：我们展示了不同基于 CNN 的方法所需的参数数量和每个样本所需的推理时间。

模型	参数 (M)	推理时间 (s)
ALC-Net	0.38	40.93
ACM-Net	0.29	18.53
DNA-Net	4.7	43.42
UIU-Net	50.54	33.98
我们的方法	4.7	43.42

D. Diff-Mosaic 增强规模的影响

我们分析了 Diff-Mosaic 生成的不同数量的样本作为增强样本对检测模型性能的影响。如表 III 所示，我们比较了仅使用真实数据训练的模型与使用额外 125、250 和 400 个样本增强的模型的性能。随着增强样本数量的增加，通过\(F_a\)、\(P_d\)和 IoU 指标衡量的模型性能有所提高。这表明 Diff-Mosaic 生成的增强样本的多样性和真实性可以提高检测模型的鲁棒性。

表 III：我们在 NUDT-SIRST 数据集上使用基线检测网络进行了实验，以比较使用不同数量的合成数据训练对准确性和检测性能的影响。我们比较了 125、250 和 400 个增强数据与真实数据结合训练的性能指标。可以看出，使用组合训练数据训练的模型的三个指标都优于使用真实训练数据训练的模型。

真实数据	663	663	663	663
合成数据	0	125	250	400
总数	663	788	913	1063
IoU(1×10⁻²)	88.38	90.11	90.21	91.18
P_d(1×10⁻²)	97.99	98.28	98.52	99.47
F_a(1×10⁻⁶)	4.04	3.24	3.1	1.91

图 7：可视化结果。我们比较了 UIU-Net 上 “不使用 Diff-Mosaic” 和 “使用 Diff-Mosaic” 的可视化结果，可以看出 UIU-Net 在使用增强和扩展的样本训练后，误报更少，对小样本目标的检测更准确。

E. 消融研究

本节的消融研究是在 NUDT-SIRST 上使用基线检测网络进行的。我们比较了使用马赛克训练的模型、使用 Pixel-Prior（“使用 Pixel-Prior”）训练的模型以及使用 Diff-Mosaic（“使用 Pixel-Prior+Diff-Prior”）训练的模型的性能。比较结果如表 IV 所示。由于 “使用马赛克” 生成的样本缺乏真实性和多样性，检测网络性能的提升非常有限。而 “使用 Pixel-Prior” 协调了增强样本，使其更真实，因此在所有三个性能指标上都有所提高。最后，“使用 Pixel-Prior+Diff-Prior” 引入了现实世界信息，并生成了多样且真实的增强样本。在 “使用 Pixel-Prior+Diff-Prior” 的帮助下，检测网络在所有三个性能指标上都显示出显著提升。

表 IV：消融研究。我们比较了我们方法的每个部分与传统马赛克之间的性能差距。可以看出，马赛克的增强样本对性能的提升不大。使用 Pixel-Prior 生成的样本训练的基本检测网络在三个评估指标上都有提升。使用 “使用 Pixel-Prior+Diff-Mosaic” 生成的样本训练的基本检测网络取得了显著的提升。

	NUDT-SIRST(Tr=50%)
	IoU(1×10⁻²) ￪	P_d(1×10⁻²) ￪	F_a(1×10⁻⁶) ￬
基线	88.38	97.99	4.04
使用马赛克	89.3	98.36	2.91
使用 Pixel-Prior	91.01	98.91	2.13
使用 Pixel-Prior+Diff-Prior	91.18(+0.17)	99.47(+0.56)	1.91(-0.22)

表 V：UIU-Net 使用和不使用 diff-mosaic 的定量结果。可以看出，“使用 Diff-Mosaic” 的结果与 “不使用 Diff-Mosaic” 相比有显著提升。

	IoU(1×10⁻²)	F_a(1×10⁻⁶)
不使用 Diff-Mosaic	89	6.02
使用 Diff-Mosaic	91.91(+2.91)	1.40(-4.62)

为了证明 Diff-Mosaic 生成的增强样本的有效性，我们将 Diff-Mosaic 的生成结果作为增强样本用于不同检测模型的训练。在本节中，我们使用 UIU-Net 和 DNA-Net 作为检测模型，比较了不使用（“不使用 Diff-Mosaic”）和使用 Diff-Mosaic 的（“使用 Diff-Mosaic”）增强样本训练的效果。如表 V 所示，我们比较了 UIU-Net 模型使用和不使用增强样本的性能。为了突出差距，我们使用红色粗体字体表示 “使用 Diff-Mosaic” 在 IoU 上的提升，使用蓝色粗体字体表示 “使用 Diff-Mosaic” 在\(F_a\)上的降低。可以看出，UIUNet 在使用 Diff-Mosaic 增强样本训练后，在 NUDT-SIRST 数据集上的性能显著提高。如图 7 所示，可以看出 “不使用 Diff-Mosaic” 的检测结果中有很多误导性目标，并且对小目标的轮廓检测不准确。此外，“使用 Diff-Mosaic” 的检测结果没有误检测目标，实现了对小的不规则目标轮廓的准确检测。

表 VI：DNA-Net 使用和不使用 diff-mosaic 的定量结果。可以看出，“使用 Diff-Mosaic” 的结果与 “不使用 Diff-Mosaic” 相比有显著提升。

	IoU(1×10⁻²)	F_a(1×10⁻⁶)
不使用 Diff-Mosaic	88.38	4.04
使用 Diff-Mosaic	91.18(+2.8)	1.91(-2.13)

如表 VI 所示，我们比较了 DNA-Net 模型使用和不使用 Diff-Mosaic 生成的增强样本的性能。结果表明，使用增强样本时，在 NUDT-SIRST 数据集上的检测性能有显著提升。如图 8 所示，很明显 “不使用 Diff-Mosaic” 方法在具有挑战性的测试样本中未能检测到目标，而 “使用 Diff-Mosaic” 方法以高精度准确识别了目标。为了展示我们生成的图像的真实性，我们将 pixel-prior 机器的结果输入到最先进的图像生成技术中，包括 SwinIR [41]、FemaSR [42] 和 DiffBIR [39]，并将它们的生成结果与我们的方法进行比较。如图 9 所示，可以观察到 SwinIR、FeMaSR 和 DiffBIR 生成的增强样本未能很好地整合马赛克策略。相反，我们的 diff-mosaic 生成的增强样本在图像分布上表现出无缝的连贯性和高度的真实性。此外，为了衡量生成图像的真实性，我们使用 KID [43] 和 FID [44] 指标来衡量它们与真实红外图像之间的差异。如表格 VII 所示，我们的方法在 FID 和 KID 方面优于其他方法，表明在增强真实性方面具有显著优势。

表 VII：不同图像生成方法的性能比较。我们使用 FID 和 KID 来衡量生成图像的真实性。可以看出，我们的方法取得了最佳结果。

	FID ￬	KID ￬
SwinIR [41]	181.39	0.167
FeMaSR [42]	200.47	0.191
DiffBir [39]	188.22	0.163
diffmosaic	126.01	0.076

图 8：可视化结果。我们比较了 DNA-Net 上 “不使用 Diff-Mosaic” 和 “使用 Diff-Mosaic” 的可视化结果。

F. Diff-Mosaic 的效果

表 VIII：将扩散先验应用于 CutMix [22]。可以看出，“使用 CutMix + 我们的扩散先验” 生成的增强样本比 “使用 CutMix” 更能激发检测模型的潜力。

方法	IoU(1×10⁻²) ￪	P_d(1×10⁻²) ￪	F_a(1×10⁻⁶) ￬
基线模型	88.38	97.99	4.04
使用 CutMix	89.24	98.62	6.25
使用 CutMix + 我们的扩散先验	90.11	99.15	2.04

图 9：不同图像生成方法的可视化结果。我们可视化了不同生成方法的生成结果。可以看出，与其他方法相比，我们的方法很好地整合了图像，使其更连贯。

G. Diffusion-Prior 的效果

为了证明 diffusion-prior 机器的有效性，我们将其与 CutMix [22] 相结合。随后，我们在 NUDT-SIRST 上生成了 400 个增强样本，并将其输入到检测网络中进行训练。如图 10 所示，我们展示了原始 CutMix（“使用 CutMix”）和与我们的方法结合（“使用 CutMix + 扩散先验”）的增强样本。很明显，原始 CutMix 生成的结果缺乏连贯性和真实性。相比之下，“使用 CutMix + 扩散先验” 生成的结果表现出更大的连贯性。为了验证我们的扩散先验的有效性，我们比较了 “使用 CutMix” 和 “使用 CutMix + 我们的扩散先验” 生成的增强样本的性能。如表 VIII 所示，可以观察到 “使用 CutMix” 对网络的提升微乎其微。相比之下，“使用 CutMix + 扩散先验” 引入了现实世界信息，导致网络的检测能力有明显提升。

图 10：我们可视化了 CutMix [22] 方法生成的增强样本与结合我们方法生成的增强样本。

五、结论

在本文中，提出了一种新型的基于扩散的数据增强方法，以解决现有数据增强方法缺乏真实性和多样性的问题。我们的方法基于图像像素分布生成协调的增强样本，使样本更真实。然后，我们使用强大的扩散模型先验，通过对样本进行重采样来引入现实世界信息。这种重采样使生成的样本真实且多样。大量实验表明，我们的方法有效地提高了模型的性能。此外，消融实验表明，我们方法的子组件是有效的，并且我们方法生成的增强样本可以提高不同基线检测网络的性能。