using System; using System.Collections.Generic; using System.Globalization; using System.IO; using System.Linq; using System.Text; using Newtonsoft.Json; using Newtonsoft.Json.Linq; using Unity.Simulation; // ReSharper disable NotAccessedField.Local // ReSharper disable CoVariantArrayConversion namespace UnityEngine.Perception.GroundTruth { partial class SimulationState { const Formatting k_Formatting = Formatting.Indented; public void WriteMetadata() { var metadata = new JObject(); metadata["version"] = DatasetCapture.SchemaVersion; metadata["per_file_size"] = k_MinPendingMetricsBeforeWrite; foreach (var md in m_Metadata) { metadata[md.Key] = JToken.FromObject(md.Value); } //metadata["dataset_size"] = 500; //metadata["per_file_size"] = 150; //metadata["image_width"] = 1024; //metadata["image_height"] = 768; WriteJObjectToFile(metadata, "metadata.json"); } /// /// Writes sensors.json, egos.json, metric_definitions.json, and annotation_definitions.json /// public void WriteReferences() { var egoReference = new JObject(); egoReference["version"] = DatasetCapture.SchemaVersion; egoReference["egos"] = new JArray(m_Egos.Select(e => { var egoObj = new JObject(); egoObj["id"] = e.Id.ToString(); if (e.Description != null) egoObj["description"] = e.Description; return egoObj; }).ToArray()); WriteJObjectToFile(egoReference, "egos.json"); var sensorReferenceDoc = new JObject(); sensorReferenceDoc["version"] = DatasetCapture.SchemaVersion; sensorReferenceDoc["sensors"] = new JArray(m_Sensors.Select(kvp => { var sensorReference = new JObject(); sensorReference["id"] = kvp.Key.Id.ToString(); sensorReference["ego_id"] = kvp.Value.egoHandle.Id.ToString(); sensorReference["modality"] = kvp.Value.modality; if (kvp.Value.description != null) sensorReference["description"] = kvp.Value.description; return sensorReference; }).ToArray()); WriteJObjectToFile(sensorReferenceDoc, "sensors.json"); if (m_AdditionalInfoTypeData.Count > 0) { var annotationDefinitionsJArray = new JArray(); var metricDefinitionsJArray = new JArray(); foreach (var typeInfo in m_AdditionalInfoTypeData) { var typeJObject = new JObject(); typeJObject.Add("id", new JValue(typeInfo.id.ToString())); typeJObject.Add("name", new JValue(typeInfo.name)); if (typeInfo.description != null) typeJObject.Add("description", new JValue(typeInfo.description)); if (typeInfo.format != null) typeJObject.Add("format", new JValue(typeInfo.format)); if (typeInfo.specValues != null) { var specValues = new JArray(); foreach (var value in typeInfo.specValues) { specValues.Add(DatasetJsonUtility.ToJToken(value)); } typeJObject.Add("spec", specValues); } switch (typeInfo.additionalInfoKind) { case AdditionalInfoKind.Annotation: annotationDefinitionsJArray.Add(typeJObject); break; case AdditionalInfoKind.Metric: metricDefinitionsJArray.Add(typeJObject); break; default: throw new NotSupportedException("Unsupported info kind"); } } if (annotationDefinitionsJArray.Count > 0) { var annotationDefinitionsJObject = new JObject(); annotationDefinitionsJObject.Add("version", DatasetCapture.SchemaVersion); annotationDefinitionsJObject.Add("annotation_definitions", annotationDefinitionsJArray); WriteJObjectToFile(annotationDefinitionsJObject, "annotation_definitions.json"); } if (metricDefinitionsJArray.Count > 0) { var metricDefinitionsJObject = new JObject(); metricDefinitionsJObject.Add("version", DatasetCapture.SchemaVersion); metricDefinitionsJObject.Add("metric_definitions", metricDefinitionsJArray); WriteJObjectToFile(metricDefinitionsJObject, "metric_definitions.json"); } } Debug.Log($"Dataset written to {Path.GetDirectoryName(OutputDirectory)}"); } void WriteJObjectToFile(JObject jObject, string filename) { m_JsonToStringSampler.Begin(); var stringWriter = new StringWriter(new StringBuilder(256), CultureInfo.InvariantCulture); using (var jsonTextWriter = new JsonTextWriter(stringWriter)) { jsonTextWriter.Formatting = k_Formatting; jObject.WriteTo(jsonTextWriter); } var contents = stringWriter.ToString(); m_JsonToStringSampler.End(); m_WriteToDiskSampler.Begin(); var path = Path.Combine(OutputDirectory, filename); File.WriteAllText(path, contents); Manager.Instance.ConsumerFileProduced(path); m_WriteToDiskSampler.End(); } void WritePendingCaptures(bool flush = false, bool writeCapturesFromThisFrame = false) { if (!flush && m_PendingCaptures.Count < k_MinPendingCapturesBeforeWrite) return; m_SerializeCapturesSampler.Begin(); var pendingCapturesToWrite = new List(m_PendingCaptures.Count); var frameCountNow = Time.frameCount; for (var i = 0; i < m_PendingCaptures.Count; i++) { var pendingCapture = m_PendingCaptures[i]; if ((writeCapturesFromThisFrame || pendingCapture.FrameCount < frameCountNow) && pendingCapture.Annotations.All(a => a.Item2.IsAssigned)) { pendingCapturesToWrite.Add(pendingCapture); m_PendingCaptures.RemoveAt(i); i--; //decrement i because we removed an element } } if (pendingCapturesToWrite.Count == 0) { m_SerializeCapturesSampler.End(); return; } void Write(List pendingCaptures, SimulationState simulationState, int captureFileIndex) { simulationState.m_SerializeCapturesAsyncSampler.Begin(); //lazily allocate for fast zero-write frames var capturesJArray = new JArray(); foreach (var pendingCapture in pendingCaptures) capturesJArray.Add(JObjectFromPendingCapture(pendingCapture)); var capturesJObject = new JObject(); capturesJObject.Add("version", DatasetCapture.SchemaVersion); capturesJObject.Add("captures", capturesJArray); simulationState.WriteJObjectToFile(capturesJObject, $"captures_{captureFileIndex:000}.json"); simulationState.m_SerializeCapturesAsyncSampler.End(); } if (flush) { Write(pendingCapturesToWrite, this, m_CaptureFileIndex); } else { var req = Manager.Instance.CreateRequest>(); req.data = new WritePendingCaptureRequestData() { CaptureFileIndex = m_CaptureFileIndex, PendingCaptures = pendingCapturesToWrite, SimulationState = this }; req.Enqueue(r => { Write(r.data.PendingCaptures, r.data.SimulationState, r.data.CaptureFileIndex); return AsyncRequest.Result.Completed; }); req.Execute(AsyncRequest.ExecutionContext.JobSystem); } m_SerializeCapturesSampler.End(); m_CaptureFileIndex++; } struct WritePendingMetricRequestData { public List PendingMetrics; public int MetricFileIndex; } void WritePendingMetrics(bool flush = false) { if (!flush && m_PendingMetrics.Count < k_MinPendingMetricsBeforeWrite) return; var pendingMetricsToWrite = new List(m_PendingMetrics.Count); m_SerializeMetricsSampler.Begin(); for (var i = 0; i < m_PendingMetrics.Count; i++) { var metric = m_PendingMetrics[i]; if (metric.IsAssigned) { pendingMetricsToWrite.Add(metric); m_PendingMetrics.RemoveAt(i); i--; //decrement i because we removed an element } } if (pendingMetricsToWrite.Count == 0) { m_SerializeMetricsSampler.End(); return; } void Write(List pendingMetrics, int metricsFileIndex) { m_SerializeMetricsAsyncSampler.Begin(); var jArray = new JArray(); foreach (var pendingMetric in pendingMetrics) jArray.Add(JObjectFromPendingMetric(pendingMetric)); var metricsJObject = new JObject(); metricsJObject.Add("version", DatasetCapture.SchemaVersion); metricsJObject.Add("metrics", jArray); WriteJObjectToFile(metricsJObject, $"metrics_{metricsFileIndex:000}.json"); m_SerializeMetricsAsyncSampler.End(); } if (flush) { Write(pendingMetricsToWrite, m_MetricsFileIndex); } else { var req = Manager.Instance.CreateRequest>(); req.data = new WritePendingMetricRequestData() { MetricFileIndex = m_MetricsFileIndex, PendingMetrics = pendingMetricsToWrite }; req.Enqueue(r => { Write(r.data.PendingMetrics, r.data.MetricFileIndex); return AsyncRequest.Result.Completed; }); req.Execute(); } m_MetricsFileIndex++; m_SerializeMetricsSampler.End(); } static JObject JObjectFromPendingMetric(PendingMetric metric) { var jObject = new JObject(); jObject["capture_id"] = metric.CaptureId == Guid.Empty ? new JRaw("null") : new JValue(metric.CaptureId.ToString()); jObject["annotation_id"] = metric.Annotation.IsNil ? new JRaw("null") : new JValue(metric.Annotation.Id.ToString()); jObject["sequence_id"] = metric.SequenceId.ToString(); jObject["step"] = metric.Step; jObject["metric_definition"] = metric.MetricDefinition.Id.ToString(); jObject["values"] = metric.Values; return jObject; } /// /// Creates the json representation of the given PendingCapture. Static because this should not depend on any SimulationState members, /// which may have changed since the capture was reported. /// static JToken JObjectFromPendingCapture(PendingCapture pendingCapture) { var sensorJObject = new JObject();//new SensorCaptureJson sensorJObject["sensor_id"] = pendingCapture.SensorHandle.Id.ToString(); sensorJObject["ego_id"] = pendingCapture.SensorData.egoHandle.Id.ToString(); sensorJObject["modality"] = pendingCapture.SensorData.modality; sensorJObject["translation"] = DatasetJsonUtility.ToJToken(pendingCapture.SensorSpatialData.SensorPose.position); sensorJObject["rotation"] = DatasetJsonUtility.ToJToken(pendingCapture.SensorSpatialData.SensorPose.rotation); if (pendingCapture.AdditionalSensorValues != null) { foreach (var(name, value) in pendingCapture.AdditionalSensorValues) sensorJObject.Add(name, DatasetJsonUtility.ToJToken(value)); } var egoCaptureJson = new JObject(); egoCaptureJson["ego_id"] = pendingCapture.SensorData.egoHandle.Id.ToString(); egoCaptureJson["translation"] = DatasetJsonUtility.ToJToken(pendingCapture.SensorSpatialData.EgoPose.position); egoCaptureJson["rotation"] = DatasetJsonUtility.ToJToken(pendingCapture.SensorSpatialData.EgoPose.rotation); egoCaptureJson["velocity"] = pendingCapture.SensorSpatialData.EgoVelocity.HasValue ? DatasetJsonUtility.ToJToken(pendingCapture.SensorSpatialData.EgoVelocity.Value) : null; egoCaptureJson["acceleration"] = pendingCapture.SensorSpatialData.EgoAcceleration.HasValue ? DatasetJsonUtility.ToJToken(pendingCapture.SensorSpatialData.EgoAcceleration.Value) : null; var capture = new JObject(); capture["id"] = pendingCapture.Id.ToString(); capture["sequence_id"] = pendingCapture.SequenceId.ToString(); capture["step"] = pendingCapture.Step; capture["timestamp"] = pendingCapture.Timestamp; capture["sensor"] = sensorJObject; capture["ego"] = egoCaptureJson; capture["filename"] = pendingCapture.Path; capture["format"] = GetFormatFromFilename(pendingCapture.Path); if (pendingCapture.Annotations.Any()) capture["annotations"] = new JArray(pendingCapture.Annotations.Select(JObjectFromAnnotation).ToArray()); return capture; } static JObject JObjectFromAnnotation((Annotation, AnnotationData) annotationInfo) { var annotationJObject = new JObject(); annotationJObject["id"] = annotationInfo.Item1.Id.ToString(); annotationJObject["annotation_definition"] = annotationInfo.Item2.AnnotationDefinition.Id.ToString(); if (annotationInfo.Item2.Path != null) annotationJObject["filename"] = annotationInfo.Item2.Path; if (annotationInfo.Item2.ValuesJson != null) annotationJObject["values"] = annotationInfo.Item2.ValuesJson; return annotationJObject; } struct WritePendingCaptureRequestData { public List PendingCaptures; public int CaptureFileIndex; public SimulationState SimulationState; } } }