使用Transformer模型进行计算机视觉任务的端对端对象检测
Transformer模型是google团队在2017在论文attention is all you need中提出的一个用于NLP领域的模型,但是随着VIT模型与Swin Transformer模型的发布,把Transformer模型成功应用到计算机视觉任务中。
上期图文,我们使用hugging face的transformers模型进行了VIT模型的对象分类任务。本期我们再次分享一个成功把Transformer模型应用到计算机对象检测任务模型。此模型是Facebook发布的基于Transformer模型的端对端对象检测任务模型-- DETR(detect Transformer模型)。

DETR模型首先使用CNN卷积神经网络搜集图片的核心特征点,然后把这些特征点整合起来,通过embedding方法,把特征图片转换到特征向量空间。然后根据标准Transformer模型的编码器与解码器进行注意力机制的计算,最后把计算后的数据进行图片对象的分类,并根据检测到的位置信息,提供对象box区域,方便我们画图。

传统的对象检测算法,如 Faster R-CNN,通过过滤大量粗略对象区域来预测对象边界框,这些区域通常是 CNN 特征图片区域。每个选定的区域都用于细化操作,包括在区域定义的位置裁剪 CNN 特征,独立对每个区域进行分类,并细化其位置。最后,应用非最大抑制步骤来删除重复框。而DETR 通过利用标准的Transformer架构来执行传统上特定于对象检测的操作,从而简化了检测管道,优化了算法步骤,提高了检测效率。

Transformers 的自注意力机制允许DETR模型对图像以及预测的特定对象执行全局推理。

例如,模型可能会查看图像的其他区域,以帮助对边界框中的对象做出检测。它还可以根据图像中对象之间的关系或相关性进行预测。例如,如果DETR预测图像中有一个人站在沙滩上,它就会知道部分被遮挡的物体更有可能是冲浪板。相比之下,其他检测模型只能独立的预测每个对象,每个对象之间并没有相关性。
代码实战DETR模型
既然是Facebook发布的模型,那么其首当其冲的便是使用pytorch框架来实现,且pytorch模型的TorchHub已经集成了此模型,我们可以直接使用TorchHub模型来实现本期代码,当然若想了解DETR模型的框架,当然最好是源代码来实现,这个我们后期进行分享。

首先我们需要导入需要的第三方库,这里主要是torch框架以及其他库。
DETR模型训练在COCO数据集上,这里我们列举一下COCO数据集的对象分类,并定义一个colors颜色列表,方便我们进行不同颜色的备注。
然后我们定义几个函数,方便显示我们经过模型检测后的box,主要是为了我们后期进行box的画图,模型对象检测后的标签与模型检测置信度。
有了以上的函数,我们就可以把图片放入我们的DETR模型进行对象检测了。
首先我们使用torch.hub.load函数来加载DETR模型的预训练模型,并传递一张图片。当然这里可以直接传递图片的URL地址,也可以直接从本地打开一张需要检测的模型。
加载图片后,我们需要把图片转换到torch变量,然后把图片数据传递到model模型来进行对象检测的预测,结果保持在output中,然后我们使用soft Max函数来挑选概率最大的对象检测数据,并挑选出置信度大于0.8(预设值)的对象标签与置信度。
得到检测结果后,最后我们使用plot_results函数来可视化检测到的对象。

DETR模型对象检测
上期图文教程,YOLO-NAS对象检测算法再一次颠覆YOLO系列算法——已超越YOLOv8,我们使用YOLO的最新模型进行了同一张图片的对象检测,我们可以使用以上代码与YOLO- NAS模型来进行对象检测的对比。

YOLO_NAS模型对象检测
从对象检测的速度与精度来讲,基于transformer注意力机制算法的DETR模型同样可以使用在对象检测任务中,让transformer模型进行模型大一统更近了一步。