使用 Unity Sentis 和 Compute Shader 进行高效人脸识别

Unity sentis就是Unity barracuda的升级版，所以使用sentis人脸识别步骤基本和barracuda一致，详情参考我的《使用 Unity Barracuda 和 Compute Shader 进行高效人脸识别》使用 Unity Barracuda 和 Compute Shader 进行高效人脸识别-CSDN博客

注意：

需要注意的是 version-RFB-320.onnx 模型在sentis框架下的inputs输入值和barracuda有所不同，形状为（1，3，240，320），barracuda里是（1，240，320，3），图片输入的时候转换一下轴顺序即可。

代码：

using System;
using Unity.Mathematics;
using Unity.Sentis;
using UnityEngine;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using UnityEngine.Rendering;
using UnityEngine.UI;
using Unity.Sentis.Layers;
 

public class Ultraface  : MonoBehaviour 
{
    public ModelAsset modelAsset;
    public Model model; 
    private IWorker worker;

    private int inputTexWidth = 320, inputTexHeight = 240;
    private int OutputCount = 4420;
     
    public  struct Detection
    {
        public   float x1, y1, x2, y2;
        public   float score;
        public   float pad1, pad2, pad3; 
        public static int Size = 8 * sizeof(float); 
    }
    

    [SerializeField]private WebCamTexture webTex;


   
    
    public ComputeShader postprocess1;
    private RenderTexture scores;
    private RenderTexture boxes;
    private ComputeBuffer post1;
    
    public ComputeShader postprocess2;
    private ComputeBuffer post2;
    private ComputeBuffer counter;
   
    
    
    
    [SerializeField]  Shader _visualizer = null;
    Material _material;
    ComputeBuffer _drawArgs;
    [SerializeField] private Texture2D _texture;
    
    [SerializeField]  RawImage _previewUI = null;
    [SerializeField] private Texture2D dTexture;
    
    public void SetIndirectDrawCount(ComputeBuffer drawArgs)
        => ComputeBuffer.CopyCount( post2, drawArgs, sizeof(uint));

    private GPUComputeBackend gpu;
    
    
  
    private void Start()
    {
        InitWebCam();
        InitModel();
        InitBuffer();
          gpu = new GPUComputeBackend();
        _material = new Material(_visualizer);
        
        _drawArgs = new ComputeBuffer(4, sizeof(uint),
            ComputeBufferType.IndirectArguments);
        _drawArgs.SetData(new int [] {6, 0, 0, 0});
        _previewUI.texture = webTex;
     
    }

    private void Update()
    {
        RunModel(webTex,0.3f);
    
    }
 
    void OnRenderObject()
    {
        SetIndirectDrawCount(_drawArgs); 
        _material.SetFloat("_Threshold", 0.3f);
        _material.SetTexture("_Texture", _texture);
        _material.SetBuffer("_Detections", post2);
        _material.SetPass(_texture == null ? 0 : 1);
        Graphics.DrawProceduralIndirectNow(MeshTopology.Triangles, _drawArgs, 0);
    }
   


    void InitWebCam()
    {
        webTex = new WebCamTexture(1920, 1080,30);
        webTex.deviceName = WebCamTexture.devices[0].name;
        webTex.Play();
       
    }

    void InitModel()
    {
        model = ModelLoader.Load(modelAsset);
        
        worker = WorkerFactory.CreateWorker(BackendType.GPUCompute, model);
        
         
    }

    void InitBuffer()
    {
        
         scores = new RenderTexture(OutputCount / 20, 20, 0);
          
         boxes = new RenderTexture(OutputCount / 20, 20, 0, RenderTextureFormat.ARGBFloat); 
    
        post1 = new ComputeBuffer(512, Detection.Size);
        post2 = new ComputeBuffer(512, Detection.Size, ComputeBufferType.Append);
        counter = new ComputeBuffer(1, sizeof(uint), ComputeBufferType.Counter);
        
    }

  
      void RunModel(Texture source, float threshold)
    {
        using ( TensorFloat input = TextureConverter.ToTensor(source, 320, 240, 3) )
        {
        
           using  TensorFloat _2 = new TensorFloat(2f);
           using  TensorFloat _1 = new TensorFloat(1f);
           using  TensorFloat input1 = TensorFloat.AllocZeros(new TensorShape(1, 3, 240, 320));
           gpu.Mul(input,_2,input1);
        
     
           using   TensorFloat input2 = TensorFloat.AllocZeros(new TensorShape(1, 3, 240, 320));
           gpu.Sub(input1,_1,input2);
     
           worker.Execute( input2);

        }  
         using var s= worker.PeekOutput("scores") as TensorFloat;
         using var b= worker.PeekOutput("boxes")as TensorFloat;
          
         s.Reshape( new TensorShape(1, 221, 20,2));
         b.Reshape( new TensorShape(1, 221, 20,4));
         
         using  TensorFloat tagetT = TensorFloat.AllocZeros(new TensorShape(1, 2, 20,221)); 
         gpu.Transpose(s,tagetT,new int[] {0, 3, 1, 2}); 
         using   TensorFloat tagetT1 = TensorFloat.AllocZeros(new TensorShape(1, 4, 20,221)); 
         gpu.Transpose(b,tagetT1,new int[] {0, 3, 1, 2}); 
            
         scores = TextureConverter.ToTexture(tagetT, 221, 20, 2);  
         boxes  = TextureConverter.ToTexture(tagetT1 , 221, 20, 4); 
         
    
        
         post2.SetCounterValue(0);
         counter.SetCounterValue(0);
        
        
        
        postprocess1.SetTexture(0, "Scores",  scores);
        postprocess1.SetTexture(0, "Boxes", boxes);
        postprocess1.SetInts("InputSize", boxes.width,boxes.height);
        postprocess1.SetFloat("Threshold", threshold);
        postprocess1.SetBuffer(0, "Output",  post1);
        postprocess1.SetBuffer(0, "OutputCount",  counter);
        postprocess1.Dispatch (0, (boxes.width+7)/16,boxes.height/4,1);
        
      
        
        postprocess2.SetFloat ("Threshold", 0.5f);
        postprocess2.SetBuffer(0, "Input",  post1);
        postprocess2.SetBuffer(0, "InputCount",  counter);
        postprocess2.SetBuffer(0, "Output",  post2);
        postprocess2.Dispatch (0, 1, 1, 1); 
    
    }

  
    private void OnDestroy()
    {
        worker.Dispose(); 
        post1.Dispose();
        post2.Dispose();
        counter.Dispose();
        _drawArgs.Dispose();
        Destroy(webTex);
        Destroy(scores);
        Destroy(boxes);
        gpu.Dispose();
    }
}