add iou_batch_obb

boxmot/utils/iou.py

@@ -1,4 +1,35 @@
 import numpy as np
+import cv2 as cv
+
+def iou_obb_pair(i, j, bboxes1, bboxes2):
+    """
+    Compute IoU for the rotated rectangles at index i and j in the batches `bboxes1`, `bboxes2` .
+    """
+    rect1 = bboxes1[int(i)]
+    rect2 = bboxes2[int(j)]
+
+    (cx1, cy1, w1, h1, angle1) = rect1[0:5]
+    (cx2, cy2, w2, h2, angle2) = rect2[0:5]
+
+
+    r1 = ((cx1, cy1), (w1, h1), angle1)
+    r2 = ((cx2, cy2), (w2, h2), angle2)
+
+    # Compute intersection
+    ret, intersect = cv.rotatedRectangleIntersection(r1, r2)
+    if ret == 0 or intersect is None:
+        return 0.0  # No intersection
+
+    # Calculate intersection area
+    intersection_area = cv.contourArea(intersect)
+
+    # Calculate union area
+    area1 = w1 * h1
+    area2 = w2 * h2
+    union_area = area1 + area2 - intersection_area
+
+    # Compute IoU
+    return intersection_area / union_area if union_area > 0 else 0.0
 
 class AssociationFunction:
     def __init__(self, w, h, asso_mode="iou"):
@@ -35,6 +66,17 @@ def iou_batch(bboxes1, bboxes2) -> np.ndarray:
         )
         return o
 
+    @staticmethod
+    def iou_batch_obb(bboxes1, bboxes2) -> np.ndarray:
+
+        N, M = len(bboxes1), len(bboxes2)
+
+        def wrapper(i, j):
+            return iou_obb_pair(i, j, bboxes1, bboxes2)
+
+        iou_matrix = np.fromfunction(np.vectorize(wrapper), shape=(N, M), dtype=int)
+        return iou_matrix
+
     @staticmethod
     def hmiou_batch(bboxes1, bboxes2):
         """
@@ -291,6 +333,7 @@ def _get_asso_func(self, asso_mode):
         """
         ASSO_FUNCS = {
             "iou": AssociationFunction.iou_batch,
+            "iou_obb": AssociationFunction.iou_batch_obb,
             "hmiou": AssociationFunction.hmiou_batch,
             "giou": AssociationFunction.giou_batch,
             "ciou": AssociationFunction.ciou_batch,