File size: 6,758 Bytes
05c9ac2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 |
using System;
using UnityEngine;
using UnityEngine.Rendering;
namespace Unity.MLAgents.Sensors
{
/// <summary>
/// A sensor that wraps a Camera object to generate visual observations for an agent.
/// </summary>
public class CameraSensor : ISensor, IBuiltInSensor, IDisposable
{
Camera m_Camera;
int m_Width;
int m_Height;
bool m_Grayscale;
string m_Name;
private ObservationSpec m_ObservationSpec;
SensorCompressionType m_CompressionType;
Texture2D m_Texture;
/// <summary>
/// The Camera used for rendering the sensor observations.
/// </summary>
public Camera Camera
{
get { return m_Camera; }
set { m_Camera = value; }
}
/// <summary>
/// The compression type used by the sensor.
/// </summary>
public SensorCompressionType CompressionType
{
get { return m_CompressionType; }
set { m_CompressionType = value; }
}
/// <summary>
/// Creates and returns the camera sensor.
/// </summary>
/// <param name="camera">Camera object to capture images from.</param>
/// <param name="width">The width of the generated visual observation.</param>
/// <param name="height">The height of the generated visual observation.</param>
/// <param name="grayscale">Whether to convert the generated image to grayscale or keep color.</param>
/// <param name="name">The name of the camera sensor.</param>
/// <param name="compression">The compression to apply to the generated image.</param>
/// <param name="observationType">The type of observation.</param>
public CameraSensor(
Camera camera, int width, int height, bool grayscale, string name, SensorCompressionType compression, ObservationType observationType = ObservationType.Default)
{
m_Camera = camera;
m_Width = width;
m_Height = height;
m_Grayscale = grayscale;
m_Name = name;
var channels = grayscale ? 1 : 3;
m_ObservationSpec = ObservationSpec.Visual(height, width, channels, observationType);
m_CompressionType = compression;
m_Texture = new Texture2D(width, height, TextureFormat.RGB24, false);
}
/// <summary>
/// Accessor for the name of the sensor.
/// </summary>
/// <returns>Sensor name.</returns>
public string GetName()
{
return m_Name;
}
/// <summary>
/// Returns a description of the observations that will be generated by the sensor.
/// The shape will be h x w x 1 for grayscale and h x w x 3 for color.
/// The dimensions have translational equivariance along width and height,
/// and no property along the channels dimension.
/// </summary>
/// <returns></returns>
public ObservationSpec GetObservationSpec()
{
return m_ObservationSpec;
}
/// <summary>
/// Generates a compressed image. This can be valuable in speeding-up training.
/// </summary>
/// <returns>Compressed image.</returns>
public byte[] GetCompressedObservation()
{
using (TimerStack.Instance.Scoped("CameraSensor.GetCompressedObservation"))
{
ObservationToTexture(m_Camera, m_Texture, m_Width, m_Height);
// TODO support more types here, e.g. JPG
var compressed = m_Texture.EncodeToPNG();
return compressed;
}
}
/// <summary>
/// Writes out the generated, uncompressed image to the provided <see cref="ObservationWriter"/>.
/// </summary>
/// <param name="writer">Where the observation is written to.</param>
/// <returns></returns>
public int Write(ObservationWriter writer)
{
using (TimerStack.Instance.Scoped("CameraSensor.WriteToTensor"))
{
ObservationToTexture(m_Camera, m_Texture, m_Width, m_Height);
var numWritten = writer.WriteTexture(m_Texture, m_Grayscale);
return numWritten;
}
}
/// <inheritdoc/>
public void Update() { }
/// <inheritdoc/>
public void Reset() { }
/// <inheritdoc/>
public CompressionSpec GetCompressionSpec()
{
return new CompressionSpec(m_CompressionType);
}
/// <summary>
/// Renders a Camera instance to a 2D texture at the corresponding resolution.
/// </summary>
/// <param name="obsCamera">Camera.</param>
/// <param name="texture2D">Texture2D to render to.</param>
/// <param name="width">Width of resulting 2D texture.</param>
/// <param name="height">Height of resulting 2D texture.</param>
public static void ObservationToTexture(Camera obsCamera, Texture2D texture2D, int width, int height)
{
if (SystemInfo.graphicsDeviceType == GraphicsDeviceType.Null)
{
Debug.LogError("GraphicsDeviceType is Null. This will likely crash when trying to render.");
}
var oldRec = obsCamera.rect;
obsCamera.rect = new Rect(0f, 0f, 1f, 1f);
var depth = 24;
var format = RenderTextureFormat.Default;
var readWrite = RenderTextureReadWrite.Default;
var tempRt =
RenderTexture.GetTemporary(width, height, depth, format, readWrite);
var prevActiveRt = RenderTexture.active;
var prevCameraRt = obsCamera.targetTexture;
// render to offscreen texture (readonly from CPU side)
RenderTexture.active = tempRt;
obsCamera.targetTexture = tempRt;
obsCamera.Render();
texture2D.ReadPixels(new Rect(0, 0, texture2D.width, texture2D.height), 0, 0);
obsCamera.targetTexture = prevCameraRt;
obsCamera.rect = oldRec;
RenderTexture.active = prevActiveRt;
RenderTexture.ReleaseTemporary(tempRt);
}
/// <inheritdoc/>
public BuiltInSensorType GetBuiltInSensorType()
{
return BuiltInSensorType.CameraSensor;
}
/// <summary>
/// Clean up the owned Texture2D.
/// </summary>
public void Dispose()
{
if (!ReferenceEquals(null, m_Texture))
{
Utilities.DestroyTexture(m_Texture);
m_Texture = null;
}
}
}
}
|