一、代码运行环境
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()