fix ai detection bugs #1

Azaion.Annotator/YOLODetector.cs

@@ -1,23 +1,24 @@
-using System.Drawing.Imaging;
-using System.IO;
+using System.IO;
 using Azaion.Annotator.DTO;
+using Azaion.Annotator.Extensions;
 using Compunet.YoloV8;
+using Compunet.YoloV8.Data;
 using SixLabors.ImageSharp;
-using SixLabors.ImageSharp.Formats.Jpeg;
 using SixLabors.ImageSharp.PixelFormats;
+using Detection = Azaion.Annotator.DTO.Detection;
 
 namespace Azaion.Annotator;
 
 public interface IAIDetector
 {
-    List<(YoloLabel Label, float Probability)> Detect(Stream stream);
+    List<Detection> 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)
+    public List<Detection> Detect(Stream stream)
     {
         stream.Seek(0, SeekOrigin.Begin);
         var image = Image.Load<Rgb24>(stream);
@@ -25,11 +26,50 @@ public class YOLODetector(Config config) : IAIDetector, IDisposable
 
         var imageSize = new System.Windows.Size(image.Width, image.Height);
 
-        return result.Select(d =>
+        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 (label, d.Confidence * 100);
+            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();