使用 Unity Sentis 和 Compute Shader，2d106det.onnx 进行高效人脸网格标记

前言

前篇：使用 Unity Sentis 和 Compute Shader，det_10g.onnx 进行高效人脸五官定位-CSDN博客

在计算机视觉领域，人脸网格标记是一项重要的任务，用于识别人脸关键点和特征。本文将介绍如何使用 Unity Sentis 和 Compute Shader，结合 2d106det.onnx 模型，实现高效的人脸网格标记。我将提供完整的代码示例。

模型分析

输入值：

模型的输入是1x3x192x192的一张图片；

输出值：

输出值为（1，212）的张量，存储的是106个特征点的位置，可以将它变形为（106，2）方便计算；

代码示例：

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

public class Landmark2D : MonoBehaviour
{
    public ModelAsset modelAsset;
    public Model model; 
    private IWorker worker;
    private GPUComputeBackend gpu;
    
    //处理图片
    public Texture2D t2d;
    public Material landMarkMat;


    private RenderTexture fc1RT;
    public ComputeShader outputCS;
    

    private void Start()
    {

        fc1RT = new RenderTexture(106, 1, 0,RenderTextureFormat.RGFloat);
        model = ModelLoader.Load(modelAsset);  
      
        gpu = new GPUComputeBackend();
        
        var model2 = Functional.Compile(input =>
        {
            var outputs = model.Forward(input);
            var pred = outputs[0].Reshape(new []{106,2});
            pred = pred[.., 0..2] + 1;
            pred = pred[.., 0..2] * (192 / 2);
            return pred;

        }, InputDef.FromModel(model)[0]
        );
        worker = WorkerFactory.CreateWorker(BackendType.GPUCompute, model2);
        Get(t2d);
    }

   
    public RenderTexture rt1;
    public RenderTexture markTexture;
    public void Get(Texture2D source )
    {
        rt1 = RenderTexture.GetTemporary(192,192,0);
        landMarkMat.SetTexture("_MainTex",source);
        Graphics.Blit(source,rt1,landMarkMat,0);   
        
        using (var input = TextureConverter.ToTensor(rt1, 192, 192, 3))
        {
            TensorFloat _255 = new TensorFloat(new TensorShape(1), new float[] {255});
            TensorFloat newInput = TensorFloat.AllocNoData(new TensorShape(1, 3, 192,192)); 
            gpu.Mul(input,_255,newInput);
            worker.Execute(newInput);
        }
   
        
        using var fc1 = worker.PeekOutput("output_0") as TensorFloat;  
        fc1.Reshape(new TensorShape(1,106,1,2));
         
        using  TensorFloat tagetT = TensorFloat.AllocNoData(new TensorShape(1, 2, 1,106)); 
        gpu.Transpose(fc1,tagetT,new int[] {0, 3, 1, 2});
        
        fc1RT = TextureConverter.ToTexture(tagetT, 106, 1, 2);  
        
        
        markTexture = new RenderTexture(192, 192, 0);
        markTexture.enableRandomWrite = true;  // 允许随机写入
        markTexture.Create();  // 创建纹理
        
        outputCS.SetTexture(0,"fcTex",fc1RT);
        outputCS.SetTexture(0,"outTex",markTexture);
        outputCS.Dispatch (0, (106 + 15)/16, 1, 1); 
        
 
        landMarkMat.SetTexture("_MarkTex",markTexture);
        
        /*fc1.CompleteOperationsAndDownload();
        for (int i = 0; i < 212; i++)
        {
            print(fc1[i]);
        }*/
    
    }

    private void OnDestroy()
    {
        worker.Dispose();
        gpu.Dispose();
        RenderTexture.ReleaseTemporary(rt1);
    }
}