135 lines
4.2 KiB
C#
135 lines
4.2 KiB
C#
using Hawkeye.VisionBuilder.Workflow.Datatypes;
|
|
using Hawkeye.VisionBuilder.Workflow.Operations.Attributes;
|
|
using ILGPU.Runtime.Cuda;
|
|
using Microsoft.ML.OnnxRuntime;
|
|
using Microsoft.ML.OnnxRuntime.Tensors;
|
|
using Microsoft.Scripting.Runtime;
|
|
using OpenCvSharp;
|
|
namespace Hawkeye.VisionBuilder.Workflow.Operations.AI;
|
|
|
|
[Category("AI")]
|
|
public class OnnxModelOperation:BaseOperation
|
|
{
|
|
|
|
private FilePath _modelFilePath = new FilePath(){Format = "ONNX file(*.onnx)|*.onnx" };
|
|
|
|
|
|
public FilePath ModelFilePath
|
|
{
|
|
get => _modelFilePath;
|
|
set
|
|
{
|
|
_modelFilePath = value;
|
|
_initialized = false;
|
|
}
|
|
}
|
|
|
|
private bool _initialized = false;
|
|
private InferenceSession _AIsession;
|
|
private string _inputName;
|
|
private int[] _dimensions;
|
|
|
|
protected override void InterpretInternal(Context context)
|
|
{
|
|
if (!CheckImageExists(context)) return;
|
|
if (!CheckColorful(context)) return;
|
|
|
|
if (!_initialized)
|
|
{
|
|
if (!File.Exists(ModelFilePath?.Path))
|
|
{
|
|
this.SetError($"Model file \"{ModelFilePath?.Path}\" not found");
|
|
return;
|
|
}
|
|
|
|
var opts = new SessionOptions();
|
|
_AIsession = new InferenceSession(ModelFilePath?.Path, opts);
|
|
_inputName = _AIsession.InputMetadata.First().Key;
|
|
_dimensions = _AIsession.InputMetadata.First().Value.Dimensions;
|
|
foreach (var input in _AIsession.InputMetadata)
|
|
{
|
|
Console.WriteLine($"Name: {input.Key}");
|
|
Console.WriteLine($" Type: {input.Value.ElementType}");
|
|
Console.WriteLine($" Shape: [{string.Join(", ", input.Value.Dimensions)}]");
|
|
}
|
|
|
|
_initialized = true;
|
|
}
|
|
|
|
var currentImage = context.ActiveImage;
|
|
|
|
var rightColor = currentImage.ImageData.CvtColor(ColorConversionCodes.BGR2RGB);
|
|
|
|
// convert to float32
|
|
var floatImage = new Mat();
|
|
rightColor.ConvertTo(floatImage, MatType.CV_32FC3, 1.0 / 255);
|
|
// to array
|
|
|
|
var width = _dimensions[1];
|
|
var height = _dimensions[2];
|
|
|
|
// check if need resize
|
|
if (floatImage.Width != width || floatImage.Height != height)
|
|
Cv2.Resize(floatImage, floatImage, new OpenCvSharp.Size(width, height));
|
|
|
|
|
|
float[] data = new float[3 * width * height];
|
|
int idx = 0;
|
|
var rows = floatImage.Rows;
|
|
var cols = floatImage.Cols;
|
|
for (int y = 0; y < rows; y++)
|
|
{
|
|
for (int x = 0; x < cols; x++)
|
|
{
|
|
Vec3f pixel = floatImage.At<Vec3f>(y, x);
|
|
for (int c = 0; c < 3; c++) // channels
|
|
{
|
|
data[idx++] = pixel[c]; // width*height*channel
|
|
}
|
|
}
|
|
}
|
|
|
|
var inputTensor = new DenseTensor<float>(
|
|
data,
|
|
new int[] { 1, floatImage.Width, floatImage.Height, 3 }
|
|
);
|
|
|
|
var inputs = new List<NamedOnnxValue> {
|
|
NamedOnnxValue.CreateFromTensor("input", inputTensor)
|
|
};
|
|
var sw = System.Diagnostics.Stopwatch.StartNew();
|
|
using var results = _AIsession.Run(inputs);
|
|
sw.Stop();
|
|
Console.WriteLine($"Inference time: {sw.ElapsedMilliseconds} ms");
|
|
var output = results.First().AsEnumerable<float>().ToArray();
|
|
|
|
// convert output to Mat and normalize to 0-255
|
|
var outputMat = new Mat(new OpenCvSharp.Size(width, height), MatType.CV_32FC1);
|
|
outputMat.SetArray<float>(output);
|
|
|
|
outputMat.ConvertTo(outputMat, MatType.CV_8UC1, 255);
|
|
|
|
// resize back to original size
|
|
Cv2.Resize(outputMat, outputMat, new OpenCvSharp.Size(currentImage.ImageData.Width, currentImage.ImageData.Height));
|
|
|
|
|
|
context.ActiveImage = new HawkeyeImage()
|
|
{
|
|
ImageData = outputMat
|
|
};
|
|
|
|
}
|
|
|
|
public override void Save(Dictionary<string, object> dict)
|
|
{
|
|
base.Save(dict);
|
|
dict[nameof(ModelFilePath)] = ModelFilePath?.Path;
|
|
}
|
|
|
|
public override void Load(Dictionary<string, object> dict)
|
|
{
|
|
base.Load(dict);
|
|
if (dict.ContainsKey(nameof(ModelFilePath)))
|
|
ModelFilePath.Path = dict[nameof(ModelFilePath)].ToString();
|
|
}
|
|
} |