Pytorch---Pytorch使用内置的目标检测算法进行目标检测

一、代码运行环境

Pytorch-gpu==1.10.1
Python==3.8

二、实现代码

import os
import numpy as np
import torch
import torchvision
from PIL import Image
from torchvision.utils import draw_bounding_boxes
import matplotlib.pyplot as plt

# 加载图片
pil_img = Image.open(os.path.join('test.jpg'))
np_img = np.array(pil_img)
tensor_img = torch.from_numpy(np_img / 255).permute(2, 0, 1).type(torch.float)

# 加载模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()

# 开始进行预测
pred = model(torch.unsqueeze(input=tensor_img, dim=0))
boxes = pred[0]['boxes']
labels = pred[0]['labels']
scores = pred[0]['scores']
threshold = 0.86
pred_index = scores > threshold
boxes = boxes[pred_index]
labels = labels[pred_index]

# 结果的展示
result = draw_bounding_boxes(image=torch.as_tensor(data=tensor_img * 255, dtype=torch.uint8),
                             boxes=boxes,
                             labels=['horse', 'car', 'car', 'person', 'car'],
                             colors=['red', 'green', 'green', 'blue', 'green'])
plt.axis('off')
plt.imshow(result.permute(1, 2, 0).numpy())
plt.savefig('result.png')
plt.show()

三、运行结果