DETR的demo使用方法详解（python）

1. 数据库调用

import torch as th
import torchvision.transforms as T
import requests
from PIL import Image, ImageDraw, ImageFont

2.加载DETR模型，并且下载预训练参数，调用GPU。

model = th.hub.load("facebookresearch/detr", "detr_resnet101", pretrained=True)
model.eval()
model = model.cuda()

3. 图像预处理和分类标签。

transform = T.Compose([
        T.ToTensor(),
        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

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']

4. 加载图片，并改变图片大小，将图片设置为RGB格式。（因为有的图可能是灰度图）

url = 'https://img2.baidu.com/it/u=3726792407,179742759&fm=253&fmt=auto&app=138&f=JPEG?w=500&h=367'
img = Image.open(requests.get(url, stream=True).raw).resize((800, 600)).convert('RGB')
# img.show()

5.对图片进行预处理，并且增加一个维度，调用GPU。不需要梯度下降

img_tens = transform(img).unsqueeze(0).cuda()
with th.no_grad():
    output = model(img_tens)

 6. 复制图片用于展示，并做出预测

im2 = img.copy()
drw = ImageDraw.Draw((im2))
pred_logits = output['pred_logits'][0][:, : len(CLASSES)]
pred_boxes = output['pred_boxes'][0]
max_output = pred_logits.softmax(-1).max(-1)
topk = max_output.values.topk(5) # 只画数值最高的五个框，其中数字可以改
pred_logits = pred_logits[topk.indices]
pred_boxes = pred_boxes[topk.indices]

 7求取坐标，并画图

for logits, box in zip(pred_logits, pred_boxes):
    cls = logits.argmax()
    if cls >= len(CLASSES):
        continue
    label = CLASSES[cls]
    box = box.cpu() * th.Tensor([800, 600, 800, 600])
    x, y, w, h = box
    x0, x1 = x-w//2, x+w//2
    y0, y1 = y-h//2, y+h//2
    drw.rectangle([x0, y0, x1, y1], outline="red", width=5)
    drw.text((x, y), label, fill="white")
    break

im2.show()

结果展示：

输入的图片：

 输出的图片：

 如有侵权请联系。