浏览代码

Coco exporter files

/coco_export
Steve Borkman 3 年前
当前提交
3b3d0a71
共有 12 个文件被更改,包括 799 次插入0 次删除
  1. 8
      com.unity.perception/Runtime/GroundTruth/Exporters.meta
  2. 8
      com.unity.perception/Runtime/GroundTruth/Exporters/Coco.meta
  3. 61
      com.unity.perception/Runtime/GroundTruth/Exporters/Coco/AnnotationHandler.cs
  4. 11
      com.unity.perception/Runtime/GroundTruth/Exporters/Coco/AnnotationHandler.cs.meta
  5. 511
      com.unity.perception/Runtime/GroundTruth/Exporters/Coco/CocoExporter.cs
  6. 11
      com.unity.perception/Runtime/GroundTruth/Exporters/Coco/CocoExporter.cs.meta
  7. 149
      com.unity.perception/Runtime/GroundTruth/Exporters/Coco/CocoTypes.cs
  8. 11
      com.unity.perception/Runtime/GroundTruth/Exporters/Coco/CocoTypes.cs.meta
  9. 18
      com.unity.perception/Runtime/GroundTruth/Exporters/IDatasetReporter.cs
  10. 11
      com.unity.perception/Runtime/GroundTruth/Exporters/IDatasetReporter.cs.meta

8
com.unity.perception/Runtime/GroundTruth/Exporters.meta


fileFormatVersion: 2
guid: 5353392e887128948bc94b5e0dd9ff73
folderAsset: yes
DefaultImporter:
externalObjects: {}
userData:
assetBundleName:
assetBundleVariant:

8
com.unity.perception/Runtime/GroundTruth/Exporters/Coco.meta


fileFormatVersion: 2
guid: b543148039c075741a0d54a25d56f573
folderAsset: yes
DefaultImporter:
externalObjects: {}
userData:
assetBundleName:
assetBundleVariant:

61
com.unity.perception/Runtime/GroundTruth/Exporters/Coco/AnnotationHandler.cs


using System;
namespace UnityEngine.Perception.GroundTruth.Exporters.Coco
{
public static class AnnotationHandler
{
public static object HandleAnnotation(AsyncAnnotation asyncAnnotation, object annotation)
{
object converted = annotation switch
{
BoundingBox2DLabeler.BoundingBoxValue bbox => CocoTypes.ObjectDetectionAnnotation.FromBoundingBoxValue(bbox),
KeypointLabeler.KeypointEntry keypoint => CocoTypes.KeypointAnnotation.FromKeypointValue(keypoint),
_ => null
};
return converted;
}
public static CocoTypes.KeypointCategory ToKeypointCategory(KeypointLabeler.KeypointJson keypointJson)
{
var keypoints = new string[keypointJson.key_points.Length];
var skeleton = new int[keypointJson.skeleton.Length * 2];
foreach (var kp in keypointJson.key_points)
{
keypoints[kp.index] = kp.label;
}
var i = 0;
foreach (var bone in keypointJson.skeleton)
{
skeleton[i++] = bone.joint1;
skeleton[i++] = bone.joint2;
}
return new CocoTypes.KeypointCategory
{
keypoints = keypoints,
skeleton = skeleton
};
}
public static object HandleCameraCapture(int id, int width, int height, string filename)
{
var image = new CocoTypes.Image()
{
id = id,
width = width,
height = height,
file_name = filename,
license = 0,
flickr_url = "",
coco_url = "",
date_captured = DateTime.Today.ToString("D")
};
return image;
}
}
}

11
com.unity.perception/Runtime/GroundTruth/Exporters/Coco/AnnotationHandler.cs.meta


fileFormatVersion: 2
guid: 0f587bc3e200ea542bd7725e7f64978c
MonoImporter:
externalObjects: {}
serializedVersion: 2
defaultReferences: []
executionOrder: 0
icon: {instanceID: 0}
userData:
assetBundleName:
assetBundleVariant:

511
com.unity.perception/Runtime/GroundTruth/Exporters/Coco/CocoExporter.cs


using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using System.Threading.Tasks;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;
namespace UnityEngine.Perception.GroundTruth.Exporters.Coco
{
public class CocoExporter : IDatasetReporter
{
bool m_ReportingObjectDetection;
bool m_ReportingKeypoints;
bool m_Initialized;
string m_DirectoryName = string.Empty;
string m_RgbCaptureFilename;
StreamWriter m_RgbCaptureStream;
string m_ObjectDetectionFilename;
StreamWriter m_ObjectDetectionStream;
Task m_ObjectDetectionWritingTask;
string m_KeypointFilename;
StreamWriter m_KeypointDetectionStream;
Task m_KeypointDetectionWritingTask;
CocoTypes.ObjectDetectionCategories m_ObjectDetectionCategories;
string m_ObjectDetectionCategoryFilename;
StreamWriter m_ObjectDetectionCategoryStream;
Task m_ObjectDetectionCategoryWritingTask;
CocoTypes.KeypointCategories m_KeypointCategories;
string m_KeypointCategoryFilename;
StreamWriter m_KeypointCategoryStream;
Task m_KeypointCategoryWritingTask;
Guid m_SessionGuid;
public void OnSimulationBegin(string directoryName)
{
m_DirectoryName = directoryName;
m_DataCaptured = false;
}
async Task AwaitAllWrites()
{
WriteOutCategories();
if (m_ObjectDetectionWritingTask != null)
{
await m_ObjectDetectionWritingTask;
await m_ObjectDetectionStream.WriteAsync("]");
}
if (m_KeypointDetectionWritingTask != null)
{
await m_KeypointDetectionWritingTask;
await m_KeypointDetectionStream.WriteAsync("]");
}
if (m_RgbCaptureWritingTask != null)
{
await m_RgbCaptureWritingTask;
await m_RgbCaptureStream.WriteAsync("]");
}
if (m_ObjectDetectionCategoryWritingTask != null)
{
await m_ObjectDetectionCategoryWritingTask;
}
if (m_KeypointCategoryWritingTask != null)
{
await m_KeypointCategoryWritingTask;
}
m_RgbCaptureStream?.Close();
m_ObjectDetectionStream?.Close();
m_ObjectDetectionCategoryStream?.Close();
m_KeypointDetectionStream?.Close();
m_KeypointCategoryStream?.Close();
}
public async void OnSimulationEnd()
{
if (!m_DataCaptured) return;
await AwaitAllWrites();
if (m_ReportingObjectDetection)
{
await WriteObjectDetectionFile(true);
}
if (m_ReportingKeypoints)
{
await WriteKeypointFile(true);
}
File.Delete(m_RgbCaptureFilename);
m_Initialized = false;
}
void InitializeCaptureFiles()
{
if (m_Initialized) return;
if (!Directory.Exists(m_DirectoryName))
Directory.CreateDirectory(m_DirectoryName);
m_SessionGuid = Guid.NewGuid();
var prefix = m_DirectoryName + Path.DirectorySeparatorChar + m_SessionGuid;
m_RgbCaptureFilename = prefix + "_coco_captures.json";
m_RgbCaptureStream = File.CreateText(m_RgbCaptureFilename);
m_ObjectDetectionFilename = prefix + "_coco_box_annotations.json";
m_ObjectDetectionStream = File.CreateText(m_ObjectDetectionFilename);
m_KeypointFilename = prefix + "_coco_keypoint_annotations.json";
m_KeypointDetectionStream = File.CreateText(m_KeypointFilename);
m_ObjectDetectionCategoryFilename = prefix + "_coco_obj_detection_categories.json";
m_KeypointCategoryFilename = prefix + "_coco_keypoint_categories.json";
m_Initialized = true;
}
static void AggregateFile(string filename, StringBuilder aggregated, bool skipFirstCharacter = false)
{
using var sr = new StreamReader(filename);
var length = (int)sr.BaseStream.Length;
var start = 0;
if (length == 0) return;
var buffer = new char[length];
sr.Read(buffer, start, length);
if (skipFirstCharacter)
{
length--;
start++;
}
aggregated.Append(buffer, start, length);
}
async Task WriteObjectDetectionFile(bool prettyPrint = false)
{
var stringBuilder = new StringBuilder();
stringBuilder.Append("{");
// Create the json header
CreateHeaderInfo(stringBuilder);
stringBuilder.Append(",");
CreateLicenseInfo(stringBuilder);
// Read in the contents of the captures
stringBuilder.Append(",\"images\":");
AggregateFile(m_RgbCaptureFilename, stringBuilder);
// Read in the contents of the object detection
stringBuilder.Append(",\"annotations\":");
AggregateFile(m_ObjectDetectionFilename, stringBuilder);
stringBuilder.Append(",");
// Read in the contents of the object detection categories
AggregateFile(m_ObjectDetectionCategoryFilename, stringBuilder, true);
var json = stringBuilder.ToString();
if (prettyPrint)
{
json = JToken.Parse(json).ToString(Formatting.Indented);
}
// Write out the files
var filename = m_DirectoryName + Path.DirectorySeparatorChar + "coco_object_detection_annotations.json";
var cocoStream = File.CreateText(filename);
await cocoStream.WriteAsync(json);
cocoStream.Close();
File.Delete(m_ObjectDetectionFilename);
File.Delete(m_ObjectDetectionCategoryFilename);
}
async Task WriteKeypointFile(bool prettyPrint = false)
{
var stringBuilder = new StringBuilder();
stringBuilder.Append("{");
// Create the json header
CreateHeaderInfo(stringBuilder);
stringBuilder.Append(",");
CreateLicenseInfo(stringBuilder);
// Read in the contents of the captures
stringBuilder.Append(",\"images\":");
AggregateFile(m_RgbCaptureFilename, stringBuilder);
// Read in the contents of the object detection
stringBuilder.Append(",\"annotations\":");
AggregateFile(m_KeypointFilename, stringBuilder);
stringBuilder.Append(",");
// Read in the contents of the object detection categories
AggregateFile(m_KeypointCategoryFilename, stringBuilder, true);
var json = stringBuilder.ToString();
if (prettyPrint)
{
json = JToken.Parse(json).ToString(Formatting.Indented);
}
// Write out the files
var filename = m_DirectoryName + Path.DirectorySeparatorChar + "coco_keypoint_annotations.json";
var cocoStream = File.CreateText(filename);
await cocoStream.WriteAsync(json);
cocoStream.Close();
File.Delete(m_KeypointFilename);
File.Delete(m_KeypointCategoryFilename);
}
bool m_DataCaptured;
void WriteOutCategories()
{
// 3 cases
// 1) Just have object detection
// 2) Just have keypoint detection
// 3) Have both, which includes writing out object detection in object detection coco output
// and merging the object detection & keypoints and writing out that data in the keypoint
// coco output data
if (m_ReportingObjectDetection)
{
m_ObjectDetectionCategoryStream = File.CreateText(m_ObjectDetectionCategoryFilename);
var json = JsonUtility.ToJson(m_ObjectDetectionCategories);
m_ObjectDetectionCategoryWritingTask = m_ObjectDetectionCategoryStream.WriteAsync(json);
}
if (m_ReportingKeypoints)
{
// TODO - revisit this, but right now we are going to add the keypoint info to all of the categories,
// we can get away with this because we only support one set of keypoint definitions per simulation
// currently.
if (m_KeypointCategories.categories.Length < 1) return;
m_KeypointCategoryStream = File.CreateText(m_KeypointCategoryFilename);
var merged = m_KeypointCategories;
if (m_ReportingObjectDetection)
{
merged = new CocoTypes.KeypointCategories
{
categories = new CocoTypes.KeypointCategory[m_ObjectDetectionCategories.categories.Count]
};
var i = 0;
foreach (var odc in m_ObjectDetectionCategories.categories)
{
merged.categories[i++] = MergeCategories(odc, m_KeypointCategories.categories[0]);
}
}
var json = JsonUtility.ToJson(merged);
m_KeypointCategoryWritingTask = m_KeypointCategoryStream.WriteAsync(json);
}
}
public void OnAnnotationRegistered<TSpec>(Guid annotationId, TSpec[] values)
{
InitializeCaptureFiles();
if (annotationId.ToString() == BoundingBox2DLabeler.annotationId)
{
m_ReportingObjectDetection = true;
m_ObjectDetectionCategories = new CocoTypes.ObjectDetectionCategories
{
categories = new List<CocoTypes.ObjectDetectionCategory>()
};
foreach (var value in values)
{
if (value is IdLabelConfig.LabelEntrySpec spec)
{
var rec = new CocoTypes.ObjectDetectionCategory
{
id = spec.label_id,
name = spec.label_name,
supercategory = spec.label_name
};
m_ObjectDetectionCategories.categories.Add(rec);
}
}
}
if (annotationId.ToString() == KeypointLabeler.annotationId)
{
m_ReportingKeypoints = true;
if (values[0] is KeypointLabeler.KeypointJson keypointJson)
{
m_KeypointCategories = new CocoTypes.KeypointCategories
{
categories = new []
{
AnnotationHandler.ToKeypointCategory(keypointJson)
}
};
}
}
}
static void CreateHeaderInfo(StringBuilder stringBuilder)
{
stringBuilder.Append("\"info\":");
var dateTime = DateTime.Today;
var info = new CocoTypes.Info
{
year = int.Parse(dateTime.ToString("yyyy")),
version = "0.0.1",
description = "temp output of coco data",
contributor = "Tyler Durden",
url = "https://ytmnd.com",
date_created = DateTime.Today.ToString("D")
};
stringBuilder.Append(JsonUtility.ToJson(info));
}
static void CreateLicenseInfo(StringBuilder stringBuilder)
{
var licenses = new CocoTypes.Licenses()
{
licenses = new[]
{
new CocoTypes.License
{
id = 0,
name = "Unity License",
url = "https://unity3d.com"
}
}
};
var tmpJson = JsonUtility.ToJson(licenses);
// Remove the start and end '{' from the licenses json
stringBuilder.Append(tmpJson.Substring(1, tmpJson.Length - 2));
}
bool m_FirstBoxAnnotation = true;
bool m_FirstKeypointAnnotation = true;
bool m_FirstCapture = true;
Task m_RgbCaptureWritingTask;
static CocoTypes.KeypointCategory MergeCategories(CocoTypes.ObjectDetectionCategory od, CocoTypes.KeypointCategory kp)
{
return new CocoTypes.KeypointCategory
{
id = od.id,
name = od.name,
supercategory = od.supercategory,
keypoints = kp.keypoints,
skeleton = kp.skeleton
};
}
public async Task ProcessPendingCaptures(List<SimulationState.PendingCapture> pendingCaptures, SimulationState simState)
{
var boxJson = string.Empty;
var keypointJson = string.Empty;
foreach (var cap in pendingCaptures)
{
var boxes = new Dictionary<int, CocoTypes.ObjectDetectionAnnotation>();
foreach (var annotation in cap.Annotations)
{
var tmp = ProcessBoundingBoxAnnotations(annotation.Item2.RawValues);
foreach (var box in tmp.Values)
{
boxes[box.id] = box;
if (m_FirstBoxAnnotation)
{
boxJson = "[";
m_FirstBoxAnnotation = false;
}
else
boxJson += ",";
boxJson += JsonUtility.ToJson(box);
}
}
foreach (var annotation in cap.Annotations)
{
var keypoints = ProcessKeypointAnnotations(annotation.Item2.RawValues, boxes);
foreach (var kp in keypoints.Values)
{
if (m_FirstKeypointAnnotation)
{
keypointJson = "[";
m_FirstKeypointAnnotation = false;
}
else
keypointJson += ",";
keypointJson += JsonUtility.ToJson(kp);
}
}
}
if (m_ObjectDetectionWritingTask != null)
await m_ObjectDetectionWritingTask;
if (boxJson != string.Empty)
m_ObjectDetectionWritingTask = m_ObjectDetectionStream.WriteAsync(boxJson);
if (m_KeypointDetectionWritingTask != null)
await m_KeypointDetectionWritingTask;
if (keypointJson != string.Empty)
m_KeypointDetectionWritingTask = m_KeypointDetectionStream.WriteAsync(keypointJson);
}
static Dictionary<int, CocoTypes.ObjectDetectionAnnotation> ProcessBoundingBoxAnnotations(IEnumerable<object> annotations)
{
var map = new Dictionary<int, CocoTypes.ObjectDetectionAnnotation>();
foreach (var annotation in annotations)
{
if (annotation is BoundingBox2DLabeler.BoundingBoxValue bbox)
{
var coco = CocoTypes.ObjectDetectionAnnotation.FromBoundingBoxValue(bbox);
map[coco.id] = coco;
}
}
return map;
}
static Dictionary<int, CocoTypes.KeypointAnnotation> ProcessKeypointAnnotations(IEnumerable<object> annotations, Dictionary<int, CocoTypes.ObjectDetectionAnnotation> boundingBoxes)
{
var map = new Dictionary<int, CocoTypes.KeypointAnnotation>();
foreach (var annotation in annotations)
{
if (annotation is KeypointLabeler.KeypointEntry keypoint)
{
var coco = CocoTypes.KeypointAnnotation.FromKeypointValue(keypoint);
if (boundingBoxes.ContainsKey(coco.id))
{
coco.CopyObjectDetectionData(boundingBoxes[coco.id]);
}
map[coco.id] = coco;
}
}
return map;
}
public async Task OnCaptureReported(int frame, int width, int height, string filename)
{
InitializeCaptureFiles();
var json = string.Empty;
var converted = AnnotationHandler.HandleCameraCapture(frame, width, height, filename);
if (m_FirstCapture)
{
json = "[";
m_FirstCapture = false;
}
else
json += ",";
json += JsonUtility.ToJson(converted);
if (m_RgbCaptureWritingTask != null)
await m_RgbCaptureWritingTask;
if (json != string.Empty)
m_RgbCaptureWritingTask = m_RgbCaptureStream.WriteAsync(json);
m_DataCaptured = true;
}
}
}

11
com.unity.perception/Runtime/GroundTruth/Exporters/Coco/CocoExporter.cs.meta


fileFormatVersion: 2
guid: d25254cc4c0ad734cb76497c70a40088
MonoImporter:
externalObjects: {}
serializedVersion: 2
defaultReferences: []
executionOrder: 0
icon: {instanceID: 0}
userData:
assetBundleName:
assetBundleVariant:

149
com.unity.perception/Runtime/GroundTruth/Exporters/Coco/CocoTypes.cs


using System;
using System.Collections.Generic;
namespace UnityEngine.Perception.GroundTruth.Exporters.Coco
{
public static class CocoTypes
{
public class CocoData
{
public Info info;
public Image[] images;
public ObjectDetectionAnnotation[] annotations;
public ObjectDetectionCategory[] categories;
public License[] licenses;
}
public class Info
{
public int year;
public string version;
public string description;
public string contributor;
public string url;
public string date_created;
}
[Serializable]
public class License
{
public int id;
public string name;
public string url;
}
[Serializable]
public class Licenses
{
public License[] licenses;
}
public class Image
{
public int id;
public int width;
public int height;
public string file_name;
public int license;
public string flickr_url;
public string coco_url;
public string date_captured;
}
[Serializable]
public class ObjectDetectionAnnotation
{
public int id;
public int image_id;
public int category_id;
public float[] segmentation;
public float area;
public float[] bbox;
public int iscrowd;
public static ObjectDetectionAnnotation FromBoundingBoxValue(BoundingBox2DLabeler.BoundingBoxValue bbox)
{
return new ObjectDetectionAnnotation
{
id = (int)bbox.instance_id,
image_id = bbox.frame,
category_id = bbox.label_id,
segmentation = new float[]{},
area = bbox.width * bbox.height,
bbox = new []{bbox.x, bbox.y, bbox.width, bbox.height},
iscrowd = 0
};
}
}
[Serializable]
public class ObjectDetectionCategory
{
public int id;
public string name;
public string supercategory;
}
[Serializable]
public class ObjectDetectionCategories
{
public List<ObjectDetectionCategory> categories;
}
public class KeypointAnnotation : ObjectDetectionAnnotation
{
public float[] keypoints;
public int num_keypoints;
public void CopyObjectDetectionData(ObjectDetectionAnnotation objDetection)
{
if (objDetection.id == this.id)
{
image_id = objDetection.image_id;
area = objDetection.area;
bbox = objDetection.bbox;
iscrowd = objDetection.iscrowd;
}
}
public static KeypointAnnotation FromKeypointValue(KeypointLabeler.KeypointEntry keypoint)
{
var outKeypoint = new KeypointAnnotation()
{
id = (int)keypoint.instance_id,
image_id = -1,
category_id = keypoint.label_id,
segmentation = new float[]{},
area = 0,
bbox = new []{0f},
iscrowd = 0,
num_keypoints = keypoint.keypoints.Length,
keypoints = new float[keypoint.keypoints.Length * 3]
};
var i = 0;
foreach (var k in keypoint.keypoints)
{
outKeypoint.keypoints[i++] = k.x;
outKeypoint.keypoints[i++] = k.y;
outKeypoint.keypoints[i++] = k.state;
}
return outKeypoint;
}
}
[Serializable]
public class KeypointCategory : ObjectDetectionCategory
{
public string[] keypoints;
public int[] skeleton;
}
[Serializable]
public class KeypointCategories
{
public KeypointCategory[] categories;
}
}
}

11
com.unity.perception/Runtime/GroundTruth/Exporters/Coco/CocoTypes.cs.meta


fileFormatVersion: 2
guid: a715e9572e6cc314098dcdf401bb9273
MonoImporter:
externalObjects: {}
serializedVersion: 2
defaultReferences: []
executionOrder: 0
icon: {instanceID: 0}
userData:
assetBundleName:
assetBundleVariant:

18
com.unity.perception/Runtime/GroundTruth/Exporters/IDatasetReporter.cs


using System;
using System.Collections.Generic;
using System.Threading.Tasks;
namespace UnityEngine.Perception.GroundTruth.Exporters
{
public interface IDatasetReporter
{
public void OnSimulationBegin(string directoryName);
public void OnSimulationEnd();
public void OnAnnotationRegistered<TSpec>(Guid annotationId, TSpec[] values);
public Task ProcessPendingCaptures(List<SimulationState.PendingCapture> pendingCaptures, SimulationState simState);
public Task OnCaptureReported(int frame, int width, int height, string filename);
}
}

11
com.unity.perception/Runtime/GroundTruth/Exporters/IDatasetReporter.cs.meta


fileFormatVersion: 2
guid: 474afd2e32730994ba1d1494b735a911
MonoImporter:
externalObjects: {}
serializedVersion: 2
defaultReferences: []
executionOrder: 0
icon: {instanceID: 0}
userData:
assetBundleName:
assetBundleVariant:
正在加载...
取消
保存