SSD

模型描述

此 SSD300 模型基于SSD: Single Shot MultiBox Detector论文,该论文将 SSD 描述为“一种使用单个深度神经网络检测图像中物体的方法”。输入大小固定为 300×300。

该模型与论文中描述的模型的主要区别在于主干网络。具体来说,VGG 模型已过时,被 ResNet-50 模型取代。

根据现代卷积目标检测器的速度/精度权衡论文,对主干网络进行了以下增强:

  • conv5_x、avgpool、fc 和 softmax 层已从原始分类模型中移除。
  • conv4_x 中的所有步长均设置为 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'
]

格式化图像以符合网络输入并将其转换为张量。

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

参考文献

用于目标检测的单发多框检测器模型

模型类型: 视觉
提交者: NVIDIA