annotations/Azaion.Annotator/YOLODetector.cs

using System.IO;
using Azaion.Annotator.DTO;
using Azaion.Annotator.Extensions;
using Compunet.YoloV8;
using Compunet.YoloV8.Data;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using Detection = Azaion.Annotator.DTO.Detection;

namespace Azaion.Annotator;

public interface IAIDetector
{
    List<Detection> Detect(Stream stream);
}

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

    public List<Detection> 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);

        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(label, (double?)d.Confidence * 100);
        }).ToList();

        return FilterOverlapping(detections);
    }

    private List<Detection> FilterOverlapping(List<Detection> detections)
    {
        var k = config.AIRecognitionConfig.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)
                {
                    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();
}