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

using System.Diagnostics;
using Hawkeye.VisionBuilder.Workflow.Datatypes;
using Hawkeye.VisionBuilder.Workflow.Operations.Attributes;
using static IronPython.Runtime.Profiler;
using System.Runtime.InteropServices;
using System.Text;
using Compunet.YoloV8;
using Compunet.YoloV8.Data;
using OpenCvSharp;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Processing;

namespace Hawkeye.VisionBuilder.Workflow.Operations.AI;
[Category("AI")]
public class YoloDetectionOperation:BaseOperation
{

    private bool _initialized = false;
    private FilePath _modelFilePath = new FilePath(){ Format = "Onnx format(*.onnx)|*.onnx" };


    public int IgnoreType { get; set; } = -1;

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

    bool EnsureInitialized()
    {
        if (!_initialized)
        {
            if (!File.Exists(ModelFilePath?.Path))
            {
                this.SetError("Model file not found");
                return false;
            }
            _predictor = YoloV8Predictor.Create(ModelFilePath?.Path);
            
            _predictor.Configuration.SuppressParallelInference=true;
            _initialized = true;
        }
        return true;
    }
    byte[] _buffer=new byte[0];
    private YoloV8Predictor _predictor;


    void EnsureBuffer(int size)
    {
        if (_buffer.Length != size)
        {
            _buffer = new byte[size];
        }
    }
    protected override void InterpretInternal(Context context)
    {
        CheckImageExists(context);
        CheckColorful(context);
        try
        {
            if (!EnsureInitialized()) return;
            var swImagePreparation = Stopwatch.StartNew();

            var image = context.ActiveImage.ImageData.CvtColor(ColorConversionCodes.BGR2RGB);
            EnsureBuffer(image.Cols * image.Rows * image.Channels());
            Marshal.Copy(image.Data, _buffer, 0, _buffer.Length);
            Image<Rgb24> img = SixLabors.ImageSharp.Image.LoadPixelData<Rgb24>(_buffer, image.Width, image.Height);
            // convert to rgb24
            
            
            swImagePreparation.Stop();
            
            var swPrediction = Stopwatch.StartNew();
            


            var result = _predictor.Detect(img);
            


            swPrediction.Stop();


            var swPostProcessing = Stopwatch.StartNew();
            Mat res = new Mat(image.Size(), MatType.CV_8UC1, new Scalar(0));

            StringBuilder sb = new StringBuilder();
            Dictionary<int, string> classes = new Dictionary<int, string>();



            foreach (BoundingBox box in result.Boxes.OrderBy(x=>x.Class.Id).ToList())
            {
                if (box.Confidence<0.7)
                {
                    continue;
                }

                if (box.Class.Id + 1 == IgnoreType)
                {
                    continue;
                }
                classes[box.Class.Id + 1] = box.Class.Name;
                var rect = box.Bounds;
                res.Rectangle(new OpenCvSharp.Rect(rect.Left, rect.Top, rect.Width, rect.Height),
                    new Scalar(box.Class.Id + 1), -1);
            }
            swPostProcessing.Stop();


            foreach (var c in classes)
            {
                sb.AppendLine($"{c.Key}={c.Value};");
            }

            Status = sb.ToString();

            context.ActiveImage = new HawkeyeImage() {ImageData = res};
        }
        catch (Exception e)
        {
            Console.WriteLine(e);
            throw;
        }
    }

    public override void Save(BinaryWriter bw)
    {
        base.Save(bw);
        bw.Write(ModelFilePath.Path);
    }

    public override void Load(BinaryReader br)
    {
        base.Load(br);
        ModelFilePath.Path = br.ReadString();
    }

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

    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(IgnoreType)))
            IgnoreType = Convert.ToInt32(dict[nameof(IgnoreType)]);
    }
}