专栏名称: 江大白

专业有趣的AI内容平台，关注后回复【算法】，获取45家大厂的《人工智能算法岗江湖武林秘籍》

基于SAM模型，红外小目标检测算法IRSAM（附论文及源码）

江大白 · 公众号 · · 2024-08-14 08:00

正文

以下文章来源于微信公众号：小小cv笔记

作者：bochen

链接：https://mp.weixin.qq.com/s/hpsD7Qc2xLV-Ns5JsgNy2A

本文仅用于学术分享，如有侵权，请联系后台作删文处理

导读

红外小目标检测（IRSTD）在计算机视觉中是一个重要的研究领域，主要用于检测红外图像中的小目标。传统的IRSTD方法和一些基于深度学习的方法在处理复杂背景和噪声时效果不佳，且通常需要大量标注数据进行训练。因此本文介绍了一种将SAM应用于IRSAM的方法，该方法在IRSTD-1k、NUDT-SIRST和NUAA-SIRST数据集上取得优异成绩！

论文链接：https://arxiv.org/abs/2407.07520

代码链接：https://github.com/IPIC-Lab/IRSAM

1. 简介

最近的分割任意模型(SAM)在自然图像分割方面取得了重大进展，表现出适用于各种下游图像分割任务的强大零射击性能。然而，由于自然图像和红外图像之间存在明显的域差距，直接利用预训练的SAM进行红外小目标检测(IRSTD)任务无法达到令人满意的性能。与可见光相机不同，热成像仪通过捕捉红外辐射来揭示物体的温度分布。小目标通常在物体的边界处显示出微妙的温度变化。为了解决这一问题，我们提出了IRSAM模型，该模型改进了SAM的编码器-解码器架构，以更好地学习红外小目标的特征表示。具体来说，我们设计了一个基于Perona-Malik扩散(PMD)的块，并将其集成到SAM编码器的多个级别中，以帮助它在抑制噪声的同时捕获基本结构特征。此外，我们设计了一个粒度感知解码器(GAD)来融合编码器的多粒度特征，以捕获可能在远程建模中丢失的结构信息。在包括NUAA-SIRST、NUDT-SIRST和IRSTD-1K在内的公共数据集上进行的大量实验验证了IRSAM的设计选择及其相对于代表性最先进方法的显著优势。

本文的贡献主要有以下几个方面:

1）首次针对IRSTD任务重新设计了通用视觉分割模型SAM，从而引入了IRSAM。在具有挑战性的基准测试中，IRSAM在客观指标和主观评估方面都优于普通SAM模型和最先进的(SOTA)方法，表现出卓越的性能。

2）设计了一个WPMD模块，以增强SAM编码器在抑制红外图像噪声的同时保持边缘相关特征的能力，有效解决IRSTD任务中的低信噪比问题。

3）设计了一个GAD，通过精心设计的边缘标记融合编码器的多粒度特征来重建在捕获远距离依赖中丢失的目标结构特征，有效增强了各种大小和形状物体的掩模表示。

2. 背景

当前红外小目标检测存在如下问题：

1）当前方法通常需要大规模的标记数据进行训练，与自然图像数据集相比，获得红外数据本质上更具挑战性。

2）SAM方法依赖于提示来指导分割过程，限制了它们在IRSTD任务中的适用性和效率，此外自然图像与红外图像之间存在固有的域间隙，这也阻碍了SAM在IRSTD中的性能。

3）直接将PMD(Perona-Malik diffusion)集成到神经网络中存在挑战，包括边界模糊和对噪声干扰的鲁棒性降低。

本文针对以上问题的改进思路：

1）鉴于对自然图像领域深度模型的广泛研究以及迁移学习在缓解下游任务中有限训练数据的泛化问题方面的有效性，利用在大规模自然图像数据集上预训练的基础分割模型来完成IRSTD任务。

2）引入了IRSAM模型，增强了SAM的编码器-解码器架构，以捕获更有效的红外小目标特征表示，同时利用其预训练的知识进行分割。

3)开发了一个WPMD模块，通过将小波变换中的高频分量作为PMD(Perona-Malik diffusion)方程中的图像微分，旨在保留结构信息同时减轻噪声影响

3. 方法

1）下图为本文提出的IRSAM的总体架构，它采用了编码器-解码器结构。IRSAM的编码器由预训练的ViT-Tiny骨干和WPMD块（Wavelet transform-based PMD Block）组成。ViT-Tiny将WPMD提取的每一层边缘特征进行聚合。对于解码器，使用Granularity-Aware解码器，与原来的SAM解码器架构不同，融合了来自编码器的不同粒度的特征，并使用输出token令牌与图像特征交互并生成最终的高质量目标掩码。

2）Perona-Malik扩散方程主要用于图像处理。各向异性扩散的特性使其能够促进平滑区域的扩散(平滑)，同时抑制边缘的扩散，从而达到提高图像质量、增强图像结构、抑制噪声的效果。红外图像经常受到噪声的干扰，目标边界模糊，探索PMD在传递过程中同时保留重要的结构信息和消除噪声。

3）WPMD块（Wavelet transform-based PMD Block），使用小波变换的高频分量作为PMD方程中的图像微分，而不是使用卷积算子。因此，WPMD的输出将是输入的平滑版本，在消除噪声的同时保留了基本的结构信息。

给定图像u，其对应的PMD方程为:

其中扩散系数g(|∇u|) = 1/(1 + |∇u|2 /k2)。K是一个正常数，用于控制扩散程度，T表示步长。由式(1)可知，当光滑区域的梯度幅度|∇u|较小时，扩散系数g(|∇u|)较大。因此，扩散强，有效地消除了噪声。在边缘部分，梯度幅度|∇u|较大，扩散系数g(|∇u|)较小。因此，扩散是弱的，边缘信息被保留。

式(1)可表示为:

另一方面，图像的二维小波变换可以表示为:

式中F(·)表示过滤操作。L和H分别代表low-pass和high-pass。用小波帧FLH(·)和FHL(·)逼近微分方程∂/∂x或∂/∂y，并将扩散步长△t设置为1，我们可以将Eq.(2)转换为离散形式:

如图3所示，经过扩散过程后，使用卷积层将得到的结构特征映射到与每层编码器特征相同的维度上。通过在不同层将多个WPMD模块集成到SAM编码器中，增强了SAM编码器对噪声的抑制能力，同时保持了红外图像的结构特征，有效地解决了IRSTD任务中低信噪比和目标边缘模糊的问题。

4）Granularity-Aware解码器，红外小目标通常具有有限的视觉特征，容易与背景目标或相似目标混淆。为了提高红外小目标分割的性能，既需要考虑全局上下文信息，这有助于提取图像的整体语义，增强对小目标的检测，也需要考虑局部边界信息，这有助于保留小目标的空间细节，提高分割边界的精度。SAM采用ViT架构，擅长捕获长期依赖关系和全局信息。在之前的研究中，ViT结构的早期层被证明可以保留更一般的图像边界细节，而深层包含更高级的语义。为了提高SAM在IRSTD任务中的性能，设计了粒度感知解码器来融合多粒度特征，通过向解码器提供全局语义上下文和局部细粒度特征，GAD拥有更丰富的多视图知识，如图2所示。

5) 损失函数

1、Dice loss:Dice loss是一种常用的评估预测掩模与真实值之间差异的方法，其定义如下:

2、BCE损失:二进制交叉熵(BCE)损失是二值分类任务中常用的损失函数。用于测量边缘预测YPreddge与地面真实边缘YEdge的差值，定义为:

3）总的损失函数，由Dice loss lDice和BCE loss LBCE组成的final loss用于监督IRSAM的训练:

其中λ是一个超参数，用于平衡两个损失，并根据经验设置为10。

4. 实验

1)与代表性方法IRSTD-1k、NUDT-SIRST和NUAA-SIRST数据集上的对比

2）和其他方法的可视化比较

3）与用SAM微调方法的比较，轻松打败最大模型的SAM

4）消融实验，WPMD 和 GAD两个模块的作用

5. 推荐阅读

江大白 | 这些年从0转行AI行业的一些感悟

AI项目应用 | 万字长文教你一步步实现离岗识别

AI行业大家有想问的，欢迎来《AI未来星球》，一起成长

扫描下方二维码即可加入~

你可以获得什么？

1、AI行业创业/工作/生活中遇到的各类问题，都可获高质量解答， 有效期一年，无限次提问，有问必答。

2、购买过的，数万元的各类数据集的分享；

3、定期邀请AI行业各类嘉宾分享，创业/商业等方面的经验！

4、AI行业研发、产品、商业落地方面专业分享（大白主业，也在负责公司AI业务板块）

帮助你解决遇到的实际问题，升职加薪！

大家一起加油！