模型描述

此 SSD300 模型基于SSD: Single Shot MultiBox Detector论文，该论文将 SSD 描述为“一种使用单个深度神经网络检测图像中对象的方法”。输入尺寸固定为 300x300。

此模型与论文中描述的模型的主要区别在于其骨干网络（backbone）。具体来说，VGG 模型已过时，并被 ResNet-50 模型取代。

根据Speed/accuracy trade-offs for modern convolutional object detectors论文，对骨干网络进行了以下改进

  • 从原始分类模型中移除了 conv5_x、avgpool、fc 和 softmax 层。
  • conv4_x 中的所有步长均设置为 1x1。

骨干网络之后是 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

参考资料