transformer模型刚发布时，主要应用于处理NLP领域任务，比如机器翻译等，但是随着注意力机制模型的大火，很多基于transformer模型的魔改模型也相继发布，且transformer模型的注意力机制也被google团队证明可以使用在计算机视觉任务上，特别是swin transformer模型的发布，更是把transformer模型带入了计算机视觉领域。前几篇文章，我们也同样介绍了另外一个基于transformer模型的，应用在计算机视觉任务上的模型BETR。且BETR模型不仅可以使用在对象检测上，还可以使用在对象分割上，本期我们就基于transformer模型来代码实现一下BETR模型。
from PIL import Image
import requests
import matplotlib.pyplot as plt
import torch
from torch import nn
from torchvision.models import resnet50
import torchvision.transforms as T
torch.set_grad_enabled(False);
基于transformer模型的对象检测算法代码实现的第一步是需要我们import python的第三方库，这里主要是torch，确保你在运行本期代码前，已经成功安装了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]]
这里我们需要建立一个list列表，一个是保存所有BETR模型能够检测的对象标签，另外一个是颜色数据，方便我们后期的可视化。
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
这里我们建立几个函数，方便我们进行对象的标注，后期进行可视化操作时比较方便。以上初始化的工作完成后，我们就可以搭建我们的BETR模型了。
根据BETR模型的框架图，我们知道这里有2层关键的框架，一个是CNN卷积神经网络层，另外一个就是transformer模型的编码器与解码器部分。
class DETR_model(nn.Module):
def __init__(self, num_classes, hidden_dim=256, nheads=8,num_encoder_layers=6, num_decoder_layers=6):
super().__init__()
# CNN卷积神经网络层
self.backbone = resnet50()
del self.backbone.fc
self.conv = nn.Conv2d(2048, hidden_dim, 1)
# transformer 层
self.transformer = nn.Transformer(hidden_dim, nheads, num_encoder_layers, num_decoder_layers)
# 预测class与box层
self.linear_class = nn.Linear(hidden_dim, num_classes 1)
self.linear_bbox = nn.Linear(hidden_dim, 4)
# 输出位置编码
self.query_pos = nn.Parameter(torch.rand(100, hidden_dim))
# 输出空间的位置编码
self.row_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))
self.col_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))
def forward(self, inputs):
# ResNet-50 CNN卷积神经网络
x = self.backbone.conv1(inputs)
x = self.backbone.bn1(x)
x = self.backbone.relu(x)
x = self.backbone.maxpool(x)
x = self.backbone.layer1(x)
x = self.backbone.layer2(x)
x = self.backbone.layer3(x)
x = self.backbone.layer4(x)
# 从 2048 转换到 256 feature
h = self.conv(x)
#positional encodings位置编码
H, W = h.shape[-2:]
pos = torch.cat([self.col_embed[:W].unsqueeze(0).repeat(H, 1, 1),self.row_embed[:H].unsqueeze(1).repeat(1, W, 1),], dim=-1).flatten(0, 1).unsqueeze(1)
# transformer 层
h = self.transformer(pos 0.1 * h.flatten(2).permute(2, 0, 1),self.query_pos.unsqueeze(1)).transpose(0, 1)
# 最终输出labels与 boxes
return {'pred_logits': self.linear_class(h), 'pred_boxes': self.linear_bbox(h).sigmoid()}
搭建好了BETR模型，我们就可以直接来进行使用了，在使用之前，我们需要下载BETR的预训练模型。
detr = DETR_model(num_classes=91)
state_dict = torch.hub.load_state_dict_from_url(
url='https://dl.fbaipublicfiles.com/detr/detr_demo-da2a99e9.pth',
map_location='cpu', check_hash=True)
detr.load_state_dict(state_dict)
detr.eval();
这里我们已经搭建好了DETR_model模型，然后我们需要使用torch.hub.load_state_dict_from_url函数来下载BETR的预训练模型，模型下载完成后，我们使用load_state_dict来加载模型，并进行eval。然后我们就可以使用BETR模型了。
def detect(im, model, transform):
img = transform(im).unsqueeze(0)
assert img.shape[-2] <= 1600 and img.shape[-1] <= 1600, ' 模型最大支持的图片1600*1600'
outputs = model(img)
probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
keep = probas.max(-1).values > 0.7
bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)
return probas[keep], bboxes_scaled
这里我们建立一个detect函数，方便我们使用模型进行图片的对象检测，然后我们就可以加载一张照片来进行模型的预测。这里我们挑选置信度大于0.7的标签，并进行数据的可视化。
im = Image.open('11.jpg').convert('RGB')
scores, boxes = detect(im, detr, transform)
def plot_results(pil_img, prob, boxes):
plt.figure(figsize=(16,10))
plt.imshow(pil_img)
ax = plt.gca()
for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), COLORS * 100):
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()
plot_results(im, scores, boxes)
首先，我们加载一张需要进行检测的图片，并把图片传递给BETR模型进行对象的检测，对象检测完成后，我们就可以得到模型预测的对象标签，置信度，对象box信息。得到这些对象信息后，我们就可以进行数据可视化操作了，这里我们建立了一个plot_results函数，方便数据的可视化操作。
transformer模型是google在attention is all you need论文中提出的一个应用于NLP领域的模型，但是随着VIT模型的发布，把transformer模型应用到计算机视觉任务上成为了可能，本期介绍的DETR模型就是基于transformer模型与CNN卷积神经网络的结合打造出的一个对象检测模型，从对象检测的结果来看，其使用transformer模型的对象检测方案，也能跟CNN卷积神经网络媲美。
