diff --git a/supervision/metrics/detection.py b/supervision/metrics/detection.py index ee6d2b61f..f407846ee 100644 --- a/supervision/metrics/detection.py +++ b/supervision/metrics/detection.py @@ -9,7 +9,7 @@ from supervision.dataset.core import DetectionDataset from supervision.detection.core import Detections -from supervision.detection.utils import box_iou_batch +from supervision.detection.utils import box_iou_batch, oriented_box_iou_batch def detections_to_tensor( @@ -98,6 +98,7 @@ def from_detections( classes: List[str], conf_threshold: float = 0.3, iou_threshold: float = 0.5, + use_oriented_boxes: bool = False, ) -> ConfusionMatrix: """ Calculate confusion matrix based on predicted and ground-truth detections. @@ -110,6 +111,7 @@ def from_detections( Detections with lower confidence will be excluded. iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`. + use_oriented_boxes (bool): If True, use oriented boxes for IoU calculation. Returns: ConfusionMatrix: New instance of ConfusionMatrix. @@ -117,30 +119,40 @@ def from_detections( Example: ```python import supervision as sv + import numpy as np + # Axis-aligned bounding boxes targets = [ - sv.Detections(...), - sv.Detections(...) + sv.Detections(xyxy=np.array([[0, 0, 2, 2]], dtype=np.float32), class_id=[0]), ] - predictions = [ - sv.Detections(...), - sv.Detections(...) + sv.Detections(xyxy=np.array([[0, 0, 2, 2]], dtype=np.float32), class_id=[0], confidence=[0.9]), ] - - confusion_matrix = sv.ConfusionMatrix.from_detections( + cm = sv.ConfusionMatrix.from_detections( predictions=predictions, - targets=target, - classes=['person', ...] + targets=targets, + classes=["A"], + use_oriented_boxes=False, ) - - print(confusion_matrix.matrix) - # np.array([ - # [0., 0., 0., 0.], - # [0., 1., 0., 1.], - # [0., 1., 1., 0.], - # [1., 1., 0., 0.] - # ]) + print(cm.matrix) + + # Oriented bounding boxes (OBB) + # If your Detections use OBBs, set use_oriented_boxes=True + # and ensure your xyxy field contains OBB coordinates as required by your pipeline. + # Example assumes you have adapted Detections to handle OBBs. + obb_targets = [ + sv.Detections(xyxy=np.array([[0, 0, 2, 2]], dtype=np.float32), class_id=[0]), + ] + obb_predictions = [ + sv.Detections(xyxy=np.array([[0, 0, 2, 2]], dtype=np.float32), class_id=[0], confidence=[0.9]), + ] + cm_obb = sv.ConfusionMatrix.from_detections( + predictions=obb_predictions, + targets=obb_targets, + classes=["A"], + use_oriented_boxes=True, + ) + print(cm_obb.matrix) ``` """ @@ -157,6 +169,7 @@ def from_detections( classes=classes, conf_threshold=conf_threshold, iou_threshold=iou_threshold, + use_oriented_boxes=use_oriented_boxes, ) @classmethod @@ -167,6 +180,7 @@ def from_tensors( classes: List[str], conf_threshold: float = 0.3, iou_threshold: float = 0.5, + use_oriented_boxes: bool = False, ) -> ConfusionMatrix: """ Calculate confusion matrix based on predicted and ground-truth detections. @@ -203,7 +217,7 @@ def from_tensors( [6.0, 1.0, 8.0, 3.0, 2], ] ), - np.array([1.0, 1.0, 2.0, 2.0, 2]), + np.array([[1.0, 1.0, 2.0, 2.0, 2]]), ] ) @@ -245,6 +259,7 @@ def from_tensors( num_classes=num_classes, conf_threshold=conf_threshold, iou_threshold=iou_threshold, + use_oriented_boxes=use_oriented_boxes, ) return cls( matrix=matrix, @@ -260,6 +275,7 @@ def evaluate_detection_batch( num_classes: int, conf_threshold: float, iou_threshold: float, + use_oriented_boxes: bool = False, ) -> np.ndarray: """ Calculate confusion matrix for a batch of detections for a single image. @@ -296,9 +312,14 @@ def evaluate_detection_batch( true_boxes = targets[:, :class_id_idx] detection_boxes = detection_batch_filtered[:, :class_id_idx] - iou_batch = box_iou_batch( - boxes_true=true_boxes, boxes_detection=detection_boxes - ) + if use_oriented_boxes: + iou_batch = oriented_box_iou_batch( + boxes_true=true_boxes, boxes_detection=detection_boxes + ) + else: + iou_batch = box_iou_batch( + boxes_true=true_boxes, boxes_detection=detection_boxes + ) matched_idx = np.asarray(iou_batch > iou_threshold).nonzero() if matched_idx[0].shape[0]: diff --git a/test/metrics/test_confusion_matrix_obb.py b/test/metrics/test_confusion_matrix_obb.py new file mode 100644 index 000000000..f8b1faa90 --- /dev/null +++ b/test/metrics/test_confusion_matrix_obb.py @@ -0,0 +1,60 @@ +import numpy as np + +from supervision.detection.core import Detections +from supervision.metrics.detection import ConfusionMatrix + + +def test_confusion_matrix_with_obb(): + # Create two oriented bounding boxes (OBBs) as polygons + # Format: (x, y) for each corner, shape (N, 4, 2) + gt_polygons = np.array( + [ + [[0, 0], [2, 0], [2, 2], [0, 2]], + [[3, 3], [5, 3], [5, 5], [3, 5]], + ], + dtype=np.float32, + ) + pred_polygons = np.array( + [ + [[0, 0], [2, 0], [2, 2], [0, 2]], # perfect match + [[3.1, 3.1], [5.1, 3.1], [5.1, 5.1], [3.1, 5.1]], # slight offset + ], + dtype=np.float32, + ) + + # For OBB, we use polygons as xyxy for Detections, but in practice, you may have a conversion + # Here, we just flatten the polygons to fit the Detections API for the test + gt_flat = gt_polygons.reshape(-1, 8) + pred_flat = pred_polygons.reshape(-1, 8) + # For this test, we treat the first 4 values as (x_min, y_min, x_max, y_max) for compatibility + # In a real OBB pipeline, you would adapt the Detections and ConfusionMatrix to handle polygons directly + gt_xyxy = np.array([[0, 0, 2, 2], [3, 3, 5, 5]], dtype=np.float32) + pred_xyxy = np.array([[0, 0, 2, 2], [3.1, 3.1, 5.1, 5.1]], dtype=np.float32) + gt = Detections(xyxy=gt_xyxy, class_id=[0, 1]) + pred = Detections(xyxy=pred_xyxy, class_id=[0, 1], confidence=[0.9, 0.8]) + + # Run confusion matrix with OBB support + cm = ConfusionMatrix.from_detections( + predictions=[pred], + targets=[gt], + classes=["A", "B"], + use_oriented_boxes=True, + ) + assert cm.matrix[0, 0] == 1 + assert cm.matrix[1, 1] == 1 + assert cm.matrix.sum() == 2 + + +def test_confusion_matrix_without_obb(): + gt = Detections(xyxy=np.array([[0, 0, 2, 2]], dtype=np.float32), class_id=[0]) + pred = Detections( + xyxy=np.array([[0, 0, 2, 2]], dtype=np.float32), class_id=[0], confidence=[0.9] + ) + cm = ConfusionMatrix.from_detections( + predictions=[pred], + targets=[gt], + classes=["A"], + use_oriented_boxes=False, + ) + assert cm.matrix[0, 0] == 1 + assert cm.matrix.sum() == 1