annotations/Azaion.Annotator/YOLODetector.cs

using System.Diagnostics;
using System.IO;
using Azaion.Annotator.Extensions;
using Azaion.Common.DTO;
using Azaion.Common.DTO.Config;
using Azaion.CommonSecurity.Services;
using Compunet.YoloV8;
using Microsoft.Extensions.Options;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using Detection = Azaion.Common.DTO.Detection;

namespace Azaion.Annotator;

public interface IAIDetector
{
    Task<List<Detection>> Detect(string fName, Stream imageStream, CancellationToken cancellationToken = default);
}

public class YOLODetector(IOptions<AIRecognitionConfig> recognitionConfig, IResourceLoader resourceLoader) : IAIDetector, IDisposable
{
    private readonly AIRecognitionConfig _recognitionConfig = recognitionConfig.Value;
    private YoloPredictor? _predictor;
    private const string YOLO_MODEL = "azaion.onnx";


    public async Task<List<Detection>> Detect(string fName, Stream imageStream, CancellationToken cancellationToken)
    {
        if (_predictor == null)
        {
            await using var stream = await resourceLoader.Load(YOLO_MODEL, cancellationToken);
            _predictor = new YoloPredictor(stream.ToArray());
        }

        imageStream.Seek(0, SeekOrigin.Begin);

        using var image = Image.Load<Rgb24>(imageStream);
        var result = await _predictor.DetectAsync(image);
        var imageSize = new System.Windows.Size(image.Width, image.Height);
        var detections = 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 new Detection(fName, label, (double?)d.Confidence * 100);
        }).ToList();

        return FilterOverlapping(detections);
    }

    private List<Detection> FilterOverlapping(List<Detection> detections)
    {
        var k = _recognitionConfig.TrackingIntersectionThreshold;
        var filteredDetections = new List<Detection>();
        for (var i = 0; i < detections.Count; i++)
        {
            var detectionSelected = false;
            for (var j = i + 1; j < detections.Count; j++)
            {
                var intersect = detections[i].ToRectangle();
                intersect.Intersect(detections[j].ToRectangle());

                var maxArea = Math.Max(detections[i].ToRectangle().Area(), detections[j].ToRectangle().Area());
                if (!(intersect.Area() > k * maxArea)) 
                    continue;

                if (detections[i].Probability > detections[j].Probability)
                {
                    filteredDetections.Add(detections[i]);
                    detections.RemoveAt(j);
                }
                else
                {
                    filteredDetections.Add(detections[j]);
                    detections.RemoveAt(i);
                }
                detectionSelected = true;
                break;
            }

            if (!detectionSelected)
                filteredDetections.Add(detections[i]);
        }

        return filteredDetections;
    }

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