SSD


模型描述
此 SSD300 模型基于 SSD: Single Shot MultiBox Detector 论文,该论文将 SSD 描述为“一种使用单个深度神经网络检测图像中物体的方法”。其输入大小固定为 300×300。
该模型与原论文所述模型的主要区别在于主干网络(backbone)。具体而言,VGG 模型已被弃用,并由 ResNet-50 模型所取代。
根据 Speed/accuracy trade-offs for modern convolutional object detectors (现代卷积目标检测器的速度/精度权衡)论文,主干网络进行了以下增强:
- 从原始分类模型中移除了 conv5_x、avgpool、fc 和 softmax 层。
- conv4_x 中的所有步长(strides)均设置为 1×1。
主干网络之后添加了 5 个额外的卷积层。除卷积层外,我们还附加了 6 个检测头。
- 第一个检测头连接到最后一个 conv4_x 层。
- 其余五个检测头分别连接到对应的 5 个额外层。
检测头与论文中引用的类似,但我们在每次卷积后增加了额外的 BatchNorm 层,从而对其进行了增强。
示例
在下面的示例中,我们将使用预训练的 SSD 模型来检测样本图像中的物体并可视化结果。
要运行此示例,您需要安装一些额外的 Python 包。这些包用于图像预处理和可视化。
pip install numpy scipy scikit-image matplotlib
加载在 COCO 数据集上预训练的 SSD 模型,以及用于方便且全面地格式化模型输入和输出的一组实用方法。
import torch
ssd_model = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd')
utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd_processing_utils')
现在,准备好加载的模型以进行推理。
ssd_model.to('cuda')
ssd_model.eval()
准备用于目标检测的输入图像。(以下示例链接对应于 COCO 数据集中的前几张测试图像,但您也可以在此处指定本地图像的路径)
uris = [
'http://images.cocodataset.org/val2017/000000397133.jpg',
'http://images.cocodataset.org/val2017/000000037777.jpg',
'http://images.cocodataset.org/val2017/000000252219.jpg'
]
格式化图像以符合网络输入要求,并将其转换为张量(tensor)。
inputs = [utils.prepare_input(uri) for uri in uris]
tensor = utils.prepare_tensor(inputs)
运行 SSD 网络以执行目标检测。
with torch.no_grad():
detections_batch = ssd_model(tensor)
默认情况下,SSD 网络对每张输入图像的原始输出包含 8732 个带有定位和类别概率分布的边界框。让我们过滤这些输出,以更全面的格式仅获取合理的检测结果(置信度 > 40%)。
results_per_input = utils.decode_results(detections_batch)
best_results_per_input = [utils.pick_best(results, 0.40) for results in results_per_input]
该模型是在 COCO 数据集上训练的,我们需要访问该数据集才能将类别 ID 转换为物体名称。首次下载标注信息可能需要一些时间。
classes_to_labels = utils.get_coco_object_dictionary()
最后,让我们可视化检测结果。
from matplotlib import pyplot as plt
import matplotlib.patches as patches
for image_idx in range(len(best_results_per_input)):
fig, ax = plt.subplots(1)
# Show original, denormalized image...
image = inputs[image_idx] / 2 + 0.5
ax.imshow(image)
# ...with detections
bboxes, classes, confidences = best_results_per_input[image_idx]
for idx in range(len(bboxes)):
left, bot, right, top = bboxes[idx]
x, y, w, h = [val * 300 for val in [left, bot, right - left, top - bot]]
rect = patches.Rectangle((x, y), w, h, linewidth=1, edgecolor='r', facecolor='none')
ax.add_patch(rect)
ax.text(x, y, "{} {:.0f}%".format(classes_to_labels[classes[idx] - 1], confidences[idx]*100), bbox=dict(facecolor='white', alpha=0.5))
plt.show()
详细信息
有关模型输入和输出、训练方法、推理和性能的详细信息,请访问: github 和/或 NGC
参考文献