03. 使用预训练的YOLO模型进行预测

本文将介绍如何使用几行代码玩转预训练的YOLO模型。

首先导入一些必要的库

from gluoncv import model_zoo, data, utils
from matplotlib import pyplot as plt

加载预训练模型

让我们获取一个以Darknet53为基础模型，并在Pascal VOC数据集上训练的YOLOv3模型。通过指定 pretrained=True，如果需要，它将自动从模型库下载模型。有关更多预训练模型，请参阅 模型库。

net = model_zoo.get_model('yolo3_darknet53_voc', pretrained=True)

输出

Downloading /root/.mxnet/models/yolo3_darknet53_voc-f5ece5ce.zip from https://apache-mxnet.s3-accelerate.dualstack.amazonaws.com/gluon/models/yolo3_darknet53_voc-f5ece5ce.zip...

  0%|          | 0/223069 [00:00<?, ?KB/s]
  0%|          | 102/223069 [00:00<04:38, 801.70KB/s]
  0%|          | 510/223069 [00:00<01:40, 2207.73KB/s]
  1%|          | 2186/223069 [00:00<00:30, 7201.25KB/s]
  3%|2         | 6478/223069 [00:00<00:18, 11542.03KB/s]
  5%|4         | 10811/223069 [00:00<00:11, 17851.78KB/s]
  7%|6         | 15152/223069 [00:00<00:10, 20602.78KB/s]
  8%|7         | 17315/223069 [00:01<00:10, 18808.40KB/s]
 12%|#1        | 25747/223069 [00:01<00:05, 33517.27KB/s]
 15%|#4        | 32975/223069 [00:01<00:04, 42930.22KB/s]
 19%|#8        | 41469/223069 [00:01<00:03, 53749.08KB/s]
 22%|##1       | 48894/223069 [00:01<00:02, 59259.88KB/s]
 26%|##5       | 56986/223069 [00:01<00:02, 65015.99KB/s]
 29%|##9       | 64881/223069 [00:01<00:02, 68959.12KB/s]
 32%|###2      | 72142/223069 [00:01<00:02, 70008.54KB/s]
 36%|###5      | 79384/223069 [00:02<00:02, 51897.63KB/s]
 38%|###8      | 85416/223069 [00:02<00:02, 47837.93KB/s]
 42%|####1     | 92824/223069 [00:02<00:02, 53845.95KB/s]
 45%|####5     | 101182/223069 [00:02<00:01, 61173.82KB/s]
 49%|####8     | 108306/223069 [00:02<00:01, 63793.18KB/s]
 52%|#####1    | 115140/223069 [00:02<00:01, 63470.24KB/s]
 55%|#####5    | 123613/223069 [00:02<00:01, 69298.68KB/s]
 59%|#####8    | 131062/223069 [00:02<00:01, 69692.98KB/s]
 62%|######2   | 139408/223069 [00:02<00:01, 73597.55KB/s]
 66%|######5   | 146927/223069 [00:03<00:01, 65994.37KB/s]
 69%|######8   | 153775/223069 [00:03<00:01, 52338.29KB/s]
 73%|#######2  | 162414/223069 [00:03<00:01, 60219.84KB/s]
 76%|#######5  | 169073/223069 [00:03<00:00, 60525.86KB/s]
 79%|#######8  | 175576/223069 [00:03<00:00, 60083.12KB/s]
 82%|########2 | 183898/223069 [00:03<00:00, 66193.20KB/s]
 86%|########5 | 191415/223069 [00:03<00:00, 68650.94KB/s]
 89%|########9 | 199476/223069 [00:03<00:00, 72013.45KB/s]
 93%|#########2| 207018/223069 [00:03<00:00, 72920.62KB/s]
 96%|#########6| 214974/223069 [00:04<00:00, 74720.19KB/s]
100%|#########9| 222546/223069 [00:04<00:00, 70806.86KB/s]
223070KB [00:04, 53088.97KB/s]

预处理图像

接下来，我们下载一张图像，并使用预设的数据转换进行预处理。这里我们指定将图像的短边调整为512像素。您可以输入任意大小的图像。YOLO的一个约束是输入的高度和宽度必须能被32整除。

如果您想一起加载多张图像，可以提供图像文件名的列表，例如 [im_fname1, im_fname2, ...] 给 gluoncv.data.transforms.presets.yolo.load_test()。

此函数返回两个结果。第一个是形状为 (batch_size, RGB_channels, height, width) 的NDArray。它可以直接输入到模型中。第二个包含numpy格式的图像，方便绘制。由于我们只加载了一张图像，因此 x 的第一维是1。

im_fname = utils.download('https://raw.githubusercontent.com/zhreshold/' +
                          'mxnet-ssd/master/data/demo/dog.jpg',
                          path='dog.jpg')
x, img = data.transforms.presets.yolo.load_test(im_fname, short=512)
print('Shape of pre-processed image:', x.shape)

输出

Downloading dog.jpg from https://raw.githubusercontent.com/zhreshold/mxnet-ssd/master/data/demo/dog.jpg...

  0%|          | 0/160 [00:00<?, ?KB/s]
100%|##########| 160/160 [00:00<00:00, 27338.93KB/s]
Shape of pre-processed image: (1, 3, 512, 683)

推理和显示

forward函数将返回所有检测到的边界框以及相应的预测类别ID和置信度得分。它们的形状分别为 (batch_size, num_bboxes, 1)、(batch_size, num_bboxes, 1) 和 (batch_size, num_bboxes, 4)。

我们可以使用 gluoncv.utils.viz.plot_bbox() 来可视化结果。我们对第一张图像的结果进行切片，并将其输入到 plot_bbox 中

class_IDs, scores, bounding_boxs = net(x)

ax = utils.viz.plot_bbox(img, bounding_boxs[0], scores[0],
                         class_IDs[0], class_names=net.classes)
plt.show()