在Unity3d使用yolov8的onnx模型

爱工作不爱上班

已于 2023-08-18 11:09:12 修改

阅读量4.6k

点赞数 3

文章标签： unity yolo

于 2023-08-17 20:22:36 首次发布

本文链接：https://blog.youkuaiyun.com/shenqingyu0605202324/article/details/132348321

版权

本文介绍了如何将YOLO模型转换为ONNX格式，并在Unity3D中使用Barracuda库解析模型输出，提取最高置信度的目标及其相关信息。

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

下载yolo以及转成onnx模型：https://github.com/ultralytics/ultralytics

(我转换onnx的命令：yolo export model=[自己的模型名字] format=onnx opset=10)

在以下网址查看onnx模型的结构：https://netron.app/

可以看到输入INPUTS是float32[1,3,640,640]，它的输入张量形状为 tensor: float32[1, 3, 640, 640]，其中的四个数字分别表示不同的维度。第一个数字 1 表示批次（batch）的大小。这表示模型一次可以处理的样本数量。在这种情况下，模型一次处理一个样本。第二个数字 3 表示通道（channel）的数量。对于彩色图像，通常有三个通道（红色、绿色、蓝色），分别对应 RGB 通道。每个通道包含图像的不同颜色信息。第三个数字 640 表示图像的高度（height）。这表示输入图像的垂直像素数量。第四个数字 640 表示图像的宽度（width）。这表示输入图像的水平像素数量。这里输入是一张640*640的图片。

输出OUTPUTS是float32[1,35,8400]，表示它会输出一个大小为[1, 35, 8400]的浮点数张量。这个输出张量的维度表示了检测到的目标的数量、每个目标的预测参数以及网格单元的数量。输出张量的第一个维度1表示检测到的目标数量。每个目标由35个值组成，表示目标的类别、置信度、边界框的坐标及10个特定点的位置。最后一个维度8400表示输入图像被划分成的网格单元数量。

在Unity3d中导入官方自带的Barracuda包: Installing Barracuda | Barracuda | 1.0.4 (unity3d.com)

下面根据onnx的输出结构来获取结果：

一共有8400个数据，每个数据有35个channel。要取得8400个数据中置信度最高的数据，然后获取该数据其他channel的数据。

using System.Collections.Generic;
using UnityEngine;
using Unity.Barracuda;
using UnityEngine.UI;
using System.Linq;

public class Test : MonoBehaviour
{
    public NNModel modelAsset;
    private Model m_RuntimeModel;
    private IWorker worker;

    public Texture2D inputTex;
    public Transform dotParent;
    public GameObject dot;
    public GameObject box;

    private void Start()
    {
        m_RuntimeModel = ModelLoader.Load(modelAsset);
        worker = WorkerFactory.CreateWorker(WorkerFactory.Type.ComputePrecompiled, m_RuntimeModel);

        Predict();
    }

    public void Predict()
    {
        using Tensor inputTensor = new Tensor(inputTex, channels: 3);

        worker.Execute(inputTensor);

        Tensor outputTensor = worker.PeekOutput();

        //get highest confidence
        var classProbabilities = new List<float>();
        for (var boxIndex = 0; boxIndex < outputTensor.width; boxIndex++)
        {
            float confidence = outputTensor[0, 0, boxIndex, 4];
            classProbabilities.Add(confidence);
        }

        var maxIndex = classProbabilities.Any() ? classProbabilities.IndexOf(classProbabilities.Max()) : 0;
        UnityEngine.Debug.Log("Highest confidence:" + outputTensor[0, 0, maxIndex, 4] + " and its index:" + maxIndex);

        Vector2 boxCenter = new Vector2(outputTensor[0, 0, maxIndex, 0], outputTensor[0, 0, maxIndex, 1]);
        Vector2 boxSize = new Vector2(outputTensor[0, 0, maxIndex, 2], outputTensor[0, 0, maxIndex, 3]);
        CreateBox(boxCenter, boxSize);

        for (int i = 5; i < outputTensor.channels; i += 3)
        {
            Vector2 pos = new Vector2(outputTensor[0, 0, maxIndex, i], outputTensor[0, 0, maxIndex, i + 1]);
            CreateRedDot(i, pos);
        }
    }

    //create red dots on specific position
    void CreateRedDot(int boxIndex, Vector2 pos)
    {
        GameObject newDot = Instantiate(dot).gameObject;
        newDot.transform.SetParent(dotParent);
        newDot.GetComponent<RectTransform>().anchoredPosition = new Vector2(pos.x, -pos.y);
        newDot.name = boxIndex.ToString();
        newDot.SetActive(true);
    }

    void CreateBox(Vector2 pos, Vector2 size)
    {
        box.GetComponent<RectTransform>().anchoredPosition = new Vector2(pos.x, -pos.y);
        box.GetComponent<RectTransform>().sizeDelta = size;
        box.SetActive(true);
    }
}

运行结果：（蓝色是检测手的边框，红点是自定义的点）