ml-agents/com.unity.ml-agents/Runtime/Sensors/CameraSensor.cs

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;
            }
        }
    }
}