用PyTorch和Torchvision进行RetinaNet物体检测

后端之巅
436 阅读4分钟

简介

物体检测是计算机视觉的一个大领域，也是计算机视觉在 "野外 "的更重要的应用之一。一方面，它可以用来建立自主系统，在环境中为代理人导航--无论是执行任务的机器人还是自动驾驶汽车，但这需要与其他领域交叉。然而，异常检测（如生产线上的缺陷产品）、定位图像内的物体、面部检测和其他各种物体检测的应用都可以在不与其他领域交叉的情况下完成。

物体检测并不像图像分类那样标准化，主要是因为大多数新的发展通常是由个人研究者、维护者和开发者完成的，而不是大型库和框架。在TensorFlow或PyTorch这样的框架中打包必要的实用脚本，并维护迄今为止指导开发的API指南是很困难的。

这使得物体检测在某种程度上更加复杂，通常更加冗长（但并不总是这样），而且不如图像分类那样容易接近。处于一个生态系统中的主要好处之一是，它为你提供了一种方法，让你不需要搜索关于良好做法、工具和使用方法的有用信息。对于物体检测--大多数人必须对该领域的景观做更多的研究，以获得一个好的把握。

用PyTorch/TorchVision的RetinaNet进行物体检测

torchvision 是PyTorch的计算机视觉项目，旨在使基于PyTorch的CV模型的开发更容易，通过提供转换和增强脚本、带有预训练权重的模型Zoo、数据集和对从业者有用的工具。

虽然仍处于测试阶段，而且非常具有实验性--torchvision 提供了一个相对简单的物体检测API，并有一些模型可供选择：

  • 更快的R-CNN
  • 视网膜网络
  • FCOS（全卷积视网膜）。
  • SSD（VGG16主干...哎呀）
  • SSDLite (MobileNetV3骨干网)

虽然该API不像其他一些第三方API那样精致或简单，但对于那些仍然喜欢在他们熟悉的生态系统中的安全的人来说，它是一个非常体面的起点。在前进之前，请确保你安装了PyTorch和Torchvision：

$ pip install torch torchvision

让我们加载一些实用的函数，比如read_image() ，draw_bounding_boxes() 和to_pil_image() ，使其更容易阅读、绘制和输出图像，接着导入RetinaNet和其预训练的权重（MS COCO：

from torchvision.io.image import read_image
from torchvision.utils import draw_bounding_boxes
from torchvision.transforms.functional import to_pil_image
from torchvision.models.detection import retinanet_resnet50_fpn_v2, RetinaNet_ResNet50_FPN_V2_Weights

import matplotlib.pyplot as plt

RetinaNet使用的是ResNet50骨干网，并在其之上使用特征金字塔网络（FPN）。虽然这个类的名字很啰嗦，但它却说明了这个架构。让我们使用requests 库来获取一个图像，并将其保存为本地驱动器上的一个文件：

import requests
response = requests.get('https://i.ytimg.com/vi/q71MCWAEfL8/maxresdefault.jpg')
open("obj_det.jpeg", "wb").write(response.content)

img = read_image("obj_det.jpeg")

有了图像，我们就可以实例化我们的模型和权重：

weights = RetinaNet_ResNet50_FPN_V2_Weights.DEFAULT
model = retinanet_resnet50_fpn_v2(weights=weights, score_thresh=0.35)
# Put the model in inference mode
model.eval()
# Get the transforms for the model's weights
preprocess = weights.transforms()

score_thresh 参数定义了一个对象被检测为一个类别的对象的阈值。直观地说，它是信心阈值，如果模型对一个物体属于某个类别的信心低于35%，我们就不会把它归入某个类别。

让我们使用我们的权重的变换对图像进行预处理，创建一个批次并运行推理：

batch = [preprocess(img)]
prediction = model(batch)[0]

就这样，我们的prediction 字典持有推断出的物体类别和位置!现在，这个结果对我们来说不是很有用--我们要从权重中提取与元数据相关的标签，并绘制边界框，这可以通过draw_bounding_boxes()：

labels = [weights.meta["categories"][i] for i in prediction["labels"]]

box = draw_bounding_boxes(img, boxes=prediction["boxes"],
                          labels=labels,
                          colors="cyan",
                          width=2, 
                          font_size=30,
                          font='Arial')

im = to_pil_image(box.detach())

fig, ax = plt.subplots(figsize=(16, 12))
ax.imshow(im)
plt.show()

这样做的结果是：

RetinaNet实际上对偷看汽车后面的人进行了分类!这是个相当困难的分类。

你可以把RetinaNet换成FCOS（完全卷积的RetinaNet），用fcos_resnet50_fpn 来代替retinanet_resnet50_fpn_v2 ，并使用FCOS_ResNet50_FPN_Weights 的权重。

from torchvision.io.image import read_image
from torchvision.utils import draw_bounding_boxes
from torchvision.transforms.functional import to_pil_image
from torchvision.models.detection import fcos_resnet50_fpn, FCOS_ResNet50_FPN_Weights

import matplotlib.pyplot as plt
import requests
response = requests.get('https://i.ytimg.com/vi/q71MCWAEfL8/maxresdefault.jpg')
open("obj_det.jpeg", "wb").write(response.content)

img = read_image("obj_det.jpeg")
weights = FCOS_ResNet50_FPN_Weights.DEFAULT
model = fcos_resnet50_fpn(weights=weights, score_thresh=0.35)
model.eval()

preprocess = weights.transforms()
batch = [preprocess(img)]
prediction = model(batch)[0]

labels = [weights.meta["categories"][i] for i in prediction["labels"]]

box = draw_bounding_boxes(img, boxes=prediction["boxes"],
                          labels=labels,
                          colors="cyan",
                          width=2, 
                          font_size=30,
                          font='Arial')

im = to_pil_image(box.detach())

fig, ax = plt.subplots(figsize=(16, 12))
ax.imshow(im)
plt.show()

结语

物体检测是计算机视觉的一个重要领域，但不幸的是，这个领域并不像它应该的那样容易接近。

在这个简短的指南中，我们看了一下torchvision ，PyTorch的计算机视觉软件包，如何使用RetinaNet更容易地对图像进行物体检测。

avatar
文章
阅读
粉丝