annotations/Azaion.Annotator/YOLODetector.cs

using System.Drawing.Imaging;
using System.IO;
using Azaion.Annotator.DTO;
using Compunet.YoloV8;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Formats.Jpeg;
using SixLabors.ImageSharp.PixelFormats;

namespace Azaion.Annotator;

public interface IAIDetector
{
    List<(YoloLabel Label, float Probability)> Detect(Stream stream);
}

public class YOLODetector(Config config) : IAIDetector, IDisposable
{
    private readonly YoloPredictor _predictor = new(config.AIRecognitionConfig.AIModelPath);

    public List<(YoloLabel Label, float Probability)> Detect(Stream stream)
    {
        stream.Seek(0, SeekOrigin.Begin);
        var image = Image.Load<Rgb24>(stream);
        var result = _predictor.Detect(image);

        var imageSize = new System.Windows.Size(image.Width, image.Height);

        return result.Select(d =>
        {
            var label = new YoloLabel(new CanvasLabel(d.Name.Id, d.Bounds.X, d.Bounds.Y, d.Bounds.Width, d.Bounds.Height), imageSize, imageSize);
            return (label, d.Confidence * 100);
        }).ToList();
    }

    public void Dispose() => _predictor.Dispose();
}