基于扩散模型的显著性物体检测

货拉拉技术
0 阅读7分钟

引言

背景

显著性目标检测（Salient Object Detection, SOD）是图像处理领域的一个重要组成部分，它通过模拟人类视觉系统的注意机制，快速从图像中提取最具吸引力或最重要的目标区域。

​

在图像编辑中，能够识别图像中的主体部分，抠出显著性主体，用于制作头像、证件照、合成创意图片等场景；

在医学图像分析里，显著性目标检测作用重大。面对 X 光、CT、MRI 等复杂影像，它能精准扫描，助力医生快速锁定如肺部结节、脑部肿瘤这类病变的显著区域；

当显著性目标检测应用于体育比赛视频时，能通过捕捉每一帧中的显著区域（如运动员的动作、球的位置等），辅助教练团队指导球队作战；比赛进行期间，实时检测如同教练的“第三只眼”。进攻时，教练依据检测结果及时向队员发出精准进攻指令；防守时，凭借对对方球员动作与球位置的检测，洞察其进攻意图，迅速调整防守阵型，有效遏制对手攻势。

挑战

  • 传统手工设计特征的方法：早期的方法依赖于颜色不一致和边缘变化等低级特征来检测显著物体。然而当显著目标与复杂背景混合时，检测模型可能会遇到在显著区域和背景之间产生清晰区分的挑战；

  • 基于深度学习的方法：利用深度神经网络来学习密集的特征交互和注意力机制，从而提高性能。尽管深度学习在显著性目标检测任务中取得了进展，但精确定位和分割仍然存在挑战，特别是在具有相似背景或复杂物体的复杂场景中，这往往导致过度自信但不准确的预测，跨领域泛化能力差；

  • 迫切需要一种新的范式和技术来解决 SOD 任务中的当前问题。

动机

现有基于编解码的SOD框架和我们提出的基于扩散模型的SOD框架

  • 扩散模型已成为生成模型中强大的范例，能够学习复杂的数据分布并生成精细的图像边缘；

  • 我们尝试利用扩散模型来解决 SOD 中的挑战，将其定义为去噪任务，利用扩散模型生成高质量掩码的能力；

  • 在前向过程（forward）中，对显著性掩码不断加噪得到随机噪声，在后向过程（reverse）中，对随机噪声迭代去噪还原出最初的显著性掩码。下图为一些对比效果图。

模型结构

扩散模型

扩散模型的由来

  • 数学理论起源于2015年发表在国际机器学习大会(ICML)的一篇论文Deep Unsupervised Learning using Nonequilibrium Thermodynamics；

  • 扩散模型流行于2020年发表在神经信息处理系统大会(NIPS)的一篇论文Denoising Diffusion Probabilistic Models(DDPM)。

扩散模型的组成

  • 前向扩散过程(Forward Diffusion Process): 不断往输入图片中添加高斯噪声；

  • 逆向扩散过程(Reverse Diffusion Process)：将噪声不断还原为原始图片；

  • 对应在我们SOD-diffusion任务，前向过程就是不断往GT加噪声，逆向过程就是将加噪的GT去噪还原出来。

创新点

  • 开启了SOD任务新范式，从判别式到生成式：不再是判别每个像素是不是显著性目标，而是将随机噪声去噪逐步生成显著性区域；

  • 显著性编码器SwinB和特征融合模块CBAM：显著性编码器学习到显著物体的部分位置信息和边缘细节，其中特征融合模块能够极大聚焦显著性物体；

  • 注意特征交互模块AFIM：引导整个扩散过程，两个输入分别为图像特征和扩散特征，其中图像特征就是上面输出的特征，扩散特征是含有GT信息的随机噪声，AFIM模块可以使得两个特征之间具有更强的耦合。

模型训练

我们使用预训练的Stable Diffusion VAE将图像x和显著性掩码y编码到潜在空间中，与此同时，显著性掩码y紧接着加入了随机噪声，另一方面构建注意力特征交互模块（AFIM） 来隐式地引导扩散过程，目的是更好地使经过SwinB和特征融合（FF）的图像特征和扩散信息之间具有更强耦合的条件语义特征，最后通过优化扩散模型的损失来使得模型收敛，详细见论文SOD-diffusion: Salient Object Detection via Diffusion-Based Image Generators。

模型推理

给定输入图像x，我们使用原始Stable Diffusion VAE对其进行编码，以获得输入图像隐编码zxz^x；随后，将此隐编码与显著性隐编码zyz^y连接，接着在每次去噪迭代期间输入到模型U-Net中。经过t步的去噪迭代后，显著性隐变量z0yz_0^y被解码，从而得到输出显著性掩码，最终的显著性掩码预测是通过对其三个通道的值取平均值得出的。

实验

实验设置

  • 与大多数现有研究一致，我们使用 DUTS-TR（10553 张图像）作为训练数据集。在 PyTorch 扩散模型的实现中，我们使用 Stable Diffusion v2 作为我们的backbone；

  • 我们在 5 个广泛使用的、在训练期间未见过的基准数据集上评估了 SOD-diffusion，包括 DUT-OMRON（5168 张图像）、ECSSD（1000 张图像）、PASCAL-S（850 张图像）、HKU-IS（4447 张图像）和 DUTS-TE（5019 张图像），我们采用通用的评估指标：F-measure、MAE、E-measure、S-measure。

评测指标

总体效果

第一行表示输入图像，第二行是GT，第三行是SOD-diffusion生成的结果

消融实验

  • SwinB有效性：SOD-diffusion 在早期去噪步骤中学习显著物体的部分位置信息和边缘细节，并在推理过程中利用辅助特征逐步细化最终的显著性掩码；

  • 特征融合CBAM模块的有效性：它证明了CBAM 在增强显著性检测方面的重要性，因为它使模型能够更好地识别和聚焦图像中最显著的区域；

  • 注意特征交互模块AFIM的有效性：我们提出的注意特征交互模块无缝地集成了来自主干的扩散特征和显著特征，性能优于原生的交叉注意力机制；

  • 辅助损失AL的有效性：我们的模型在辅助损失的帮助下取得了最佳性能。

采样过程可视化

总结

局限性

  • 虽然 SOD-diffusion 在处理物体边缘方面取得了出色的成绩，但平衡推理时间和检测准确性仍然是一个关键挑战，需要进行更大的优化；

  • 未来的工作应侧重于整合轻量级神经网络架构并探索知识蒸馏和模型修剪等训练策略，以增强实时适用性。

展望

  • 在本研究中，我们提出了基于扩散的显著物体检测框架 SOD-diffusion。为了克服现有 SOD 方法中存在的挑战，我们设计了三个关键模块来补偿扩散模型：特征融合与卷积块注意模块 (CBAM)、注意特征交互模块 (AFIM) 和辅助损失 (AL)；

  • 此外，我们提出的SOD-diffusion 通过采用生成模型进行显著物体检测，彻底改变了 SOD 任务当前的检测范式，并提高了性能；

  • 后续研究将探讨如何解决当前方法中的限制，包括推理效率的提高以及多种模态的SOD任务。作为显著物体检测的新范例，我们设想我们提出的方法将作为坚实的基础，激发未来的进一步探索。

作者简介

黄佳明，货拉拉技术中心智能平台部图像组高级算法工程师

avatar
文章
阅读
粉丝