您最多选择25个主题 主题必须以中文或者字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

103 行
3.8 KiB

using System;
using System.Linq;
using Unity.Collections;
using Unity.Simulation;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;
namespace UnityEngine.Perception.GroundTruth
{
/// <summary>
/// RenderTextureReader reads a RenderTexture from the GPU whenever Capture is called and passes the data back through a provided callback.
/// </summary>
/// <typeparam name="T">The type of the raw texture data to be provided.</typeparam>
class RenderTextureReader<T> : IDisposable where T : struct
{
RenderTexture m_Source;
Texture2D m_CpuTexture;
/// <summary>
/// Creates a new <see cref="RenderTextureReader{T}"/> for the given <see cref="RenderTexture"/>, <see cref="Camera"/>, and image readback callback
/// </summary>
/// <param name="source">The <see cref="RenderTexture"/> to read from.</param>
public RenderTextureReader(RenderTexture source)
{
m_Source = source;
}
public void Capture(ScriptableRenderContext context, Action<int, NativeArray<T>, RenderTexture> imageReadCallback)
{
if (!GraphicsUtilities.SupportsAsyncReadback())
{
RenderTexture.active = m_Source;
if (m_CpuTexture == null || m_CpuTexture.width != m_Source.width || m_CpuTexture.height != m_Source.height)
{
if (m_CpuTexture != null)
m_CpuTexture.Resize(m_Source.width, m_Source.height);
m_CpuTexture = new Texture2D(m_Source.width, m_Source.height, m_Source.graphicsFormat, TextureCreationFlags.None);
}
m_CpuTexture.ReadPixels(new Rect(
Vector2.zero,
new Vector2(m_Source.width, m_Source.height)),
0, 0);
RenderTexture.active = null;
var data = m_CpuTexture.GetRawTextureData<T>();
imageReadCallback(Time.frameCount, data, m_Source);
}
else
{
var commandBuffer = CommandBufferPool.Get("RenderTextureReader");
var frameCount = Time.frameCount;
commandBuffer.RequestAsyncReadback(m_Source, r => OnGpuReadback(r, frameCount, imageReadCallback));
context.ExecuteCommandBuffer(commandBuffer);
context.Submit();
CommandBufferPool.Release(commandBuffer);
}
}
void OnGpuReadback(AsyncGPUReadbackRequest request, int frameCount,
Action<int, NativeArray<T>, RenderTexture> imageReadCallback)
{
if (request.hasError)
{
Debug.LogError("Error reading segmentation image from GPU");
}
else if (request.done && imageReadCallback != null)
{
imageReadCallback(frameCount, request.GetData<T>(), m_Source);
}
}
/// <summary>
/// Synchronously wait for all image requests to complete.
/// </summary>
public void WaitForAllImages()
{
AsyncGPUReadback.WaitAllRequests();
}
/// <summary>
/// Shut down the reader, waiting for all requests to return.
/// </summary>
public void Dispose() => Dispose(true);
/// <summary>
/// Shut down the reader, optionally waiting for all requests to return.
/// </summary>
/// <param name="waitForAllImages">Whether this should block on waiting for asynchronous readbacks</param>
public void Dispose(bool waitForAllImages)
{
if (waitForAllImages)
WaitForAllImages();
if (m_CpuTexture != null)
{
Object.Destroy(m_CpuTexture);
m_CpuTexture = null;
}
}
}
}