Unity3d 基于Barracuda推理库和YOLO算法实现对象检测功能

前言

近年来，随着AI技术的发展，在游戏引擎中实现和运行机器学习模型的需求也逐渐显现。Unity3d引擎官方推出深度学习推理框架–Barracuda ，旨在帮助开发者在Unity3d中轻松地实现和运行机器学习模型，它的主要功能是支持在 Unity 中加载和推理训练好的深度学习模型，尤其适用于需要人工智能（AI）或机器学习（ML）推理的游戏或应用。

YOLO（You Only Look Once）是一种用于目标检测的深度学习模型，它是由Joseph Redmon等人在2015年提出的。YOLO的核心思想是将目标检测问题转化为一个回归问题，在单一的神经网络中同时预测图像中的多个目标位置和类别标签。它通过将目标检测转化为回归问题，极大地提高了检测速度，并且在精度上也能达到非常好的水平。随着版本的更新和技术的不断进步，YOLO逐渐成为了计算机视觉领域中最重要和最广泛应用的模型之一，特别适用于实时处理、嵌入式设备和大规模部署。

本文依托上述两个技术，在Unity3d中实现YOLO的目标检测功能，基于Barracuda(2.0.0)的跨平台性，将实现包含移动端（目前测试了安卓）的目标检测功能，能检测出日常物体桌、椅、人、狗、羊、马等对象。

理论上本工程可以在Windows/Mac/iPhone/Android/Magic Leap/Switch/PS4/Xbox等系统和平台正常工作，目前仅测试了Windows和Android平台，相比Windows平台的流畅，Android手机上运行有明显的掉帧和卡顿，具体可以对比效果图。

官方给出支持的平台说明：
CPU 推理：支持所有 Unity 平台。
GPU 推理：支持所有 Unity 平台，但以下平台：
OpenGL ESon ：使用 Vulkan/Metal。Android/iOS
OpenGL Core上：使用 Metal。Mac
WebGL：使用 CPU 推理。

关注并私信 U3D目标检测免费获取应用包(底部公众号)。

效果

手机端效果：
在这里插入图片描述

在这里插入图片描述

PC端效果：
在这里插入图片描述

在这里插入图片描述

实现

Barracuda 是一个简单、对开发人员友好的API，只需编写少量代码即可开始使用Barracuda：

var model = ModelLoader.Load(filename);
var engine = WorkerFactory.CreateWorker(model, WorkerFactory.Device.GPU);
var input = new Tensor(1, 1, 1, 10);
var output = engine.Execute(input).PeekOutput();

Barracuda 神经网络导入管道基于ONNX（Open Neural Network Exchange）格式的模型，允许您从各种外部框架（包括Pytorch、TensorFlow和Keras）引入神经网络模型。

关于模型

Barracuda目前仅支持推理，所以模型靠TensorFlow/Pytorch/Keras训练、导入，而且必须先将其转换为 ONNX，然后将其加载到 Unity中。ONNX（Open Neural Network Exchange）是一种用于ML 模型的开放格式。它允许您在各种ML框架和工具之间轻松交换模型。
Pytorch将模型导出到ONNX很容易

# network
net = ...

# Input to the model
x = torch.randn(1, 3, 256, 256)

# Export the model
torch.onnx.export(net,                       # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  "example.onnx",            # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=9,           # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['X'],       # the model's input names
                  output_names = ['Y']       # the model's output names
                  )

我这里准备的是很简单的模型，如下图：
在这里插入图片描述

确保ONNX模型的输入尺寸、通道顺序（NCHW）与Barracuda兼容。
因为要兼顾移动端效果，所以模型检测识别对象较少，以防止在移动设备上的推理慢。

UI搭建

运行时的UI相对简单，两个button用于打开摄像头和打开视频功能，一个Slider用于控制标记框的显示阈值，就是检测的可信度从0-1（0%-100%）的范围；一个rawImage组件用于显示检测的画面：
在这里插入图片描述

其次是标记框的UI，由一个图片和Text构成：
在这里插入图片描述

编码

加载模型

var model = ModelLoader.Load(resources.model);

其中模型类型是NNModel。

创建推理引擎（Worker）并执行模型：

_worker = model.CreateWorker();
using (var t = new Tensor(_config.InputShape, _buffers.preprocess))
  _worker.Execute(t);

提取神经网络输出：

_worker.CopyOutput("Identity", _buffers.feature1);
_worker.CopyOutput("Identity_1", _buffers.feature2);

将网络的两个输出复制到缓冲区。

第一阶段后处理，检测数据：

var post1 = _resources.postprocess1;
post1.SetInt("ClassCount", _config.ClassCount);
post1.SetFloat("Threshold", threshold);
post1.SetBuffer(0, "Output", _buffers.post1);
post1.SetBuffer(0, "OutputCount", _buffers.counter);

var width1 = _config.FeatureMap1Width;
post1.SetTexture(0, "Input", _buffers.feature1);
post1.SetInt("InputSize", width1);
post1.SetFloats("Anchors", _config.AnchorArray1);
post1.DispatchThreads(0, width1, width1, 1);

var width2 = _config.FeatureMap2Width;
post1.SetTexture(0, "Input", _buffers.feature2);
post1.SetInt("InputSize", width2);
post1.SetFloats("Anchors", _config.AnchorArray2);
post1.DispatchThreads(0, width2, width2, 1);

聚合检测结果，使用两个特征图进行目标检测，执行目标定位（Bounding Box）预测。

第二阶段后处理，重叠移除：

var post2 = _resources.postprocess2;
post2.SetFloat("Threshold", 0.5f);
post2.SetBuffer(0, "Input", _buffers.post1);
post2.SetBuffer(0, "InputCount", _buffers.counter);
post2.SetBuffer(0, "Output", _buffers.post2);
post2.Dispatch(0, 1, 1, 1);

移除重叠的边界框。

上面的复杂处理是借Compute Shader的Preprocess、Postprocess1和postprocess2来实现的，Compute Shader 是一种图形编程中的着色器类型，专门用于执行计算任务，而不直接参与渲染。详细内容如下。

Common.hlsl：

// Compile-time constants
#define MAX_DETECTION 512
#define ANCHOR_COUNT 3

// Detection data structure - The layout of this structure must be matched
// with the one defined in Detection.cs.
struct Detection
{
    float x, y, w, h;
    uint classIndex;
    float score;
};

// Misc math functions

float CalculateIOU(in Detection d1, in Detection d2)
{
    float x0 = max(d1.x - d1.w / 2, d2.x - d2.w / 2);
    float x1 = min(d1.x + d1.w / 2, d2.x + d2.w / 2);
    float y0 = max(d1.y - d1.h / 2, d2.y - d2.h / 2);
    float y1 = min(d1.y + d1.h / 2, d2.y + d2.h / 2);

    float area0 = d1.w * d1.h;
    float area1 = d2.w * d2.h;
    float areaInner = max(0, x1 - x0) * max(0, y1 - y0);

    return areaInner / (area0 + area1 - areaInner);
}

float Sigmoid(float x)
{
    return 1 / (1 + exp(-x));
}

#endif

Postprocess1.compute：

#pragma kernel Postprocess1

#include "Common.hlsl"

// Input
Texture2D<float> Input;
uint InputSize;
float2 Anchors[ANCHOR_COUNT];
uint ClassCount;
float Threshold;

// Output buffer
RWStructuredBuffer<Detection> Output;
RWStructuredBuffer<uint> OutputCount; // Only used as a counter

[numthreads(8, 8, 1)]
void Postprocess1(uint2 id : SV_DispatchThreadID)
{
    if (!all(id < InputSize)) return;

    // Input reference point:
    // We have to read the input tensor in reversed order.
    uint ref_y = (InputSize - 1 - id.y) * InputSize + (InputSize - 1 - id.x);

    for (uint aidx = 0; aidx < ANCHOR_COUNT; aidx++)
    {
        uint ref_x = aidx * (5 + ClassCount);

        // Bounding box / confidence
        float x = Input[uint2(ref_x + 0, ref_y)];
        float y = Input[uint2(ref_x + 1, ref_y)];
        float w = Input[uint2(ref_x + 2, ref_y)];
        float h = Input[uint2(ref_x + 3, ref_y)];
        float c = Input[uint2(ref_x + 4, ref_y)];

        // ArgMax[SoftMax[classes]]
        uint maxClass = 0;
        float maxScore = exp(Input[uint2(ref_x + 5, ref_y)]);
        float scoreSum = maxScore;
        for (uint cidx = 1; cidx < ClassCount; cidx++)
        {
            float score = exp(Input[uint2(ref_x + 5 + cidx, ref_y)]);
            if (score > maxScore)
            {
                maxClass = cidx;
                maxScore = score;
            }
            scoreSum += score;
        }

        // Output structure
        Detection data;
        data.x = (id.x + Sigmoid(x)) / InputSize;
        data.y = (id.y + Sigmoid(y)) / InputSize;
        data.w = exp(w) * Anchors[aidx].x;
        data.h = exp(h) * Anchors[aidx].y;
        data.classIndex = maxClass;
        data.score = Sigmoid(c) * maxScore / scoreSum;

        // Thresholding
        if (data.score > Threshold)
        {
            // Detected: Count and output
            uint count = OutputCount.IncrementCounter();
            if (count < MAX_DETECTION) Output[count] = data;
        }
    }
}

Postprocess2.compute：

#pragma kernel Postprocess2

#include "Common.hlsl"

// Input
StructuredBuffer<Detection> Input;
RWStructuredBuffer<uint> InputCount; // Only used as a counter
float Threshold;

// Output
AppendStructuredBuffer<Detection> Output;

// Local arrays for data cache
groupshared Detection _entries[MAX_DETECTION];
groupshared bool _flags[MAX_DETECTION];

[numthreads(1, 1, 1)]
void Postprocess2(uint3 id : SV_DispatchThreadID)
{
    // Initialize data cache arrays
    uint entry_count = min(MAX_DETECTION, InputCount.IncrementCounter());
    if (entry_count == 0) return;

    for (uint i = 0; i < entry_count; i++)
    {
        _entries[i] = Input[i];
        _flags[i] = true;
    }

    for (i = 0; i < entry_count - 1; i++)
    {
        if (!_flags[i]) continue;

        for (uint j = i + 1; j < entry_count; j++)
        {
            if (!_flags[j]) continue;
            if (CalculateIOU(_entries[i], _entries[j]) < Threshold)
                continue;
            if (_entries[i].score < _entries[j].score)
            {
                _flags[i] = false;
                break;
            }
            else
                _flags[j] = false;
        }
    }
    for (i = 0; i < entry_count; i++)
        if (_flags[i]) Output.Append(_entries[i]);
}

Postprocess.compute：

#pragma kernel Preprocess

sampler2D Image;
RWStructuredBuffer<float> Tensor;
uint Size;

[numthreads(8, 8, 1)]
void Preprocess(uint2 id : SV_DispatchThreadID)
{
    // UV (vertically flipped)
    float2 uv = float2(0.5 + id.x, Size - 0.5 - id.y) / Size;

    // UV gradients
    float2 duv_dx = float2(1.0 / Size, 0);
    float2 duv_dy = float2(0, -1.0 / Size);

    // Texture sample
    float3 rgb = tex2Dgrad(Image, uv, duv_dx, duv_dy).rgb;

    // Tensor element output
    uint offs = (id.y * Size + id.x) * 3;
    Tensor[offs + 0] = rgb.r;
    Tensor[offs + 1] = rgb.g;
    Tensor[offs + 2] = rgb.b;
}

通过以上的处理，最后输出了一个目标检测的对象结果数组，主要包含如下数据：

public readonly struct Detection
{
    public readonly float x, y, w, h;
    public readonly uint classIndex;
    public readonly float score;
}

通过遍历这个数组，并将结果标记框和对象名称等信息显示出来：

public void SetAttributes(in Detection d)
{
    var rect = _parent.rect;

    var x = d.x * rect.width;
    var y = (1 - d.y) * rect.height;
    var w = d.w * rect.width;
    var h = d.h * rect.height;

    _xform.anchoredPosition = new Vector2(x, y);
    _xform.SetSizeWithCurrentAnchors(RectTransform.Axis.Horizontal, w);
    _xform.SetSizeWithCurrentAnchors(RectTransform.Axis.Vertical, h);

    var name = _labels[(int)d.classIndex];
    _label.text = $"{name} {(int)(d.score * 100)}%";

    var hue = d.classIndex * 0.073f % 1.0f;
    var color = Color.HSVToRGB(hue, 1, 1);

    _panel.color = color;

    transform.localScale = Vector3.one;
}