GluonCV-物体检测教程-11.使用预先训练的CenterNet模型进行预测

本文展示了如何只用几行代码就可以使用经过预先训练的CenterNet模型。

首先，让我们导入一些必要的库：

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

加载预训练模型

让我们获得一个在Pascal VOC数据集上经过训练的CenterNet模型，并使用resnet18_v1b作为基础模型。通过指定 pretrained=True，它将在必要时从Model Zoo自动下载模型。有关更多预训练的模型，请参考 Model Zoo。

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

预处理图像

接下来，我们下载图像，并使用预设的数据转换进行预处理。在这里，我们指定将图像的短边尺寸调整为512像素。您可以提供任意大小的图像，但是，由于模型是使用512x512图像训练的，因此在某些输入分辨率下，模型的性能可能会更好。

您可以提供图像文件名的列表，如到，如果你想一起加载多个图像。[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.center_net.load_test(im_fname, short=512)

print('Shape of pre-processed image:', x.shape)

输出显示：

Shape of pre-processed image: (1, 3, 512, 683)

推理和显示

前向功能将返回所有检测到的边界框，以及相应的预测类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()

脚本的总运行时间：（0分钟0.974秒）

下载Python源码: demo_center_net.py

下载Jupyter notebook: demo_center_net.ipynb

---

翻译采用：谷歌在线翻译+人工修正。

原文链接：https://gluon-cv.mxnet.io/build/examples_detection/demo_center_net.html