现有的深度边缘检测网络通常基于包含了上下采样模块的编码器 - 解码器架构,以更好的提取多层次的特征,但这也限制了网络输出既准又细的边缘检测结果。
针对这一问题,一篇 AAAI 2024 上的论文给出了新的解决方案。
- 论文题目:DiffusionEdge: Diffusion Probabilistic Model for Crisp Edge Detection
- 作者:叶云帆(国防科技大学),徐凯(国防科技大学),黄雨行(国防科技大学),易任娇(国防科技大学),蔡志平(国防科技大学)
- 论文链接: https://arxiv.org/abs/2401.02032
- 开源代码: https://github.com/GuHuangAI/DiffusionEdge
国防科技大学 iGRAPE Lab 提出了首个用于二维边缘检测任务的扩散概率模型方法,通过学习迭代的去噪过程获得边缘结果图,去噪过程如图 1 所示。为了在保留最终性能的同时减少计算资源的消耗,该方法在隐空间中训练网络,并引入了不确定性蒸馏模块以更好的优化。同时,本方法还采用解耦架构来加速去噪过程,并提出了相应的自适应傅立叶滤波器来调整特征。基于这些设计,本方法得以用有限的资源进行稳定的训练,用更少的增强策略来预测清晰准确的边缘图。在四个公共基准数据集上的大量实验表明,本文提出的方法在准确度和粗细度上全面超越了其他方法。
图 1 基于扩散概率模型的边缘检测过程与优势示例
本文的创新点包括:
1、提出了领域内第一个针对边缘检测任务的扩散模型 DiffusionEdge,无需任何后处理即可预测出更细更准确的边缘图。
2、针对扩散模型应用时的难点,设计了多种技术以确保方法能在隐空间中稳定学习,同时保留像素级的不确定性先验知识并自适应地过滤傅立叶空间中的隐特征。
3、在四个边缘检测公共基准数据集上开展的大量对比实验展示了 DiffusionEdge 在准确度和细度方面均具有卓越的性能优势。
相关工作
基于深度学习的方法通常采用包含上下采样的编解码结构集成多层特征 [1-2],或是整合多个标注的不确定性信息以提升边缘检测的准确度 [3]。然而,天然受限于这样的结构,其生成的边缘结果图对于下游任务来说太过粗厚而严重依赖后处理的问题仍然亟待解决。尽管许多工作已经在损失函数 [4-5] 和标签修正策略 [6] 方面做出了探索以使网络能输出更细的边缘,但本文认为该领域仍然需要一个可以不借助任何额外模块,就能直接满足准确度和细度的边缘检测器,而无需任何后处理步骤。
扩散模型是一类基于马尔可夫链的生成模型,通过学习去噪过程逐渐恢复目标数据样本。扩散模型在计算机视觉、自然语言处理和音频生成等领域都表现出了卓越的性能。不仅如此,通过将图像或是其他模态的输入作为额外条件时,其在感知任务中也表现出了巨大的潜力,例如图像分割 [7]、目标检测 [8] 和姿态估计 [9] 等。
方法描述
本文所提出的 DiffusionEdge 方法总体框架如图 2 所示。受以往工作的启发,该方法在隐空间中训练具有解耦结构的扩散模型,并将图像作为额外的条件线索输入。该方法引入了自适应傅里叶滤波器来进行频率解析,且为了保留来自多个标注者的像素级不确定性信息并减少对计算资源的要求,还以蒸馏的方式直接使用交叉熵损失优化隐空间。
图 2 DiffusionEdge 的整体结构示意
针对目前的扩散模型受到采样步数太多,推理时间太长等问题的困扰,本方法受 DDM [10] 的启发,同样使用解耦的扩散模型架构来加速采样推理过程。其中,解耦的前向扩散过程由显式的转移概率和标准 Wiener 过程的组合来控制:
如图 2 所示,该方法首先训练了一对自编码器和解码器的网络,该编码器将边缘标注压缩为一个隐变量,而解码器则用于从这个隐变量中恢复出原来的边缘标注。如此一来,在训练基于 U-Net 结构的去噪网络阶段,该方法便固定这一对自编码和解码器网络的权重,并在隐空间中训练去噪过程,这样可以大幅减少网络对计算资源的消耗,同时维持不错的性能。
实验结果
本方法在四个领域内被广泛使用的边缘检测公共标准数据集上进行了实验:BSDS、NYUDv2、Multicue 和 BIPED。由于边缘检测数据标注比较难,标注数据量都比较少,以往的方法通常会使用各种策略来增强数据集。例如说,BSDS 中的图像通过水平翻转 (2×)、缩放 (3×) 和旋转 (16×) 进行增强,能够生成比原始版本扩大了 96 倍的训练集。以往的方法在其他数据集上使用的通用增强策略总结在了表格 1 中,其中 F 代表水平翻转,S 代表缩放,R 代表旋转,C 代表裁剪,G 代表伽马校正。不同的是,本方法仅需要使用随机裁剪的 320320 的图像块来训练所有数据。在 BSDS 数据集中,本方法则仅仅采用随机的翻转和缩放,其定量对比结果展示在了表 2 中。在 NYUDv2、Mu lticue 和 BIPED 数据集中,该方法仅需采用随机翻转训练。在使用更少增强策略的情况下,本方法在各个数据集,各个指标上的表现都优于之前的方法。通过观察图 3-5 的预测结果可以看出,DiffusionEdge 能够学习并预测出跟 gt 分布几乎一样的边缘检测结果图来,预测结果准确而清晰的优势对于有精细化需求的下游任务来说非常重要,也展示了其能直接应用于后续任务的巨大潜力。
表 1 以往方法在四个边缘检测数据集上使用的增强策略
表 2 不同方法在 BSDS 数据集上的定量对比
图 3 不同方法在 BSDS 数据集上的定性对比
图 4 不同方法在 NYUDv2 数据集上的定性对比
图 5 不同方法在 BIPED 数据集上的定性对比