代码: https://github.com/jhjie/edgenat
论文: https://arxiv.org/abs/2408.10527v1
论文 EdgeNAT: Transformer for Efficient Edge Detection 介绍了一种名为EdgeNAT的基于Transformer的边缘检测方法。
1. 背景与动机
EdgeNAT预测结果示例。(a, b):来自BSDS500的数据集的输入图像。(c, d):对应的真实标签。(e, f):由EdgeNAT检测到的边缘。(e)显示了由于颜色变化剧烈,EdgeNAT没有提取出鹿颈部区域的边缘。(f)显示了EdgeNAT准确提取出了远处模糊物体的边缘。
边缘检测是许多计算机视觉任务的基础,旨在从输入图像中精确提取物体边界和视觉显著的边缘。然而,由于图像中存在远距离物体、复杂背景中的模糊边界以及物体内部的颜色变化等挑战,边缘检测任务变得十分困难。传统的边缘检测方法主要依赖于颜色和纹理等局部信息,而基于卷积神经网络(CNN)的深度学习方法虽然可以扩展感受野以捕捉全局特征,但容易丢失细节信息。
近年来,Transformer在视觉任务中展示出了强大的特征提取能力,特别是具有层次结构的Dilated Neighborhood Attention Transformer (DiNAT)在有效捕捉全局和局部特征方面表现出色。然而,Transformer在边缘检测中的应用尚未得到充分开发。为了解决这一问题,论文提出了一种基于DiNAT的单阶段边缘检测模型——EdgeNAT,以期在精确提取物体边界和重要边缘的同时,提升检测效率。
2. EdgeNAT的模型设计
我们提出的EdgeNAT的整体框架。
(a) SCAFM的详细架构。(b) 空间注意力模块(SAM)。(c) 通道注意力模块(CAM)。
2.1 编码器:DiNAT
EdgeNAT的编码器采用了Dilated Neighborhood Attention Transformer (DiNAT),该模型能够同时保留局部信息的翻译等变性,并通过扩展感受野来捕捉长距离的特征依赖。DiNAT的结构通过在每个层次之间使用3x3卷积进行下采样,逐层减小空间分辨率,同时增加通道数。这种设计使得DiNAT在保持局部性和扩展感受野方面表现优异,从而适用于边缘检测任务。
2.2 解码器:SCAF-MLA
为了充分利用Transformer编码器生成的丰富特征,论文提出了一种新的解码器结构——空间和通道注意力融合多级聚合(SCAF-MLA)。这个解码器的核心组件是SCAF模块,它结合了空间注意力模块(SAM)和通道注意力模块(CAM),以同时计算特征图的空间和通道权重。这种设计能够在保持当前层次特征的独特性的同时,捕捉更高级别的特征,有助于在更高层次提取全局语义信息,同时在低层次保留细致的局部信息。
此外,解码器还通过一种称为预融合的技术来进一步提高性能,即在融合操作之前,将每一层的特征通道减少到与编码器第一层相同的数量,而不是直接减少到1。这种方法能够更好地集成不同层次的特征,提高边缘检测的精度。
3. 实验与结果
3.1 数据集与训练
论文在两个主流数据集上对EdgeNAT进行了评估:BSDS500和NYUDv2。BSDS500包含500张RGB图像,通过数据增强扩展到28,800张图像,模型在PASCAL VOC Context数据集上预训练后,使用BSDS500数据集进行微调并评估。NYUDv2数据集包含1449对RGB和深度图像,模型在经过扩展的训练和验证集上训练,并在测试集上进行评估。
对BSDS500测试集中三个具有挑战性的样本进行了定性比较。有趣的是,在第三个示例中,尽管帽子的边缘即使对于人眼也难以分辨,但我们的L模型仍然能够完整地检测出右侧帽子的边缘。这个前所未有的现象表明,我们的模型在全局语义理解方面优于之前的工作。
BSDS500测试集上的结果。最佳的两个结果分别用红色和蓝色标出,其他表格中也是如此。†表示使用额外的PASCAL VOC数据进行训练,‡表示多尺度测试。
3.2 消融实验
消融实验验证了不同设计对模型性能的影响。结果表明,预融合比最终融合对性能提升更大,而底部路径(Bottom-up Path)对DiNAT编码器的边缘检测效果反而有负面影响。此外,实验还验证了SCAF模块的有效性,相比于PPM模块,SCAF模块能够在不增加计算复杂度的情况下显著提升边缘检测的性能。
3.3 模型可扩展性
为了适应不同的应用场景,论文设计了五个不同参数规模的EdgeNAT模型(L、S0、S1、S2、S3),其中L模型拥有最多的参数。实验结果表明,随着模型规模的减小,边缘检测的性能略有下降,但推理速度显著提升,特别是S0模型,由于其第三层只有6层,处理速度远高于其他模型。
3.4 与现有方法的比较
在BSDS500数据集上,EdgeNAT-L模型在单尺度输入情况下取得了84.3%的ODS,超过了所有现有的边缘检测方法。在多尺度输入情况下,EdgeNAT-L模型的ODS达到了86.0%,进一步证明了其在准确性和效率方面的优越性。相比于传统的Canny和gPb-UCM方法,EdgeNAT在精度和召回率上均有显著提高。同时,EdgeNAT的推理速度在RTX 4090 GPU上达到20.87 FPS,远高于EDTER的2.2 FPS,展示了Transformer在边缘检测任务中的巨大潜力。
在NYUDv2数据集上,EdgeNAT-L模型在RGB、HHA和RGB-HHA三种输入类型下的ODS分别为78.9%、72.6%和79.4%,同样超过了现有的最佳方法,展示了模型的广泛适用性。
4. 结论
论文总结了EdgeNAT的主要贡献:通过引入DiNAT作为编码器,不仅提升了边缘检测的精度,还显著提高了计算效率;提出了SCAF模块,用于生成更丰富、更准确的特征表示;设计了适应不同应用场景的五个模型版本,并在BSDS500和NYUDv2数据集上进行了广泛的实验,证明了EdgeNAT在效率和准确性方面的优势。