65.9K
CodeProject 正在变化。 阅读更多。
Home

将 ONNX 对象检测模型转换为 iOS Core ML

Jarek Szczegielniak
starIconstarIconstarIconstarIconstarIcon

5.00/5 (7投票s)

2020 年 11 月 20 日

CPOL

5分钟阅读

viewsIcon

9718

downloadIcon

139

在本系列文章中，我们将使用一个预训练模型来创建一个 iOS 应用程序，该应用程序可以在实时摄像头馈送中检测多个人和物体，而不是在静态图片中。

引言

在前一篇文章系列中，我们创建了一个用于图像分类的 iOS 应用程序。然而，在许多情况下，为图片分配单个标签是不够的。例如，在安全系统中，您不仅想知道图片中某个地方是否有行李箱，还想知道它在哪里，以及它是否被人群包围或被遗弃。

对象检测模型可以帮助我们实现这一点。在本系列文章中，我们将使用一个预训练模型来创建一个 iOS 应用程序，该应用程序可以在实时摄像头馈送中检测多个人和物体，而不是在静态图片中。

本系列假设您熟悉 Python、Conda 和 ONNX，并且有在 Xcode 中开发 iOS 应用程序的经验。欢迎您 下载此项目的源代码。我们将使用 macOS 10.15+、Xcode 11.7+ 和 iOS 13+ 来运行代码。

要创建 coreml Conda 环境，请使用下载代码中的 env-coreml.yml 文件和以下命令

$ conda env create -f env-coreml.yml
$ conda activate coreml
$ jupyter notebook

在本文中，我们将开始将 Open Neural Network Exchange 格式（ONNX）的 YOLO 对象检测模型转换为 Core ML。虽然您可以从头开始训练自己的对象检测器，但这将是一项令人沮丧的练习。除非您拥有海量训练数据集并花费大量资金购买 GPU 时间，否则使用预训练的 YOLO 对象检测模型将获得更准确的结果。

使用 YOLO 进行对象检测

对象检测模型接受图像作为输入，并返回已检测到的对象列表，其中包含相应的标签和边界框。

上图 来自 Open Images 数据集。

在本系列文章中，我们将使用 YOLO 架构，该架构首次在文章 "YOLO9000: Better, Faster, Stronger" 中引入，更确切地说，是 YOLO v2。YOLO 是 "You Only Look Once" 的缩写，它反映了其基本思想：图像数据通过神经网络的单次传递即可产生检测结果。许多文章详细解释了 YOLO 的工作原理，例如 "YOLO: Real-Time Object Detection" 或 "Real-time object detection with YOLO"。在这里，我们将专注于基本原理。

YOLO 将输入图像划分为一个网格。对于每个网格单元，它会计算带有相应对象标签的多个框。为了使其更容易（对模型而言，而不是对使用它的人而言），它不计算绝对框坐标，而是计算一组用于缩放预定义“锚点”框的因子。

我们将在本文系列中使用的 YOLO v2 使用一个包含 169 个单元（13 x 13）的网格，每个单元包含五个框。总而言之，模型为每张图像处理 845 个边界框。

在深入研究 YOLO 的输出之前，让我们先关注第一步：将 YOLO 模型转换为 Core ML 格式。

查找合适的模型

在 ONNX Model Zoo 中有许多可用的模型，尽管并非所有模型都同样易于使用或转换，这主要是由于缺乏文档。在几次尝试转换 YOLO v4 或 v3 失败后，我们选择了 YOLO v2。这是一个宝贵的教训：即使 Core ML 支持源 ONNX 模型声明的相同操作集（"opset"）版本，转换也可能失败。

所选模型是在 COCO 数据集上训练的，该数据集包含 80 种不同的对象类别，如人、狗、刀等。

将 ONNX 转换为 Core ML

一旦模型已下载到 onnx_model_path 变量中存储的路径，我们就可以加载它

import os
import coremltools as ct
import onnx
import urllib

onnx_model_path = './models/yolov2-coco-9.onnx'

with open(onnx_model_path, 'rb') as f:
    model_onnx = onnx.load(f)

接下来，让我们检查模型的输入

print(model_onnx.graph.input)

因此，输入节点的名称是“input.1”，它接受一个维度为 (1, 3, 416, 416) 的数组，这对应于一批 3 通道（RGB）图像，每张图像为 416 x 416 像素。

ONNX Model Zoo 页面上该特定模型的文档非常含糊。然而，由于 YOLO v3 和 v4 的可用预处理代码都将输入值缩放到 [0 - 1] 范围，因此我们假设 v2 也一样。现在我们可以运行转换了

cml_model = ct.converters.onnx.convert(
    model=onnx_model_path,
    image_input_names=['input.1'],
    preprocessing_args={
        'image_scale': 1/255.0,
            'is_bgr': False
    },
    minimum_ios_deployment_target='13', 
)

确保您为 image_input_names 参数分配一个单元素列表。如果您在那里省略了方括号，一切都会看起来正常工作，但 preprocessing_args 参数定义的输入缩放将在预测期间被静默忽略。这看起来可能没什么大不了——但实际上很重要。这会弄乱模型返回的结果，您可能会花很长时间试图弄清楚出了什么问题，并质疑世界的确定性。我们有过这样的经历！

当转换完成后，您应该会在日志中看到确认信息

我们可以保存转换后的模型并查看其详细信息

cml_model_path = './models/yolov2-coco-9.mlmodel'
cml_model.save(cml_model_path)
print(cml_model)

很好：输入是一个 RGB 图像，尺寸为 416 x 416。现在让我们快速检查一下是否可以使用它进行预测。为此，我们需要将图像裁剪并缩放到 416 x 416 像素

def load_and_scale_image(image_url):
    image = Image.open(urlopen(image_url))
    w,h = image.size
    min_dim = min(w,h)
    x0 = int((w - min_dim)/2)
    y0 = int((h - min_dim)/2)
    box = (x0, y0, x0 + min_dim, y0 + min_dim)
    return image.crop(box=box).resize((416,416))

image = load_and_scale_image('https://c2.staticflickr.com/4/3393/3436245648_c4f76c0a80_o.jpg')
pred = cml_model.predict(data={INPUT_NODE: image})

如您所见，模型的输出是一个形状为 (1, 425, 13, 13) 的数组。第一个维度代表批次，最后两个维度代表我们上面提到的 13 x 13 的 YOLO 网格。但是每个单元格中的“425”值是什么意思呢？请继续阅读本系列的下一篇文章以了解更多信息。

后续步骤

在本系列的 下一篇文章中，我们将讨论如何将输出值转换为检测到的对象标签、置信度分数和相应的框。

历史

  • 2020 年 11 月 20 日：初版
© . All rights reserved.