HawkeyeVision/Hawkeye.VisionBuilder.Workflow/Operations/AI/RawModelAIOperation.cs

using Hawkeye.VisionBuilder.Workflow.Datatypes;
using OpenCvSharp;
using System.Runtime.InteropServices;
using Hawkeye.VisionBuilder.Workflow.Operations.Attributes;
using Hawkeye.VisionBuilder.Workflow.DataTransfer;

namespace Hawkeye.VisionBuilder.Workflow.Operations.AI;

[Category("AI")]
public class RawModelAIOperation: BaseOperation
{


    private string _modelName;

    public RawModelAIOperation()
    {
        
        
        CanHaveProcessingError = false;
    }

    private CSharpDataTransferHTTP _cSharpDataTransferMQRPC;
    private int _imageWidth;
    private int _imageHeight;



    public string FilterClasses { get; set; } = "1";

    private bool _initialized = false;
    private FilePath _modelFilePath=new FilePath();


    public FilePath ModelFilePath
    {
        get => _modelFilePath;
        set
        {
            _modelFilePath = value;
            _initialized = false;
        }
    }

    public bool IsCategorical { get; set; }

    protected override void InterpretInternal(Context context)
    {
        _cSharpDataTransferMQRPC = PythonModelProxyHTTP.GetInterface();
        if (!CheckImageExists(context)) return;
        if (!CheckColorful(context)) return;

        if (!_initialized)
        {
            if (!File.Exists(ModelFilePath?.Path))
            {
                this.SetError($"Model file \"{ModelFilePath?.Path}\" not found");
                return;
            }

            _modelName = Path.GetFileNameWithoutExtension(ModelFilePath.Path);

            // is relative path
            if (ModelFilePath.Path.StartsWith("."))
            {
                var dir = Directory.GetCurrentDirectory();
                var absPath = Path.GetFullPath(Path.Combine(dir, ModelFilePath.Path));
                ModelFilePath.Path = absPath;
            }
            _cSharpDataTransferMQRPC.LoadModel(ModelFilePath.Path, _modelName);
            _initialized = true;
        }


        _cSharpDataTransferMQRPC.ActivateModel(_modelName);


        int[] sizeEncoded=[];
        int[] outputSize = [];
        sizeEncoded = _cSharpDataTransferMQRPC.GetAcceptSize();
        outputSize = _cSharpDataTransferMQRPC.GetOutputSize();

        
        var currentImage = context.ActiveImage;

        var rightColor = currentImage.ImageData.CvtColor(ColorConversionCodes.BGR2RGB);
        
        
        var size = sizeEncoded.Select(Convert.ToInt32).ToArray();
        _imageWidth = size[0];
        _imageHeight = size[1];

        var resized = rightColor.Resize(new Size(_imageHeight, _imageWidth));



        int dataSize = resized.Rows * resized.Cols * resized.ElemSize();
        var byteArray = new byte[dataSize];
        Marshal.Copy(resized.Data, byteArray, 0, dataSize);

        var result = _cSharpDataTransferMQRPC.PredictRaw(byteArray);

        byte[] resultColor;
        resultColor = result.Select(x => (byte)(x) ).ToArray();

        
        var outputSizeArray = outputSize.Select(Convert.ToInt32).ToArray();
        var mask = new Mat(outputSizeArray[0], outputSizeArray[1], MatType.CV_8UC1, resultColor);

        var maskResized = mask.Resize(new Size(currentImage.ImageData.Width, currentImage.ImageData.Height));

        context.ActiveImage = new HawkeyeImage() { ImageData = maskResized };


        Result = true;
    }

    public override void Save(BinaryWriter bw)
    {
        base.Save(bw);
        bw.Write(_modelFilePath.Path);
        bw.Write(_modelFilePath.Format);
        bw.Write(FilterClasses);
        bw.Write(IsCategorical);
    }

    public override void Load(BinaryReader br)
    {
        base.Load(br);
        _modelFilePath.Path = br.ReadString();
        _modelFilePath.Format = br.ReadString();
        FilterClasses = br.ReadString();
        IsCategorical = br.ReadBoolean();
    }

    public override void Save(Dictionary<string, object> dict)
    {
        base.Save(dict);
        dict[nameof(ModelFilePath)] = ModelFilePath.Path;
        dict[nameof(FilterClasses)] = FilterClasses;
        dict[nameof(IsCategorical)] = IsCategorical;
    }

    public override void Load(Dictionary<string, object> dict)
    {
        base.Load(dict);
        if (dict.ContainsKey(nameof(ModelFilePath)))
            ModelFilePath.Path = dict[nameof(ModelFilePath)].ToString();
        if (dict.ContainsKey(nameof(FilterClasses)))
            FilterClasses = dict[nameof(FilterClasses)].ToString();
        if (dict.ContainsKey(nameof(IsCategorical)))
            IsCategorical = Convert.ToBoolean(dict[nameof(IsCategorical)]);
    }
}