用于增强现实的实时可穿带目标检测：基于YOLOv8进行ONNX转换和部署

计算机视觉研究院

发布于 2023-08-23 13:21:48

6460

发布于 2023-08-23 13:21:48

文章被收录于专栏：计算机视觉战队计算机视觉战队

从此不迷路

计算机视觉研究院

Column of Computer Vision Institute

今天给大家介绍了一种在增强现实（AR）环境中使用机器学习（ML）进行实时目标检测的软件体系结构。

01

前景概要

我们的方法使用了最近最先进的YOLOv8网络，该网络在Microsoft HoloLens 2头戴式显示器（HMD）上运行。这项研究背后的主要动机是通过可穿戴、免提的AR平台，应用先进的ML模型来增强感知和情景感知。我们展示了YOLOv8模型的图像处理流水线，以及在耳机的资源有限的边缘计算平台上使其实时的技术。实验结果表明，我们的解决方案在不需要将任务卸载到云或任何其他外部服务器的情况下实现了实时处理，同时在通常的mAP度量和测量的定性性能方面保持了令人满意的准确性。

02

背景&动机

属于沉浸式技术类别的增强现实（AR）技术通过在用户视野（FoV）中叠加数字内容，提供了将数字伪像与物理环境融合的能力（如下图）。

目前，在智能手机或平板电脑等移动设备上运行的流行AR应用程序可以通过机器学习（ML）进一步增强。得益于这种方法，我们可以在视频和图像数据上包括基于视觉的目标检测和跟踪特征。然而，移动AR解决方案有很大的局限性，例如相对较小的FoV受到屏幕画布的限制或需要手动控制。后者缩小了我们可以成功部署AR的潜在场景，如手动组装、设备修复任务或老年人使用AR增强剂。在这种情况下，用户不仅能够自由移动双手，而且能够迅速改变无约束的FoV或身体姿势，这对于安全问题和任务完成至关重要。

可穿戴智能头戴式显示器（HMD）的替代技术规避了这些警告。AR耳机，如被广泛认为是最先进的Microsoft HoloLens 2（HL2），提供免提AR体验。不幸的是，HL2和其他类似的耳机并不能为基于ML的处理提供令人满意的支持，这可以增强用户与环境交互的能力。因此，在耳机的边缘计算平台中运行板载实时ML模型对于开发新的AR应用领域至关重要。

我们解决了HL2上的实时对象检测问题，包括最新的YOLOv8框架。我们专注于定义必要的步骤，以实现板载ML模型所需的图像处理帧速率，同时确定HL2计算平台的约束条件。克服这些限制使得能够在耳机上使用广泛可用的ML算法。我们还相信，AR开发人员可以利用我们在YOLOv8 for HL2上的工作来创建新的应用程序，扩展该耳机的当前用例。

03

新框架设计

整体的思路如下图所示。我们首先为HL2准备YOLOv8神经网络模型。这些模型可以选择性地重新训练（微调）以包括不同的检测类别。下一步涉及将模型导出为开放神经网络交换（ONNX）格式。然后，Unity引擎中的Barracuda库使用该模型对HL2执行目标检测，并提供检测到的对象的可视化。我们决定使用Unity平台，因为它是AR和VR（虚拟现实）研究中使用最广泛的软件框架之一。

用于ML推理的Barracuda库

检测管道的神经网络部分基于Barracuda库。它是Unity开发的一个开源库，用于在游戏引擎中使用神经网络。它支持最常见的深度学习层，并提供GPU和CPU推理引擎。通过使用ONNX格式加载预训练的神经网络，确保了对不同机器学习库的跨框架支持。它实现了不同ML框架之间的互操作性，提供了一组用于深度学习的标准操作。

模型准备

在线操作中使用的每个模型都可以使用相同的管道进行准备。我们将PyTorch序列化的.pt文件中的每个模型导出为ONNX格式。由于当前的Barracuda版本支持高版本9的ONNX深度学习操作（opset），因此导出具有正确opset标志的模型至关重要。除了导出之外，还可以使用ONNX简化工具缩小模型。该操作使用常数折叠合并冗余运算符，从而加快推理速度。我们成功测试了导出和部署公开可用的原始YOLOv8目标检测模型。此外，我们可以为任何具有足够数据的自定义类训练YOLOv8，同时遵循对自定义数据集进行模型微调的指导原则。

HL2上的目标检测流水线