使用Transformer模型进行计算机视觉任务的端对端对象检测

liftword11小时前技术文章1

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模型的框架,当然最好是源代码来实现,这个我们后期进行分享。

import math
from PIL import Image
import requests
import matplotlib.pyplot as plt
%config InlineBackend.figure_format = 'retina'
import ipywidgets as widgets
from IPython.display import display, clear_output
import torch
from torch import nn
from torchvision.models import resnet50
import torchvision.transforms as T
torch.set_grad_enabled(False);

首先我们需要导入需要的第三方库,这里主要是torch框架以及其他库。

CLASSES = [
    'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
    'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
    'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
    'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
    'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
    'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
    'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
    'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A',
    'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
    'toothbrush'
]
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]

DETR模型训练在COCO数据集上,这里我们列举一下COCO数据集的对象分类,并定义一个colors颜色列表,方便我们进行不同颜色的备注。

transform = T.Compose([
    T.Resize(800),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=1)
def rescale_bboxes(out_bbox, size):
    img_w, img_h = size
    b = box_cxcywh_to_xyxy(out_bbox)
    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
    return b
def plot_results(pil_img, prob, boxes):
    plt.figure(figsize=(16,10))
    plt.imshow(pil_img)
    ax = plt.gca()
    colors = COLORS * 100
    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
                                   fill=False, color=c, linewidth=3))
        cl = p.argmax()
        text = f'{CLASSES[cl]}: {p[cl]:0.2f}'
        ax.text(xmin, ymin, text, fontsize=15,
                bbox=dict(facecolor='yellow', alpha=0.5))
    plt.axis('off')
    plt.show()

然后我们定义几个函数,方便显示我们经过模型检测后的box,主要是为了我们后期进行box的画图,模型对象检测后的标签与模型检测置信度。

有了以上的函数,我们就可以把图片放入我们的DETR模型进行对象检测了。

model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True)
model.eval();
#url = 'http://images.cocodataset.org/val2017/00000003769.jpg'
#im = Image.open(requests.get(url, stream=True).raw)
im = Image.open('11.png').convert('RGB')

img = transform(im).unsqueeze(0)
outputs = model(img)
probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
keep = probas.max(-1).values > 0.8
bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)

plot_results(im, probas[keep], bboxes_scaled)

首先我们使用torch.hub.load函数来加载DETR模型的预训练模型,并传递一张图片。当然这里可以直接传递图片的URL地址,也可以直接从本地打开一张需要检测的模型。

加载图片后,我们需要把图片转换到torch变量,然后把图片数据传递到model模型来进行对象检测的预测,结果保持在output中,然后我们使用soft Max函数来挑选概率最大的对象检测数据,并挑选出置信度大于0.8(预设值)的对象标签与置信度。

得到检测结果后,最后我们使用plot_results函数来可视化检测到的对象。

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

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




相关文章

Python 和 R 数据分析/挖掘工具互查

写在前面在此总结一些在数据分析/挖掘中可能用到的功能,方便大家索引或者从一种语言迁移到另一种。当然,这篇博客还会随时更新(不会另起一篇,为了方便大家索引),请大家如果有需要收藏到书签中。如果大家已经熟...

Python 裸体检测

人脑知道如何区分穿衣服的人、穿比基尼的人和裸体的人。然后我有一个我无法立即回答的问题:一个简单的程序,几行代码就能知道照片是否裸体吗?让我们进入正题……一个可以知道你是否裸体的人工智能?Nudity...

化学方向:EI源刊实验数据写法

推荐期刊:Journal of Analytical Chemistry- 刊号:ISSN 1061-9348,CN 11-2345/O6- 影响因子:2.1(最新版)- 分区:中科院3区,JCR...

想学好Python数据分析,一定要掌握的重要模块之Statsmodels

Statsmodels是Python中一个功能强大且广泛使用的统计建模和经济计量学库。它允许用户构建各种统计模型,执行假设检验,进行预测,并对模型结果进行详细的诊断和可视化。以下是statsmodel...

「医学统计学之二」卡方检验之独立性检验:python的scipy实现

△ 概念:1 独立性检验:是对两个或两个以上因子彼此之间是否相互独立做检验时,就要用到卡方检验。2 列联表中变量独立性的卡方检验,采用python中scipy包的scipy.stats.chi2_co...

R数据分析:倾向性评分匹配实例操作

倾向性评分匹配是一种控制混杂的技巧,本文主要写操作,具体原理大家自己摸索哈。propensity score matching (PSM) is a “statistical matching tec...