在Pytorch提供的已经训练好的图像目标检测中,均是R-CNN系列的网络,并且针对目标检测和人体关键点检测分别提供了容易调用的方法。针对目标检测的网络,输入图像均要求使用相同的预处理方式,即先将每张图像的像素值预处理到0~1之间,且输入的图像尺寸不是很小即可直接调用。已经预训练的可供使用的网络模型如下所示:
网络类 | 描述 |
detection.fasterrcnn_resnet50_fpn | 具有ResNet-50-FPN结构的Fast R-CNN网络模型 |
detection.maskrcnn_resnet50_fpn | 具有ResNet-50-FPN结构的Mask R-CNN网络模型 |
detection.keypointrcnn_resnet50_fpn | 具有ResNet-50-FPN结构的Keypoint R-CNN网络模型 |
这些网络同样是在COC O2017数据上进行训练的。下面分别展示如何使用已经训练好的网络进行图像目标检测以及人体关键点检测。首先导入相关库和模块,程序如下:
import numpy as np
import torchvision
import torch
import torchvision.transforms as transforms
from PIL import Image,ImageDraw
import matplotlib.pyplot as plt
import matplotlib.image as mping
在进行图像目标检测时,使用已经预训练好的具有ResNet-50-FPN结构的Fast-R-CNN模型,该网络同样是通过COCO数据集进行与训练,导入已与训练的网络,程序如下所示:
model=torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()
下面从文件夹中读取一张照片,并将其转化为张量,像素值在0~1之间,然后使用导入模型对其进行预测,程序如下:
image=Image.open("data/chap10/2012_004308.jpg")
transform_d=transforms.Compose([transforms.ToTensor()])
image_t=transform_d(image)
pred=model([image_t])
# print(pred)
运行结果如下:
在pred输出的结果中主要包括检测到每个目标的边界框(boxes坐标).下面将检测到的目标可视化,并观察检测的具体结果。
首先定义每个类别所对应的标签COCO_INSTANCE_CATEGORY_NAMES,程序如下:
COCO_INSTANCE_CATEGORY_NAMES=[
'__background__','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'
]
针对预测结果,在可视化之前,需要分析有效的预测目标数据解读出来,需要提取的信息有个目标的位置、类别和得分,然后将得分大于0.5的目标作为检测到的有效目标,并将检测到的目标在图像上显示出来,程序如下:
pred_class=[COCO_INSTANCE_CATEGORY_NAMES[ii] for ii in list(pred[0]['labels'].numpy())]
pred_score=list(pred[0]['scores'].detach().numpy())
pred_boxes=[[ii[0],ii[1],ii[2],ii[3]] for ii in list (pred[0] ['boxes'].detach().numpy())]
pred_index=[pred_score.index(x) for x in pred_score if x >0.5]
# fontsize=np.int16(image.size[1]/30)
# font1=ImageFont.truetype
draw=ImageDraw.Draw(image)
for index in pred_index:
box=pred_boxes[index]
draw.rectangle(box,outline="red")
texts=pred_class[index]+":"+str(np.round(pred_score[index],2))
draw.text((box[0],box[1]),texts,fill="red")
image.show()
上面的程序在可视化图像时,使用ImageDraw.Draw(image)方法,表示要在原始的image图像上相应位置添加一些元素,draw.rectangle()表示要添加矩形框,draw.text()表示在图像上的制定位置添加文本。运行结果如下: