在Unity上使用Sentis部署yolov8与添加NMS，并编译到手机上

paboluo

已于 2024-04-20 19:53:43 修改

阅读量3.5k

点赞数 15

文章标签： unity YOLO 游戏引擎

于 2024-04-20 19:46:07 首次发布

本文链接：https://blog.csdn.net/woshi212534/article/details/135705159

版权

本文介绍了如何在Unity中使用Sentis库对YOLOv8的ONNX模型进行改造，添加内置NMS层进行目标检测，以及在实际项目中处理和优化输出结果的过程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

项目地址

项目：https://github.com/212534/Unity-Sentis-YOLOv8

Demo apk：链接：https://pan.baidu.com/s/1agTZRhnCzgT5P5HtuUvgWQ?pwd=ydj7
提取码：ydj7
–来自百度网盘超级会员V5的分享

效果展示

这是在电脑上的测试，用的摄像头拍屏幕
在这里插入图片描述

安装Sentis

可以把Sentis看作Barracuda的升级版。
在Package里装com.unity.sentis
在这里插入图片描述

使用Sentis推理yolov8的onnx

这一部分比较简单，教程非常多

using System.Collections;
using System.Collections.Generic;
using Unity.Sentis;
using UnityEngine;
using Unity.Sentis.Layers;

public class Test: MonoBehaviour
{
   
    public ModelAsset modelAsset;
    private Model model;
    private IWorker worker;
    Ops ops;

    void Start()
    {
   
        model = ModelLoader.Load(modelAsset);
        worker = WorkerFactory.CreateWorker(BackendType.GPUCompute, model);
        ops = WorkerFactory.CreateOps(BackendType.GPUCompute, null);
    }

    public void Predict(WebCamTexture camImage) //我这里使用的是摄像头图像，你也可以用普通图片。
    {
   
        using Tensor inputImage = TextureConverter.ToTensor(camImage,width:640,height:640, channels: 3); //对输入的图像做处理

        var m_Inputs = new Dictionary<string, Tensor> 
        {
   
            {
   "images", inputImage }
        };

        worker.Execute(m_Inputs);//执行推理

        var output0 = worker.PeekOutput("output0") as TensorFloat;  //获取输出结果
        output0.MakeReadable(); //从GPU中取出数据，经过这一步之后就可以读取output0中的数据了
    }
}

直接把onnx文件拖上去就行了。
在这里插入图片描述

YOLOv8的输出格式

需要注意的是，yolov8的onnx的输出结果是一个1x84x8400的张量（当类别数不同时会有变化），具体含义可以参考这篇文章https://blog.csdn.net/yangkai6121/article/details/133843368
简单来讲，这里的8400是8400个预测框。而每个预测框包含以下信息：84 = 边界框预测4 + 数据集类别80
本文使用的yolov8是在coco数据集上训练的，包含80个类别，所以这里是4+80，具体情况与你的训练用的数据集类别有关。
1x84x8400这个输出很明显没法直接使用。所以需要加一个nms来后处理

NMS

yolo属于单阶段的目标检测算法，因此还要对预测框进行后处理。常见的做法是在Unity里用C#自己实现nms，但nms本身就是一个性能开销很大的后处理，自己用C#实现在性能上不划算。

如果是了解目标检测部署的同学，应该知道onnx在opset11及以上版本支持了内置nms层。那么理论上，只要在yolov8的检测头里把torch.ops.torchvision.nms加进去就可以导出自带nms层的ONNX模型。像下图一样
在这里插入图片描述
实际上经过测试，Sentis确实可以识别出onnx的内置nms层，但是会报错。

原因可能在于，Sentis内实现的nms层与torch.ops.torchvision.nms在input的shape上存在区别。

举例来讲，在python中对yolov8使用torch.ops.torchvision.nms时，它要求的boxes输入shape为8400x4，scores为8400.
但是Sentis实现的nms层要求的shape分别为，1x8400x4与1x1x8400。

使用Sentis给模型加入NMS层

既然torch.ops.torchvision.nms不能用，那我们可以在Unity里用代码把Sentis实现的nms层加到模型里。
先看一下Sentis中NonMaxSuppression()的输入参数。详细的可以看这里
https://docs.unity3d.com/Packages/com.unity.sentis@1.2/api/Unity.Sentis.Layers.NonMaxSuppression.html
在这里插入图片描述
大体来讲与torch.ops.torchvision.nms的输入参数比较一致。
我们需要解决的问题便是把yolov8输出的1x84x8400张量拆解为boxes与scores，也就是NonMaxSuppression()输入。而Sentis是支持在Unity中对模型结构进行修改的。

1.首先我们给模型加入三个新的输入，这也是NonMaxSuppression()所需要的输入。使用新增输入的形式设置这三个值，可以在程序运行时对它们动态修改。

		//Set input
        model.AddInput("maxOutputBoxesPerClass",DataType.Int, new SymbolicTensorShape(1)); //每个类别最多返回的边框数量。
        model.AddInput("iouThreshold",DataType.Float, new SymbolicTensorShape(1)); //iou阈值
        model.AddInput("scoreThreshold",DataType.Float, new SymbolicTensorShape(1)); //置信度阈值

2.对1x8400x4进行拆分。比较反人类的是，Sentis在添加层时，有些层的初始化只能输入字符串作为参数，这些字符串实际上是网络里的输出名，输入名，常量名等。所以这里要给模型加入几个常量。
我在注释里加入了每一层的shape变化，方便各位理解这个过程。如果你没搞过深度学习，不理解每一层的含义，可以去官方文档里看具体解释。
https://docs.unity3d.com/Packages/com.unity.sentis@1.2/api/Unity.Sentis.Layers.html

 		//Set constants   
        model.AddConstant(new Constant("0", new int[] {
    0 })

最低0.47元/天解锁文章