Image Classification
unity-sentis
ONNX
sentis-mobilenet-v2 / RunMobileNet.cs
UnityPaul's picture
Upload 4 files
05abce4 verified
raw
history blame
3.41 kB
using System.Collections.Generic;
using Unity.Sentis;
using UnityEngine;
using System.IO;
using FF = Unity.Sentis.Functional;
/*
* MovileNetV2 Inference Script
* ============================
*
* Place this script on the Main Camera
*
* Drag an image to the inputImage field
*
* When run the prediction of what the image is will output to the console window.
* You can modify the script to make it do something more interesting.
*
*/
public class RunMobileNet : MonoBehaviour
{
//draw the sentis file here:
public ModelAsset modelAsset;
const string modelName = "mobilenet_v2.sentis";
//The image to classify here:
public Texture2D inputImage;
//Link class_desc.txt here:
public TextAsset labelsAsset;
//All images are resized to these values to go into the model
const int imageHeight = 224;
const int imageWidth = 224;
const BackendType backend = BackendType.GPUCompute;
private IWorker engine;
private string[] labels;
//Used to normalise the input RGB values
TensorFloat mulRGB = new TensorFloat(new TensorShape(1, 3, 1, 1), new float[] { 1 / 0.229f, 1 / 0.224f, 1 / 0.225f });
TensorFloat shiftRGB = new TensorFloat(new TensorShape(1, 3, 1, 1), new float[] { 0.485f, 0.456f, 0.406f });
void Start()
{
//Parse neural net labels
labels = labelsAsset.text.Split('\n');
//Load model from file or asset
//var model = ModelLoader.Load(Path.Join(Application.streamingAssetsPath, modelName));
var model = ModelLoader.Load(modelAsset);
//We modify the model to normalise the input RGB values and select the highest prediction
//probability and item number
var model2 = FF.Compile(
input =>
{
var probability = model.Forward(NormaliseRGB(input))[0];
return (FF.ReduceMax(probability, 1), FF.ArgMax(probability, 1));
},
model.inputs[0]
);
//Setup the engine to run the model
engine = WorkerFactory.CreateWorker(backend, model2);
//Execute inference
ExecuteML();
}
public void ExecuteML()
{
//Preprocess image for input
using var input = TextureConverter.ToTensor(inputImage, imageWidth, imageHeight, 3);
//Execute neural net
engine.Execute(input);
//Read output tensor
var probability = engine.PeekOutput("output_0") as TensorFloat;
var item = engine.PeekOutput("output_1") as TensorInt;
item.CompleteOperationsAndDownload();
probability.CompleteOperationsAndDownload();
//Select the best output class and print the results
var ID = item[0];
var accuracy = probability[0];
//The result is output to the console window
int percent = Mathf.FloorToInt(accuracy * 100f + 0.5f);
Debug.Log($"Prediction: {labels[ID]} {percent}﹪");
//Clean memory
Resources.UnloadUnusedAssets();
}
//This scales and shifts the RGB values for input into the model
FunctionalTensor NormaliseRGB(FunctionalTensor image)
{
return (image - FunctionalTensor.FromTensor(shiftRGB)) * FunctionalTensor.FromTensor(mulRGB);
}
private void OnDestroy()
{
mulRGB?.Dispose();
shiftRGB?.Dispose();
engine?.Dispose();
}
}