using System.Collections.Generic;
using Unity.Sentis;
using UnityEngine;

/*
 *  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
{
    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:
    const int imageHeight = 224;
    const int imageWidth = 224;

    const BackendType backend = BackendType.GPUCompute;
    private Model model;
    private IWorker engine;
    private string[] labels;

    static Ops ops;
    ITensorAllocator allocator;

    void Start()
    {
        //These are used for tensor operations
        allocator = new TensorCachingAllocator();
        ops = WorkerFactory.CreateOps(backend, allocator);

        //Parse neural net labels
        labels = labelsAsset.text.Split('\n');

        //Load model
        model = ModelLoader.Load(Application.streamingAssetsPath + "/" + modelName);

        //Setup the engine to run the model
        engine = WorkerFactory.CreateWorker(backend, model);

        //Execute inference
        ExecuteML();
    }

    public void ExecuteML()
    {
        //Preprocess image for input
        using var rawinput = TextureConverter.ToTensor(inputImage, 224, 224, 3);
        using var input = Normalise(rawinput);
        
        //Execute neural net
        engine.Execute(input);

        //Read output tensor
        var output = engine.PeekOutput() as TensorFloat;
        var argmax = ops.ArgMax(output, 1, false);
        argmax.MakeReadable();

        //Select the best output class and print the results
        var res = argmax[0];
        var label = labels[res];

        output.MakeReadable();
        var accuracy = output[res];

        //The result is output to the console window
        int percent = Mathf.FloorToInt(accuracy * 100f + 0.5f);
        Debug.Log($"{label} {percent}﹪");

        //Clean memory
        Resources.UnloadUnusedAssets();
    }

    //This scales and shifts the RGB values for input into the model
    TensorFloat Normalise(TensorFloat image)
    {
        using var M = new TensorFloat(new TensorShape(1, 3, 1, 1), new float[]
        {
           1/0.229f, 1/0.224f, 1/0.225f
        });
        using var P = new TensorFloat(new TensorShape(1, 3, 1, 1), new float[]
        {
            0.485f, 0.456f, 0.406f
        });
        using var image2 = ops.Sub(image, P);
        return ops.Mul(image2, M);
    }
    
    private void OnDestroy()
    {
        engine?.Dispose();
        ops?.Dispose();
        allocator?.Dispose();
    }
}