Skip to content

Latest commit

 

History

History
366 lines (302 loc) · 15.1 KB

File metadata and controls

366 lines (302 loc) · 15.1 KB

Custom Endpoints

The Perception package from version 0.11.0-preview.1 onwards supports the concept of custom endpoints. Previously all data generated by Perception was automatically written to disk using our proprietary output format. However, with custom endpoints, you can use our intermediary representation to write all data in the output format of your choice!

We demonstrate this with a hypothetical endpoint we call the FlatEndpoint which will write all data into one single unique folder per simulation run. To do so, we create a new class and implement the IConsumerEndpoint. This allows our custom endpoint to be selected in the Perception Settings UI and be used as an endpoint for the simulation. We further implement the IFileSystemEndpoint interface to hook into how Perception saves where the last created dataset was located and enable features like the Show Latest Dataset button to work properly.

Example Flat Endpoint

Let's define how our files will be structured. Firstly, we want our FlatEndpoint to dump all generated data into small separate JSON and image files into a single unique folder. Secondly, each new dataset will have a unique folder with an increasing suffix. For example, for 3 different runs, we would put each runs files in {basePath}/FLAT_25, {basePath}/FLAT_26, and {basePath}/FLAT_27. Lastly, the file structure for each would look similar to:

- {BASE_PATH}/FLAT_{num}
  - frame-0.json
  - frame-1.json
  - ...
  - annotation-definition-bbox2d.json
  - annotation-definition-obj-count.json
  - ...
  - camera.sem-seg_1.0.json

Implementation in Code

Note on IMessageProducers and IMessageBuilders

Before we dive into the code, there are two important concepts we should go over that'll make it easier to understand the consumer endpoint logic. IMessageProducer's and IMessageBuilder's in conjunction dictate how a piece of information will be serialized. IMessageProducer's create a generic representation (called a message) of some struct or class while IMessageBuilder's are responsible for building, combining, and optionally serializing the message. Let's take a deeper look at these two interfaces below and then see how they're utilized as part of a full code example.

IMessageProducers
Any data structure (i.e. class/struct) can implement the IMessageProducer interface to specify how the encapsulated data can be represented in a generic way. Consider the following class for example:

class Point: IMessageProducer {
    public int x;
    public int y;
    
    void ToMessage(IMessageBuilder builder) {
        builder.AddInt("x", this.x);
        builder.AddInt("y", this.y);
    }   
}

In this case, we are informing any IMessageBuilder's that want to represent the Point class that it should be represented as two integer fields with the values specified. This allows us to pair it with different MessageBuilders to customize how this data is structured and finally serialized.

IMessageBuilder
A class that implements the IMessageBuilder interface can then define for example how the AddInt field is internally represented and eventually serialized. The JsonMessageBuilder included with Perception might write serialize it as:

{
  "x": 2,
  "y": 4
}

However, you may instead have a custom YamlMessageBuilder which instead serializes the same message as:

point:
  x: 2
  y: 4

Here is a full code example of a minimal version of the FlatEndpoint along with brief comments over important pieces of functionality.

using System;
using System.Collections.Generic;
using System.IO;
using Newtonsoft.Json.Linq;
using UnityEngine;
using UnityEngine.Perception.GroundTruth;
using UnityEngine.Perception.GroundTruth.Consumers;
using UnityEngine.Perception.GroundTruth.DataModel;
using UnityEngine.Perception.Settings;

namespace MyCustomNamespace
{
    /// <summary>
    /// Example endpoint which outputs all information in small files to one single directory.
    /// We demonstrate two methods of serializing data: (1) Custom serialization; (2) IMessageProducers
    /// </summary>
    /// <remarks>For more complex examples, checkout SoloEndpoint</remarks>
    [Serializable]
    public class FlatEndpoint : IConsumerEndpoint, IFileSystemEndpoint
    {
        public string prefix = "FLAT";

        DateTimeOffset m_SimulationStartTime;

        /// Helper function to create a new JsonMessageBuilder
        JsonMessageBuilder GetNewJsonMessageBuilder() => new JsonMessageBuilder();

        #region IFileSystemEndpoint

        /// <summary>
        /// Allows the user to set the base path from the Perception Settings UI.
        /// </summary>
        public string basePath
        {
            get => PerceptionSettings.GetOutputBasePath();
            set => PerceptionSettings.SetOutputBasePath(value);
        }

        string m_CachedCurrentPath = string.Empty;
        /// <summary>
        /// The root directory to use for all files that we output.
        /// </summary>
        public string currentPath
        {
            get
            {
                // Check if we already reserved the output path for this simulation run
                if (!string.IsNullOrWhiteSpace(m_CachedCurrentPath))
                    return m_CachedCurrentPath;

                // A small piece of logic to get the next available directory name
                //   get: {basePath}/FLAT_0
                //   if above already exists, then get: {basePath}/FLAT_1
                //   ... and so on

                var availableSuffix = 0;
                m_CachedCurrentPath = string.Empty;
                do
                {
                    m_CachedCurrentPath = Path.Combine(basePath, $"{prefix}_{availableSuffix}");
                    availableSuffix++;
                }
                while (Directory.Exists(m_CachedCurrentPath));

                // actually create the directory we decided upon above
                Directory.CreateDirectory(m_CachedCurrentPath);

                return m_CachedCurrentPath;
            }
        }

        /// <summary>
        /// The path used when "Reset to Default" is used in the Perception Settings UI.
        /// </summary>
        public string defaultPath => Path.Combine(Application.persistentDataPath);
        #endregion

        #region IConsumerEndpoint

        public string description => "Example endpoint that puts all the files in one single directory";

        /// <summary>
        /// Validate the configuration of your endpoint before the simulation runs.
        /// </summary>
        public bool IsValid(out string errorMessage)
        {
            // Check if the prefix supplied by the user is empty or whitespace
            if (string.IsNullOrWhiteSpace(prefix))
            {
                errorMessage = "Prefix must not be empty.";
                return false;
            }

            errorMessage = $"Directory {basePath} does not exist. Please create the directory.";
            // To create {basePath}/FLAT_{xyz}, we need to ensure that the {basePath} directory exists.
            // If it doesn't, the Perception Settings UI will show the above error message.
            return Directory.Exists(basePath);
        }

        public void SimulationStarted(SimulationMetadata metadata)
        {
            // record when the simulation started so we can use this to calculate
            // duration of the simulation on the SimulationEnded function
            m_SimulationStartTime = DateTimeOffset.Now;
        }

        public void SensorRegistered(SensorDefinition sensor)
        {
            //// Using Method 1 (Custom Serialization)
            // 1. Create a new JsonMessageBuilder class
            var builder = GetNewJsonMessageBuilder();

            // 2. Add all relevant fields to the builder as unique key-value pairs
            builder.AddString("model_type", sensor.modelType);
            builder.AddString("capture_mode", sensor.captureTriggerMode.ToString());
            builder.AddString("id", sensor.id);
            builder.AddInt("frames_between_captures", sensor.framesBetweenCaptures);
            builder.AddFloat("first_capture_frame", sensor.firstCaptureFrame);

            // Invariant: The builder now contains a representation of the sensor class
            // 3. We can use the ToJson function in the builder to write that representation to JSON
            PathUtils.WriteAndReportJsonFile(
                Path.Combine(currentPath, $"sensor-{sensor.id}.json"),
                builder.ToJson()
            );
        }

        public void AnnotationRegistered(AnnotationDefinition annotationDefinition)
        {
            //// Using Method 2 (IMessageProducer Serialization)

            // 1. Create a new JsonMessageBuilder class
            var builder = GetNewJsonMessageBuilder();

            // 2. Allow the annotation definition to convert itself to a message and
            // add it to our builder
            annotationDefinition.ToMessage(builder);

            // Invariant: The builder now contains a representation of the annotation definition class
            // 3. We can use the ToJson function in the builder to write that representation to JSON
            PathUtils.WriteAndReportJsonFile(
                Path.Combine(currentPath, $"annotation-definition-{annotationDefinition.id}.json"),
                builder.ToJson()
            );
        }

        public void MetricRegistered(MetricDefinition metricDefinition)
        {
            // Using Method 2 (IMessageProducer Serialization)

            // Similar to SensorDefinition, MetricDefinition also inherits from IMessageProducer
            // so it can tell the builder how it should be serialized.
            var builder = GetNewJsonMessageBuilder();
            metricDefinition.ToMessage(builder);
            PathUtils.WriteAndReportJsonFile(
                Path.Combine(currentPath, $"annotation-definition-{metricDefinition.id}.json"),
                builder.ToJson()
            );
        }

        public void FrameGenerated(Frame frame)
        {
            // Using Method 2 (IMessageProducer Serialization)

            // By default, the JsonMessageBuilder class does not know how to process image files referenced in the
            // Frame class. So we need to make a new FlatFrameMessageBuilder that inherits from JsonMessageBuilder
            // and specify how to handle image files. We can conveniently use the ToMessage function of the Frame
            // class and pass in our new FlatFrameMessageBuilder class.
            var builder = new FlatFrameMessageBuilder(this, frame);
            frame.ToMessage(builder);

            PathUtils.WriteAndReportJsonFile(
                Path.Combine(currentPath, $"frame-{frame.id}.json"),
                builder.ToJson()
            );
        }

        public void SimulationCompleted(SimulationMetadata metadata)
        {
            //// Using Method 2 (IMessageProducer Serialization)

            // 1. Create a new JsonMessageBuilder class
            var metadataBuilder = GetNewJsonMessageBuilder();

            // 2. Add metadata as a message into the metadataBuilder
            metadata.ToMessage(metadataBuilder);

            // 3. Write the metadata parameter to {currentPath}/metadata.json
            PathUtils.WriteAndReportJsonFile(
                Path.Combine(currentPath, "metadata.json"),
                metadataBuilder.ToJson()
            );

            //// Using Method 1 (Custom Serialization)

            // 1. Create a new JsonMessageBuilder class
            var completeBuilder = GetNewJsonMessageBuilder();
            var simulationEndTime = DateTimeOffset.Now;
            var simulationDuration = simulationEndTime - m_SimulationStartTime;

            // 2. Add all relevant key-value pairs
            completeBuilder.AddLong("start_timestamp", m_SimulationStartTime.ToUnixTimeMilliseconds());
            completeBuilder.AddLong("end_timestamp", m_SimulationStartTime.ToUnixTimeMilliseconds());
            completeBuilder.AddDouble("duration_seconds", simulationDuration.TotalSeconds);

            // 3. Convert data to json and write to file
            PathUtils.WriteAndReportJsonFile(
                Path.Combine(currentPath, "simulation-complete.json"),
                completeBuilder.ToJson()
            );
        }

        /// <summary>
        /// Placeholder for crash resumption logic.
        /// </summary>
        /// <remarks>Unsupported for FlatEndpoint</remarks>
        public (string, int) ResumeSimulationFromCrash(int maxFrameCount)
        {
            Debug.LogError("Crash resumption not supported for FlatEndpoint output.");
            return (string.Empty, 0);
        }

        public object Clone()
        {
            return new FlatEndpoint();
        }

        #endregion
    }

    /// <summary>
    /// A MessageBuilder that extends JsonMessageBuilder to add support for serializing images and tensors.
    /// </summary>
    class FlatFrameMessageBuilder : JsonMessageBuilder
    {
        Frame m_Frame;
        FlatEndpoint m_Endpoint;

        public FlatFrameMessageBuilder(FlatEndpoint endpoint, Frame frame)
        {
            m_Endpoint = endpoint;
            m_Frame = frame;
        }

        /// <summary>
        /// Write out the byte array as an image and append sequence and step number to the key to construct
        /// the final file-name.
        /// </summary>
        public override void AddEncodedImage(string key, string extension, byte[] value)
        {
            if (value.Length > 0)
            {
                var filename = $"{key}_{m_Frame.sequence}-{m_Frame.step}.{extension.ToLower()}";
                // write out the file
                PathUtils.WriteAndReportImageFile(
                    Path.Combine(m_Endpoint.currentPath, filename),
                    value
                );
            }
        }

        /// <summary>
        /// A nested message adds the output of an IMessageBuilder to a specific key.
        /// </summary>
        public override IMessageBuilder AddNestedMessage(string key)
        {
            var nested = new FlatFrameMessageBuilder(m_Endpoint, m_Frame);
            if (nestedValue.ContainsKey(key))
            {
                Debug.LogWarning($"Report data with key [{key}] will be overridden by new values");
            }

            nestedValue[key] = nested;
            return nested;
        }

        /// <summary>
        /// Adds the output of an IMessageBuilder as an element of an array
        /// identified by the key <see cref="arraykey"/>.
        /// </summary>
        public override IMessageBuilder AddNestedMessageToVector(string arraykey)
        {
            if (!nestedArrays.TryGetValue(arraykey, out var nestedList))
            {
                nestedList = new List<JsonMessageBuilder>();
                nestedArrays[arraykey] = nestedList;
            }
            var nested = new FlatFrameMessageBuilder(m_Endpoint, m_Frame);
            nestedList.Add(nested);
            return nested;
        }

        // A tensor is a multi-dimensional array
        public override void AddTensor(string key, Tensor tensor)
        {
            // By default, write the tensor as a flattened array
            currentJToken[key] = new JArray(tensor.buffer);
        }
    }
}