Merge pull request #46 from bnemetchek2/main

Add OBB

Source/YoloV8/Data/ObbBoundingBox.cs

@@ -0,0 +1,6 @@
+namespace Compunet.YoloV8.Data;
+
+public class ObbBoundingBox : BoundingBox
+{
+    public required float Angle { get; init; }
+}

Source/YoloV8/Data/ObbDetectionResult.cs

@@ -0,0 +1,8 @@
+namespace Compunet.YoloV8.Data;
+
+public class ObbDetectionResult : YoloV8Result
+{
+    public required ObbBoundingBox[] Boxes { get; init; }
+
+    public override string ToString() => Boxes.Summary();
+}

Source/YoloV8/Extensions/OrientedBoundingBoxExtensions.cs

@@ -0,0 +1,57 @@
+namespace Compunet.YoloV8.Extensions;
+
+internal static class OrientedBoundingBoxExtensions
+{
+    public static Point[] GetCornerPoints(this ObbBoundingBox obb)
+    {
+        return GetCornerPoints(obb.Bounds, obb.Angle);
+    }
+
+    public static Point[] GetCornerPoints(this ObbIndexedBoundingBox obb)
+    {
+        return GetCornerPoints(obb.Bounds, obb.Angle);
+    }
+
+    private static Point[] GetCornerPoints(Rectangle bounds, float _angle)
+    {
+        var angle = _angle * Math.PI / 180.0; // Radians
+
+        var b = (float)Math.Cos(angle) * 0.5f;
+        var a = (float)Math.Sin(angle) * 0.5f;
+
+        var x = bounds.X;
+        var y = bounds.Y;
+        var w = bounds.Width;
+        var h = bounds.Height;
+
+        var points = new Point[4];
+
+        points[0].X = (int)Math.Round(x - a * h - b * w, 0);
+        points[0].Y = (int)Math.Round(y + b * h - a * w, 0);
+
+        points[1].X = (int)Math.Round(x + a * h - b * w, 0);
+        points[1].Y = (int)Math.Round(y - b * h - a * w, 0);
+
+        points[2].X = (int)Math.Round(2d * x - points[0].X, 0);
+        points[2].Y = (int)Math.Round(2d * y - points[0].Y, 0);
+
+        points[3].X = (int)Math.Round(2d * x - points[1].X, 0);
+        points[3].Y = (int)Math.Round(2d * y - points[1].Y, 0);
+
+        // Calculate the distances of each point from the origin (0, 0)
+        var distance1 = Math.Sqrt(Math.Pow(points[0].X, 2) + Math.Pow(points[0].Y, 2));
+        var distance2 = Math.Sqrt(Math.Pow(points[1].X, 2) + Math.Pow(points[1].Y, 2));
+
+        // Rotate if necessary to ensure pt[0] is the top-left point
+        if (distance2 < distance1)
+        {
+            var temp = points[0];
+            points[0] = points[1];
+            points[1] = points[2];
+            points[2] = points[3];
+            points[3] = temp;
+        }
+
+        return points;
+    }
+}

Source/YoloV8/Parsers/ObbDetectionOutputParser.cs

@@ -0,0 +1,26 @@
+namespace Compunet.YoloV8;
+
+internal readonly struct ObbDetectionOutputParser(YoloV8Metadata metadata, YoloV8Configuration configuration)
+{
+    public ObbBoundingBox[] Parse(Tensor<float> output, Size originSize)
+    {
+        var boxes = new ObbIndexedBoundingBoxParser(metadata, configuration).Parse(output, originSize);
+
+        var result = new ObbBoundingBox[boxes.Length];
+
+        for (int i = 0; i < boxes.Length; i++)
+        {
+            var box = boxes[i];
+
+            result[i] = new ObbBoundingBox
+            {
+                Class = box.Class,
+                Bounds = box.Bounds,
+                Angle = box.Angle,
+                Confidence = box.Confidence,
+            };
+        }
+
+        return result;
+    }
+}

Source/YoloV8/Parsers/ObbIndexedBoundingBox.cs

@@ -0,0 +1,18 @@
+namespace Compunet.YoloV8.Parsers;
+
+internal readonly struct ObbIndexedBoundingBox : IComparable<ObbIndexedBoundingBox>
+{
+    public bool IsEmpty => Bounds == default;
+
+    public required int Index { get; init; }
+
+    public required YoloV8Class Class { get; init; }
+
+    public required Rectangle Bounds { get; init; }
+
+    public required float Angle { get; init; }
+
+    public required float Confidence { get; init; }
+
+    public int CompareTo(ObbIndexedBoundingBox other) => Confidence.CompareTo(other.Confidence);
+}

Source/YoloV8/Parsers/ObbIndexedBoundingBoxParser.cs

@@ -0,0 +1,128 @@
+// https://github.com/ultralytics/ultralytics/issues/7667
+
+namespace Compunet.YoloV8;
+
+internal readonly struct ObbIndexedBoundingBoxParser(YoloV8Metadata metadata, YoloV8Configuration configuration)
+{
+    public ObbIndexedBoundingBox[] Parse(Tensor<float> output, SixLabors.ImageSharp.Size originSize)
+    {
+        int xPadding;
+        int yPadding;
+
+        if (configuration.KeepOriginalAspectRatio)
+        {
+            var reductionRatio = Math.Min(metadata.ImageSize.Width / (float)originSize.Width,
+                                          metadata.ImageSize.Height / (float)originSize.Height);
+
+            xPadding = (int)((metadata.ImageSize.Width - originSize.Width * reductionRatio) / 2);
+            yPadding = (int)((metadata.ImageSize.Height - originSize.Height * reductionRatio) / 2);
+        }
+        else
+        {
+            xPadding = 0;
+            yPadding = 0;
+        }
+
+        return Parse(output, originSize, xPadding, yPadding);
+    }
+
+    public ObbIndexedBoundingBox[] Parse(Tensor<float> output, SixLabors.ImageSharp.Size originSize, int xPadding, int yPadding)
+    {
+        var xRatio = (float)originSize.Width / metadata.ImageSize.Width;
+        var yRatio = (float)originSize.Height / metadata.ImageSize.Height;
+
+        if (configuration.KeepOriginalAspectRatio)
+        {
+            var maxRatio = Math.Max(xRatio, yRatio);
+
+            xRatio = maxRatio;
+            yRatio = maxRatio;
+        }
+
+        return Parse(output, originSize, xPadding, yPadding, xRatio, yRatio);
+    }
+
+    public ObbIndexedBoundingBox[] Parse(Tensor<float> output, Size originSize, int xPadding, int yPadding, float xRatio, float yRatio)
+    {
+        var _metadata = metadata;
+        var _parameters = configuration;
+
+        var detectionDataSize = output.Dimensions[1];
+        var boxes = new ObbIndexedBoundingBox[output.Dimensions[2]];
+
+        Parallel.For(0, output.Dimensions[2], i =>
+        {
+            var maxConfidence = _parameters.Confidence;
+            var maxConfidenceIndex = -1;
+
+            for (int j = 0; j < _metadata.Names.Count; j++)
+            {
+                var confidence = output[0, j + 4, i];
+
+                if (confidence > maxConfidence)
+                {
+                    maxConfidence = confidence;
+                    maxConfidenceIndex = j;
+                }
+            }
+
+            if (maxConfidenceIndex == -1)
+            {
+                return;
+            }
+
+            var x = (int)((output[0, 0, i] - xPadding) * xRatio);
+            var y = (int)((output[0, 1, i] - yPadding) * yRatio);
+            var w = (int)(output[0, 2, i] * xRatio);
+            var h = (int)(output[0, 3, i] * yRatio);
+
+            var bounds = new Rectangle(x, y, w, h);
+
+            var angle = (output[0, detectionDataSize - 1, i]); // Radians
+                                                               // Angle in [-pi/4,3/4 pi) --》 [-pi/2,pi/2)
+            if (angle >= Math.PI && angle <= 0.75 * Math.PI)
+            {
+                angle -= (float)Math.PI;
+            }
+
+            var name = _metadata.Names[maxConfidenceIndex];
+
+            boxes[i] = new ObbIndexedBoundingBox
+            {
+                Index = i,
+                Class = name,
+                Bounds = bounds,
+                Angle = (float)(angle * 180 / Math.PI), // Degrees
+                Confidence = maxConfidence
+            };
+        });
+
+        var count = 0;
+
+        for (int i = 0; i < boxes.Length; i++)
+        {
+            if (boxes[i].IsEmpty == false)
+            {
+                count++;
+            }
+        }
+
+        var topBoxes = new ObbIndexedBoundingBox[count];
+
+        var topIndex = 0;
+
+        for (int i = 0; i < boxes.Length; i++)
+        {
+            var box = boxes[i];
+
+            if (box.IsEmpty)
+            {
+                continue;
+            }
+
+            topBoxes[topIndex++] = box;
+        }
+
+        return ObbNonMaxSuppressionHelper.Suppress(topBoxes, configuration.IoU);
+    }
+}