transformer模型刚发布时,主要应用于处理NLP领域任务,比如机器翻译等,但是随着注意力机制模型的大火,很多基于transformer模型的魔改模型也相继发布,且transformer模型的注意力机制也被google团队证明可以使用在计算机视觉任务上,特别是swin transformer模型的发布,更是把transformer模型带入了计算机视觉领域。前几篇文章,我们也同样介绍了另外一个基于transformer模型的,应用在计算机视觉任务上的模型BETR。且BETR模型不仅可以使用在对象检测上,还可以使用在对象分割上,本期我们就基于transformer模型来代码实现一下BETR模型。
基于transformer模型的对象检测算法代码实现的第一步是需要我们import python的第三方库,这里主要是torch,确保你在运行本期代码前,已经成功安装了torch库。
这里我们需要建立一个list列表,一个是保存所有BETR模型能够检测的对象标签,另外一个是颜色数据,方便我们后期的可视化。
这里我们建立几个函数,方便我们进行对象的标注,后期进行可视化操作时比较方便。以上初始化的工作完成后,我们就可以搭建我们的BETR模型了。
根据BETR模型的框架图,我们知道这里有2层关键的框架,一个是CNN卷积神经网络层,另外一个就是transformer模型的编码器与解码器部分。
搭建好了BETR模型,我们就可以直接来进行使用了,在使用之前,我们需要下载BETR的预训练模型。
这里我们已经搭建好了DETR_model模型,然后我们需要使用torch.hub.load_state_dict_from_url函数来下载BETR的预训练模型,模型下载完成后,我们使用load_state_dict来加载模型,并进行eval。然后我们就可以使用BETR模型了。
这里我们建立一个detect函数,方便我们使用模型进行图片的对象检测,然后我们就可以加载一张照片来进行模型的预测。这里我们挑选置信度大于0.7的标签,并进行数据的可视化。
首先,我们加载一张需要进行检测的图片,并把图片传递给BETR模型进行对象的检测,对象检测完成后,我们就可以得到模型预测的对象标签,置信度,对象box信息。得到这些对象信息后,我们就可以进行数据可视化操作了,这里我们建立了一个plot_results函数,方便数据的可视化操作。
领取专属 10元无门槛券
私享最新 技术干货