基于transformer模型的对象检测算法-BETR模型

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函数，方便数据的可视化操作。