Speed up evaluation with r-trees to find overlapping detections

fiftyone/utils/eval/coco.py

@@ -527,12 +527,9 @@ def _coco_evaluation_setup(
         # Sort ground truth so crowds are last
         gts = sorted(gts, key=iscrowd)
 
-        # Compute ``num_preds x num_gts`` IoUs
+        # Compute IoUs for overlapping detections
         ious = foui.compute_ious(preds, gts, **iou_kwargs)
-
-        gt_ids = [g.id for g in gts]
-        for pred, gt_ious in zip(preds, ious):
-            pred_ious[pred.id] = list(zip(gt_ids, gt_ious))
+        pred_ious.update(ious)
 
     return cats, pred_ious, iscrowd
 

fiftyone/utils/eval/openimages.py

@@ -564,10 +564,7 @@ def _open_images_evaluation_setup(
 
         # Compute ``num_preds x num_gts`` IoUs
         ious = foui.compute_ious(preds, gts, **iou_kwargs)
-
-        gt_ids = [g.id for g in gts]
-        for pred, gt_ious in zip(preds, ious):
-            pred_ious[pred.id] = list(zip(gt_ids, gt_ious))
+        pred_ious.update(ious)
 
     return cats, pred_ious, iscrowd
 

fiftyone/utils/iou.py

@@ -11,6 +11,7 @@
 
 import numpy as np
 import scipy.spatial as sp
+import rtree.index as rti
 
 import eta.core.numutils as etan
 import eta.core.utils as etau
@@ -19,7 +20,7 @@
 import fiftyone.core.utils as fou
 import fiftyone.core.validation as fov
 
-from .utils3d import compute_cuboid_iou
+from .utils3d import compute_cuboid_iou, _Box
 
 sg = fou.lazy_import("shapely.geometry")
 so = fou.lazy_import("shapely.ops")
@@ -84,8 +85,11 @@ def compute_ious(
     Returns:
         a ``num_preds x num_gts`` array of IoUs
     """
-    if not preds or not gts:
-        return np.zeros((len(preds), len(gts)))
+    if not preds:
+        return {}
+
+    if not gts:
+        return dict((p.id, []) for p in preds)
 
     if etau.is_str(iscrowd):
         iscrowd = lambda l: bool(l.get_attribute_value(iscrowd, False))
@@ -485,6 +489,37 @@ def compute_bbox_iou(gt, pred, gt_crowd=False):
     return min(etan.safe_divide(inter, union), 1)
 
 
+def _get_detection_box(det, dimension=None):
+    if dimension is None:
+        dimension = _get_bbox_dim(det)
+
+    if dimension == 2:
+        px, py, pw, ph = det.bounding_box
+        xmin = px
+        xmax = px + pw
+        ymin = py
+        ymax = py + ph
+        result = (xmin, xmax, ymin, ymax)
+    elif dimension == 3:
+        box = _Box(det.rotation, det.location, det.dimensions)
+        xmin = np.min(box.vertices[:, 0])
+        xmax = np.max(box.vertices[:, 0])
+        ymin = np.min(box.vertices[:, 1])
+        ymax = np.max(box.vertices[:, 1])
+        zmin = np.min(box.vertices[:, 2])
+        zmax = np.max(box.vertices[:, 2])
+        result = (xmin, xmax, ymin, ymax, zmin, zmax)
+    else:
+        raise Exception(f"dimension should be 2 or 3, but got {dimension}")
+
+    return result
+
+
+def _get_poly_box(x):
+    detection = x.to_detection()
+    return _get_detection_box(detection)
+
+
 def _compute_bbox_ious(preds, gts, iscrowd=None, classwise=False):
     is_symmetric = preds is gts
 
@@ -501,28 +536,36 @@ def _compute_bbox_ious(preds, gts, iscrowd=None, classwise=False):
         else:
             gts = _polylines_to_detections(gts)
 
+    index_property = rti.Property()
     if _get_bbox_dim(gts[0]) == 3:
         bbox_iou_fcn = compute_cuboid_iou
+        index_property.dimension = 3
     else:
         bbox_iou_fcn = compute_bbox_iou
-
-    ious = np.zeros((len(preds), len(gts)))
-
-    for j, (gt, gt_crowd) in enumerate(zip(gts, gt_crowds)):
-
-        for i, pred in enumerate(preds):
-            if is_symmetric and i < j:
-                iou = ious[j, i]
-            elif is_symmetric and i == j:
-                iou = 1
-            elif classwise and pred.label != gt.label:
+        index_property.dimension = 2
+
+    rtree_index = rti.Index(properties=index_property, interleaved=False)
+    for i, gt in enumerate(gts):
+        box = _get_detection_box(gt, dimension=index_property.dimension)
+        rtree_index.insert(i, box)
+
+    pred_ious = {}
+    for i, pred in enumerate(preds):
+        box = _get_detection_box(pred, dimension=index_property.dimension)
+        indices = rtree_index.intersection(box)
+        pred_ious[pred.id] = []
+        for j in indices:  # pylint: disable=not-an-iterable
+            gt = gts[j]
+            gt_crowd = gt_crowds[j]
+            if classwise and pred.label != gt.label:
                 continue
-            else:
-                iou = bbox_iou_fcn(gt, pred, gt_crowd=gt_crowd)
 
-            ious[i, j] = iou
+            iou = bbox_iou_fcn(gt, pred, gt_crowd=gt_crowd)
+            pred_ious[pred.id].append((gt.id, iou))
+            if is_symmetric:
+                pred_ious[gt.id].append((pred.id, iou))
 
-    return ious
+    return pred_ious
 
 
 def _compute_polygon_ious(
@@ -564,57 +607,71 @@ def _compute_polygon_ious(
         elif gt_crowds is None:
             gt_crowds = [False] * num_gt
 
-        ious = np.zeros((num_pred, num_gt))
+        rtree_index = rti.Index(interleaved=False)
+        for i, gt in enumerate(gts):
+            box = _get_poly_box(gt)
+            rtree_index.insert(i, box)
 
-        for j, (gt_poly, gt_label, gt_area, gt_crowd) in enumerate(
-            zip(gt_polys, gt_labels, gt_areas, gt_crowds)
+        pred_ious = {}
+        for i, (pred, pred_poly, pred_label, pred_area) in enumerate(
+            zip(preds, pred_polys, pred_labels, pred_areas)
         ):
-            for i, (pred_poly, pred_label, pred_area) in enumerate(
-                zip(pred_polys, pred_labels, pred_areas)
-            ):
-                if is_symmetric and i < j:
-                    iou = ious[j, i]
-                elif is_symmetric and i == j:
-                    iou = 1
-                elif classwise and pred_label != gt_label:
+            pred_ious[pred.id] = []
+
+            if is_symmetric:
+                pred_ious[pred.id].append((pred.id, 1))
+
+            box = _get_poly_box(pred)
+            indices = rtree_index.intersection(box)
+            for j in indices:  # pylint: disable=not-an-iterable
+                gt = gts[j]
+                gt_poly = gt_polys[j]
+                gt_label = gt_labels[j]
+                gt_area = gt_areas[j]
+                gt_crowd = gt_crowds[j]
+
+                if classwise and pred_label != gt_label:
                     continue
-                else:
-                    try:
-                        inter = gt_poly.intersection(pred_poly).area
-                    except Exception as e:
-                        inter = 0.0
-                        fou.handle_error(
-                            ValueError(
-                                "Failed to compute intersection of predicted "
-                                "object '%s' and ground truth object '%s'"
-                                % (preds[i].id, gts[j].id)
-                            ),
-                            error_level,
-                            base_error=e,
-                        )
-
-                    if gt_crowd:
-                        union = pred_area
-                    else:
-                        union = pred_area + gt_area - inter
 
-                    iou = min(etan.safe_divide(inter, union), 1)
+                try:
+                    inter = gt_poly.intersection(pred_poly).area
+                except Exception as e:
+                    inter = 0.0
+                    fou.handle_error(
+                        ValueError(
+                            "Failed to compute intersection of predicted "
+                            "object '%s' and ground truth object '%s'"
+                            % (preds[i].id, gts[j].id)
+                        ),
+                        error_level,
+                        base_error=e,
+                    )
+
+                if gt_crowd:
+                    union = pred_area
+                else:
+                    union = pred_area + gt_area - inter
 
-                ious[i, j] = iou
+                iou = min(etan.safe_divide(inter, union), 1)
+                pred_ious[pred.id].append((gt.id, iou))
+                if is_symmetric:
+                    pred_ious[gt.id].append((pred.id, iou))
 
-    return ious
+    return pred_ious
 
 
 def _compute_polyline_similarities(preds, gts, classwise=False):
-    sims = np.zeros((len(preds), len(gts)))
-    for j, gt in enumerate(gts):
-        for i, pred in enumerate(preds):
+    sims = {}
+    for pred in preds:
+        sims[pred.id] = []
+        for gt in gts:
             if classwise and pred.label != gt.label:
                 continue
 
             gtp = list(itertools.chain.from_iterable(gt.points))
             predp = list(itertools.chain.from_iterable(pred.points))
-            sims[i, j] = _compute_object_keypoint_similarity(gtp, predp)
+            sim = _compute_object_keypoint_similarity(gtp, predp)
+            sims[pred.id].append((gt.id, sim))
 
     return sims
 
@@ -687,15 +744,17 @@ def _compute_segment_ious(preds, gts):
 
 
 def _compute_keypoint_similarities(preds, gts, classwise=False):
-    sims = np.zeros((len(preds), len(gts)))
-    for j, gt in enumerate(gts):
-        for i, pred in enumerate(preds):
+    sims = {}
+    for pred in preds:
+        sims[pred.id] = []
+        for gt in gts:
             if classwise and pred.label != gt.label:
                 continue
 
             gtp = gt.points
             predp = pred.points
-            sims[i, j] = _compute_object_keypoint_similarity(gtp, predp)
+            sim = _compute_object_keypoint_similarity(gtp, predp)
+            sims[pred.id].append((gt.id, sim))
 
     return sims
 

fiftyone/utils/utils3d.py

@@ -167,7 +167,7 @@ def __init__(self, rotation, location, scale):
         vertices = np.zeros((_NUM_KEYPOINTS, 3))
         for i in range(_NUM_KEYPOINTS):
             vertices[i, :] = (
-                np.matmul(rotation, scaled_identity_box[i, :])
+                np.matmul(self.rotation, scaled_identity_box[i, :])
                 + location.flatten()
             )
 

setup.py

@@ -60,6 +60,7 @@ def get_version():
     "PyYAML",
     "regex",
     "retrying",
+    "rtree",
     "scikit-learn",
     "scikit-image",
     "scipy",

tests/unittests/evaluation_tests.py

@@ -1936,9 +1936,15 @@ def _make_box(self, dimensions, location, rotation):
     def _check_iou(self, dataset, field1, field2, expected_iou):
         dets1 = dataset.first()[field1].detections
         dets2 = dataset.first()[field2].detections
-        actual_iou = foui.compute_ious(dets1, dets2)[0][0]
+        ious = foui.compute_ious(dets1, dets2)
+        result = list(ious.values())[0]
 
-        self.assertTrue(np.isclose(actual_iou, expected_iou))
+        if expected_iou == 0:
+            self.assertTrue(len(result) == 0)
+
+        else:
+            _, actual_iou = result[0]
+            self.assertTrue(np.isclose(actual_iou, expected_iou))
 
     @drop_datasets
     def test_non_overlapping_boxes(self):

tests/unittests/utils3d_tests.py

@@ -10,6 +10,7 @@
 import unittest
 
 import numpy as np
+import numpy.testing as nptest
 import open3d as o3d
 from PIL import Image
 
@@ -389,6 +390,16 @@ def test_get_scene_asset_paths(self):
             )
 
 
+class BoxTests(unittest.TestCase):
+    def test_box_vertices(self):
+        rotation = [0, 0, 0]
+        location = [0, 0, 0]
+        scale = [1, 1, 1]
+        box = fou3d._Box(rotation, location, scale)
+        expected = fou3d._UNIT_BOX
+        nptest.assert_equal(expected, box.vertices)
+
+
 if __name__ == "__main__":
     fo.config.show_progress_bars = False
     unittest.main(verbosity=2)