Merge pull request #1502 from patel-zeel/feat/oriented_box_iou_batch

Add `oriented_box_iou_batch` function to `detection.utils`

docs/detection/utils.md

@@ -16,6 +16,12 @@ comments: true
 
 :::supervision.detection.utils.mask_iou_batch
 
+<div class="md-typeset">
+  <h2><a href="#supervision.detection.utils.oriented_box_iou_batch">oriented_box_iou_batch</a></h2>
+</div>
+
+:::supervision.detection.utils.oriented_box_iou_batch
+
 <div class="md-typeset">
   <h2><a href="#supervision.detection.utils.polygon_to_mask">polygon_to_mask</a></h2>
 </div>

supervision/__init__.py

@@ -65,6 +65,7 @@ from supervision.detection.utils import (
     mask_to_xyxy,
     move_boxes,
     move_masks,
+    oriented_box_iou_batch,
     pad_boxes,
     polygon_to_mask,
     polygon_to_xyxy,

supervision/detection/utils.py

@@ -140,6 +140,45 @@ def mask_iou_batch(
     return np.vstack(ious)
 
 
+def oriented_box_iou_batch(
+    boxes_true: np.ndarray, boxes_detection: np.ndarray
+) -> np.ndarray:
+    """
+    Compute Intersection over Union (IoU) of two sets of oriented bounding boxes -
+    `boxes_true` and `boxes_detection`. Both sets of boxes are expected to be in
+    `((x1, y1), (x2, y2), (x3, y3), (x4, y4))` format.
+
+    Args:
+        boxes_true (np.ndarray): a `np.ndarray` representing ground-truth boxes.
+            `shape = (N, 4, 2)` where `N` is number of true objects.
+        boxes_detection (np.ndarray): a `np.ndarray` representing detection boxes.
+            `shape = (M, 4, 2)` where `M` is number of detected objects.
+
+    Returns:
+        np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`.
+            `shape = (N, M)` where `N` is number of true objects and
+            `M` is number of detected objects.
+    """
+
+    boxes_true = boxes_true.reshape(-1, 4, 2)
+    boxes_detection = boxes_detection.reshape(-1, 4, 2)
+
+    max_height = max(boxes_true[:, :, 0].max(), boxes_detection[:, :, 0].max()) + 1
+    # adding 1 because we are 0-indexed
+    max_width = max(boxes_true[:, :, 1].max(), boxes_detection[:, :, 1].max()) + 1
+
+    mask_true = np.zeros((boxes_true.shape[0], max_height, max_width))
+    for i, box_true in enumerate(boxes_true):
+        mask_true[i] = polygon_to_mask(box_true, (max_width, max_height))
+
+    mask_detection = np.zeros((boxes_detection.shape[0], max_height, max_width))
+    for i, box_detection in enumerate(boxes_detection):
+        mask_detection[i] = polygon_to_mask(box_detection, (max_width, max_height))
+
+    ious = mask_iou_batch(mask_true, mask_detection)
+    return ious
+
+
 def clip_boxes(xyxy: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
     """
     Clips bounding boxes coordinates to fit within the frame resolution.

supervision/metrics/utils/object_size.py

@@ -101,17 +101,20 @@ def get_obb_size_category(xyxyxyxy: npt.NDArray[np.float32]) -> npt.NDArray[np.i
     Get the size category of a oriented bounding boxes array.
 
     Args:
-        xyxyxyxy (np.ndarray): The bounding boxes array shaped (N, 8).
+        xyxyxyxy (np.ndarray): The bounding boxes array shaped (N, 4, 2).
 
     Returns:
         (np.ndarray) The size category of each bounding box, matching
         the enum values of ObjectSizeCategory. Shaped (N,).
     """
-    if len(xyxyxyxy.shape) != 2 or xyxyxyxy.shape[1] != 8:
-        raise ValueError("Oriented bounding boxes must be shaped (N, 8)")
+    if len(xyxyxyxy.shape) != 3 or xyxyxyxy.shape[1] != 4 or xyxyxyxy.shape[2] != 2:
+        raise ValueError("Oriented bounding boxes must be shaped (N, 4, 2)")
 
     # Shoelace formula
-    x1, y1, x2, y2, x3, y3, x4, y4 = xyxyxyxy.T
+    x = xyxyxyxy[:, :, 0]
+    y = xyxyxyxy[:, :, 1]
+    x1, x2, x3, x4 = x.T
+    y1, y2, y3, y4 = y.T
     areas = 0.5 * np.abs(
         (x1 * y2 + x2 * y3 + x3 * y4 + x4 * y1)
         - (x2 * y1 + x3 * y2 + x4 * y3 + x1 * y4)