Unity异步把图片数据从显存下载到内存（GPU-＞CPU）

Unity异步把图片数据从显存下载到内存（GPU->CPU）

1.c#核心代码

using System.Collections;
using System.Collections.Generic;
using Unity.Collections;
using UnityEditor.PackageManager.Requests;
using UnityEngine;
using UnityEngine.Rendering;

namespace Tools
{
    /// <summary>
    /// RenderTexture转Bytes
    /// </summary>
    public class TextureOrBytes
    {
        private ComputeShader computeShader;
        private Texture srcRrenderTexture;
        private RenderTexture dstRenderTexture;
        private int kernelHandle;
        private bool requestPending;
        private NativeArray<byte> srcNativeArray;
        private byte[] dstRTData;
        private int width;
        private int height;
        private int memorySize;
        private int threadGroupsX;
        private int threadGroupsY;
        private object lockObject = new object();

        public TextureOrBytes(ComputeShader _computeShader, Texture _srcRrenderTexture)
        {
            computeShader = _computeShader;
            srcRrenderTexture = _srcRrenderTexture; 
            Init();
        }

        public void Update()
        {
            if (!requestPending)
            {
                computeShader.Dispatch(kernelHandle, threadGroupsX, threadGroupsY, 1);

                AsyncGPUReadback.Request(dstRenderTexture, 0, TextureFormat.RGBA32, OnCompleteReadback);
                requestPending = true;
            }
        }

        public void Clear()
        {
            if (dstRenderTexture != null)
            {
                dstRenderTexture.Release();
                dstRenderTexture = null;
            }
            if (srcNativeArray != null)
            {
                srcNativeArray.Dispose();
            }
            
            dstRTData = null;
            lockObject = null;
        }

        private void Init()
        {
            width = srcRrenderTexture.width;
            height = srcRrenderTexture.height;
            memorySize = width * height * 4;
            srcNativeArray = new NativeArray<byte>(memorySize, Allocator.Temp);
            dstRTData = new byte[memorySize];

            //创建临时RT
            dstRenderTexture = new RenderTexture(width, height, 0, RenderTextureFormat.ARGB32);
            dstRenderTexture.enableRandomWrite = true;
            dstRenderTexture.Create();

            kernelHandle = computeShader.FindKernel("CSMain");
            computeShader.SetTexture(kernelHandle, "sourceTex", srcRrenderTexture);
            computeShader.SetTexture(kernelHandle, "destTex", dstRenderTexture);
            threadGroupsX = Mathf.CeilToInt(width / 8.0f);
            threadGroupsY = Mathf.CeilToInt(height / 8.0f);
        }

        private void OnCompleteReadback(AsyncGPUReadbackRequest request)
        {
            if (srcNativeArray == null || dstRTData == null || lockObject == null)
            {
                requestPending = false;
                return;
            }

            if (!request.done)
            {
                Debug.Log("GPU readback hasnt done yet");
                return;
            }

            if (request.hasError)
            {
                Debug.LogError("GPU readback error detected.");
                requestPending = false;
                return;
            }

            srcNativeArray = request.GetData<byte>(); 

            lock (lockObject)
            {
                if (srcNativeArray.Length == dstRTData.Length)
                {
                    srcNativeArray.CopyTo(dstRTData);
                    
                }
            }
            requestPending = false;
        }

        public byte[] GetTexGPUBytesData()
        {
            lock (lockObject) 
            {
                return dstRTData;
            }
        }

        public Texture GetTex()
        {
            return srcRrenderTexture;
        }

        public RenderTexture GetRWTex()
        {
            return dstRenderTexture;
        }


        #region 递归方式..

        private void StartRequestReadbacks()
        {
            computeShader.Dispatch(kernelHandle, threadGroupsX, threadGroupsY, 1);

            AsyncGPUReadback.Request(dstRenderTexture, 0, TextureFormat.RGBA32, OnCompleteReadback_Recursive);
        }

        private void OnCompleteReadback_Recursive(AsyncGPUReadbackRequest request)
        {
            if (srcNativeArray == null || dstRTData == null || lockObject == null)
            {
                requestPending = false;
                return;
            }

            if (!request.done)
            {
                Debug.Log("GPU readback hasnt done yet");
                return;
            }

            if (request.hasError)
            {
                Debug.LogError("GPU readback error detected.");
                requestPending = false;
                return;
            }

            srcNativeArray = request.GetData<byte>();

            //lock (lockObject)
            {
                if (srcNativeArray.Length == dstRTData.Length)
                {
                    srcNativeArray.CopyTo(dstRTData);
                }
            }

            StartRequestReadbacks();
        }

        #endregion


    }

}

2.ComputeShader核心代码

// Each #kernel tells which function to compile; you can have many kernels
#pragma kernel CSMain

// Texture to read from
Texture2D<float4> sourceTex;
// Texture to write to
RWTexture2D<float4> destTex;

[numthreads(8, 8, 1)]
void CSMain(uint3 id : SV_DispatchThreadID)
{
    float4 color = sourceTex[id.xy];
    destTex[id.xy] = color;
}

3.使用

• 在外部实例化TextureOrBytes
• 在update中调用 TextureOrBytes.Update();
• 通过TextureOrBytes.GetTexGPUBytesData()获取CPU数据