142 lines
3.7 KiB
C#
142 lines
3.7 KiB
C#
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 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)
|
|
{
|
|
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 == 10)
|
|
{
|
|
|
|
}
|
|
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();
|
|
}
|
|
} |