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 runx,                         # 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 fileopset_version=9,           # the ONNX version to export the model todo_constant_folding=True,  # whether to execute constant folding for optimizationinput_names = ['X'],       # the model's input namesoutput_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 functionsfloat 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 / confidencefloat 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 structureDetection 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;// Thresholdingif (data.score > Threshold){// Detected: Count and outputuint 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 arraysuint 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 Preprocesssampler2D 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 gradientsfloat2 duv_dx = float2(1.0 / Size, 0);float2 duv_dy = float2(0, -1.0 / Size);// Texture samplefloat3 rgb = tex2Dgrad(Image, uv, duv_dx, duv_dy).rgb;// Tensor element outputuint 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;
}