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

在本系列中,我们将使用预训练的模型来创建一个iOS应用程序,该应用程序将在实时摄像头Feed中(而非静态图片中)检测多个人和物体。
在本文中,我们将从YOLO对象检测模型从开放神经网络交换格式(ONNX)转换为Core ML开始。
下载源350.1 KB
介绍
在上一篇文章系列中,我们创建了一个用于图像分类的iOS应用程序。但是,在许多情况下,仅给图片分配一个标签是不够的。例如,在安全系统中,您不仅想知道图片中是否有手提箱,还想知道手提箱在哪里,以及它是否被人包围或被遗弃。

对象检测模型可以帮助我们实现这一目标。在本系列中,我们将使用预训练的模型来创建一个iOS应用程序,该应用程序将在实时摄像头Feed中(而非静态图片中)检测多个人和物体。

本系列假定您熟悉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 comer$ jupyter notebook123

在本文中,我们将从YOLO对象检测模型从开放式神经网络交换格式(ONNX)到Core ML的转换开始。尽管您可以从头开始训练自己的对象检测器,但这将是一项令人沮丧的练习。除非您拥有庞大的训练数据集和数千美元的GPU时间花费,否则使用预训练的YOLO对象检测模型将获得更加准确的结果。

使用YOLO进行物体检测
对象检测模型接受图像作为输入,并返回带有相应标签和边界框的检测对象列表。

上面的图像取自Open Images数据集。

在本系列中,我们将使用YOLO架构,更精确地说,该架构是在文章“ YOLO9000:更好,更快,更强”或YOLO v2中首次引入的。YOLO代表“您只看一次”,它反映了其基本思想:图像数据通过神经网络的单次传递即可产生检测结果。许多文章都在详细解释YOLO的工作方式。例如“ YOLO:实时对象检测”或“使用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 osimport coremltools as ctimport onnximport urllibonnx_model_path = './models/yolov2-coco-9.onnx'with open(onnx_model_path, 'rb') as f:    model_onnx = onnx.load(f)123456789

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

print(model_onnx.graph.input)1

因此,输入节点的名称为“ 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', )123456789

确保为该image_input_names参数分配一个元素列表。如果您错过那里的方括号,一切似乎都可以正常工作,但是preprocessing_args在预测期间,由参数定义的输入比例将被忽略。它看起来似乎不多,但是确实如此。它弄乱了模型返回的结果,您可能需要花费很长时间才能找出问题出在哪里,并质疑世界的确定性。我们去过那里!

转换完成后,您应该在日志中看到一个确认:

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

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

很好:输入是尺寸为416 x 416的RGB图像。现在让我们快速检查是否可以将其用于预测。为此,我们需要将图像裁剪并缩放为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))12345678

image = load_and_scale_image(‘https://www.xiaoyuani.com/4/3393/3436245648_c4f76c0a80_o.jpg’)
pred = cml_model.predict(data={INPUT_NODE: image})

如您所见,模型的输出为数组形状(1、425、13、13)。第一个维度代表批次,后两个维度代表我们上面提到的13 x 13 YOLO网格。但是每个单元格中的“ 425”值如何?继续阅读本系列的下一篇文章以进行查找。

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

(0)

相关推荐