NNModel LoadOnnxModel(byte[] rawModel)
{
var converter = new ONNXModelConverter(true);
var onnxModel = converter.Convert(rawModel);
NNModelData assetData = ScriptableObject.CreateInstance <NNModelData>();
using (var memoryStream = new MemoryStream())
using (var writer = new BinaryWriter(memoryStream))
{
ModelWriter.Save(writer, onnxModel);
assetData.Value = memoryStream.ToArray();
}
assetData.name = "Data";
assetData.hideFlags = HideFlags.HideInHierarchy;
var asset = ScriptableObject.CreateInstance <NNModel>();
asset.modelData = assetData;
return(asset);
}
public NNModel GetModelForBehaviorName(string behaviorName)
{
if (m_CachedModels.ContainsKey(behaviorName))
{
return(m_CachedModels[behaviorName]);
}
string assetPath = null;
if (m_BehaviorNameOverrides.ContainsKey(behaviorName))
{
assetPath = m_BehaviorNameOverrides[behaviorName];
}
else if (!string.IsNullOrEmpty(m_BehaviorNameOverrideDirectory))
{
assetPath = Path.Combine(m_BehaviorNameOverrideDirectory, $"{behaviorName}.{m_OverrideExtension}");
}
if (string.IsNullOrEmpty(assetPath))
{
Debug.Log($"No override for BehaviorName {behaviorName}, and no directory set.");
return(null);
}
byte[] rawModel = null;
try
{
rawModel = File.ReadAllBytes(assetPath);
}
catch (IOException)
{
Debug.Log($"Couldn't load file {assetPath} at full path {Path.GetFullPath(assetPath)}", this);
// Cache the null so we don't repeatedly try to load a missing file
m_CachedModels[behaviorName] = null;
return(null);
}
NNModel asset;
var isOnnx = m_OverrideExtension.Equals("onnx");
if (isOnnx)
{
var converter = new ONNXModelConverter(true);
var onnxModel = converter.Convert(rawModel);
NNModelData assetData = ScriptableObject.CreateInstance <NNModelData>();
using (var memoryStream = new MemoryStream())
using (var writer = new BinaryWriter(memoryStream))
{
ModelWriter.Save(writer, onnxModel);
assetData.Value = memoryStream.ToArray();
}
assetData.name = "Data";
assetData.hideFlags = HideFlags.HideInHierarchy;
asset = ScriptableObject.CreateInstance <NNModel>();
asset.modelData = assetData;
}
else
{
// Note - this approach doesn't work for onnx files. Need to replace with
// the equivalent of ONNXModelImporter.OnImportAsset()
asset = ScriptableObject.CreateInstance <NNModel>();
asset.modelData = ScriptableObject.CreateInstance <NNModelData>();
asset.modelData.Value = rawModel;
}
asset.name = "Override - " + Path.GetFileName(assetPath);
m_CachedModels[behaviorName] = asset;
return(asset);
}
