valor-lite 0.33.2__py3-none-any.whl → 0.33.4__py3-none-any.whl

valor_lite/detection/manager.py

@@ -3,12 +3,20 @@ from dataclasses import dataclass
 
 import numpy as np
 from numpy.typing import NDArray
+from shapely.geometry import Polygon as ShapelyPolygon
 from tqdm import tqdm
-from valor_lite.detection.annotation import Detection
+from valor_lite.detection.annotation import (
+    Bitmask,
+    BoundingBox,
+    Detection,
+    Polygon,
+)
 from valor_lite.detection.computation import (
+    compute_bbox_iou,
+    compute_bitmask_iou,
     compute_detailed_counts,
-    compute_iou,
     compute_metrics,
+    compute_polygon_iou,
     compute_ranked_pairs,
 )
 from valor_lite.detection.metric import (
@@ -35,7 +43,7 @@ Usage
 -----
 
 loader = DataLoader()
-loader.add_data(
+loader.add_bounding_boxes(
     groundtruths=groundtruths,
     predictions=predictions,
 )
@@ -51,6 +59,103 @@ filtered_metrics = evaluator.evaluate(iou_thresholds=[0.5], filter_mask=filter_m
 """
 
 
+def _get_valor_dict_annotation_key(
+    annotation_type: type[BoundingBox] | type[Polygon] | type[Bitmask],
+) -> str:
+    """Get the correct JSON key to extract a given annotation type."""
+
+    if issubclass(annotation_type, BoundingBox):
+        return "bounding_box"
+    if issubclass(annotation_type, Polygon):
+        return "polygon"
+    else:
+        return "raster"
+
+
+def _get_annotation_representation(
+    annotation_type: type[BoundingBox] | type[Polygon] | type[Bitmask],
+) -> str:
+    """Get the correct representation of an annotation object."""
+
+    representation = (
+        "extrema"
+        if issubclass(annotation_type, BoundingBox)
+        else ("mask" if issubclass(annotation_type, Bitmask) else "shape")
+    )
+
+    return representation
+
+
+def _get_annotation_representation_from_valor_dict(
+    data: list,
+    annotation_type: type[BoundingBox] | type[Polygon] | type[Bitmask],
+) -> tuple[float, float, float, float] | ShapelyPolygon | NDArray[np.bool_]:
+    """Get the correct representation of an annotation object from a valor dictionary."""
+
+    if issubclass(annotation_type, BoundingBox):
+        x = [point[0] for shape in data for point in shape]
+        y = [point[1] for shape in data for point in shape]
+        return (min(x), max(x), min(y), max(y))
+    if issubclass(annotation_type, Polygon):
+        return ShapelyPolygon(data)
+    else:
+        return np.array(data)
+
+
+def _get_annotation_data(
+    keyed_groundtruths: dict,
+    keyed_predictions: dict,
+    annotation_type: type[BoundingBox] | type[Polygon] | type[Bitmask] | None,
+    key=int,
+) -> np.ndarray:
+    """Create an array of annotation pairs for use when calculating IOU. Needed because we unpack bounding box representations, but not bitmask or polygon representations."""
+    if annotation_type == BoundingBox:
+        return np.array(
+            [
+                np.array([*gextrema, *pextrema])
+                for _, _, _, pextrema in keyed_predictions[key]
+                for _, _, gextrema in keyed_groundtruths[key]
+            ]
+        )
+    else:
+        return np.array(
+            [
+                np.array([groundtruth_obj, prediction_obj])
+                for _, _, _, prediction_obj in keyed_predictions[key]
+                for _, _, groundtruth_obj in keyed_groundtruths[key]
+            ]
+        )
+
+
+def compute_iou(
+    data: NDArray[np.floating],
+    annotation_type: type[BoundingBox] | type[Polygon] | type[Bitmask],
+) -> NDArray[np.floating]:
+    """
+    Computes intersection-over-union (IoU) calculations for various annotation types.
+
+    Parameters
+    ----------
+    data : NDArray[np.floating]
+        A sorted array of bounding box, bitmask, or polygon pairs.
+    annotation_type: type[BoundingBox] | type[Polygon] | type[Bitmask]
+        The type of annotation contained in the data.
+
+
+    Returns
+    -------
+    NDArray[np.floating]
+        Computed IoU's.
+    """
+
+    if annotation_type == BoundingBox:
+        return compute_bbox_iou(data=data)
+    elif annotation_type == Bitmask:
+        return compute_bitmask_iou(data=data)
+    else:
+        return compute_polygon_iou(data=data)
+
+
 @dataclass
 class Filter:
     indices: NDArray[np.int32]
@@ -74,6 +179,10 @@ class Evaluator:
         self.uid_to_index: dict[str, int] = dict()
         self.index_to_uid: dict[int, str] = dict()
 
+        # annotation reference
+        self.groundtruth_examples: dict[int, NDArray[np.float16]] = dict()
+        self.prediction_examples: dict[int, NDArray[np.float16]] = dict()
+
         # label reference
         self.label_to_index: dict[tuple[str, str], int] = dict()
         self.index_to_label: dict[int, tuple[str, str]] = dict()
@@ -84,10 +193,10 @@ class Evaluator:
         self.label_index_to_label_key_index: dict[int, int] = dict()
 
         # computation caches
-        self._detailed_pairs = np.array([])
-        self._ranked_pairs = np.array([])
-        self._label_metadata = np.array([])
-        self._label_metadata_per_datum = np.array([])
+        self._detailed_pairs: NDArray[np.floating] = np.array([])
+        self._ranked_pairs: NDArray[np.floating] = np.array([])
+        self._label_metadata: NDArray[np.int32] = np.array([])
+        self._label_metadata_per_datum: NDArray[np.int32] = np.array([])
 
     @property
     def ignored_prediction_labels(self) -> list[tuple[str, str]]:
@@ -163,46 +272,35 @@ class Evaluator:
                     [self.uid_to_index[uid] for uid in datum_uids],
                     dtype=np.int32,
                 )
-            mask = np.zeros_like(mask_pairs, dtype=np.bool_)
-            mask[
-                np.isin(self._ranked_pairs[:, 0].astype(int), datum_uids)
-            ] = True
-            mask_pairs &= mask
-
-            mask = np.zeros_like(mask_datums, dtype=np.bool_)
-            mask[datum_uids] = True
-            mask_datums &= mask
+            mask_pairs[
+                ~np.isin(self._ranked_pairs[:, 0].astype(int), datum_uids)
+            ] = False
+            mask_datums[~np.isin(np.arange(n_datums), datum_uids)] = False
 
         if labels is not None:
             if isinstance(labels, list):
                 labels = np.array(
                     [self.label_to_index[label] for label in labels]
                 )
-            mask = np.zeros_like(mask_pairs, dtype=np.bool_)
-            mask[np.isin(self._ranked_pairs[:, 4].astype(int), labels)] = True
-            mask_pairs &= mask
-
-            mask = np.zeros_like(mask_labels, dtype=np.bool_)
-            mask[labels] = True
-            mask_labels &= mask
+            mask_pairs[
+                ~np.isin(self._ranked_pairs[:, 4].astype(int), labels)
+            ] = False
+            mask_labels[~np.isin(np.arange(n_labels), labels)] = False
 
         if label_keys is not None:
             if isinstance(label_keys, list):
                 label_keys = np.array(
                     [self.label_key_to_index[key] for key in label_keys]
                 )
-            label_indices = np.where(
-                np.isclose(self._label_metadata[:, 2], label_keys)
-            )[0]
-            mask = np.zeros_like(mask_pairs, dtype=np.bool_)
-            mask[
-                np.isin(self._ranked_pairs[:, 4].astype(int), label_indices)
-            ] = True
-            mask_pairs &= mask
-
-            mask = np.zeros_like(mask_labels, dtype=np.bool_)
-            mask[label_indices] = True
-            mask_labels &= mask
+            label_indices = (
+                np.where(np.isclose(self._label_metadata[:, 2], label_keys))[0]
+                if label_keys.size > 0
+                else np.array([])
+            )
+            mask_pairs[
+                ~np.isin(self._ranked_pairs[:, 4].astype(int), label_indices)
+            ] = False
+            mask_labels[~np.isin(np.arange(n_labels), label_indices)] = False
 
         mask = mask_datums[:, np.newaxis] & mask_labels[np.newaxis, :]
         label_metadata_per_datum = self._label_metadata_per_datum.copy()
@@ -224,8 +322,10 @@ class Evaluator:
 
     def evaluate(
         self,
+        metrics_to_return: list[MetricType] = MetricType.base_metrics(),
         iou_thresholds: list[float] = [0.5, 0.75, 0.9],
         score_thresholds: list[float] = [0.5],
+        number_of_examples: int = 0,
         filter_: Filter | None = None,
     ) -> dict[MetricType, list]:
         """
@@ -233,10 +333,14 @@ class Evaluator:
 
         Parameters
         ----------
+        metrics_to_return : list[MetricType]
+            A list of metrics to return in the results.
         iou_thresholds : list[float]
             A list of IoU thresholds to compute metrics over.
         score_thresholds : list[float]
             A list of score thresholds to compute metrics over.
+        number_of_examples : int, default=0
+            Number of annotation examples to return in DetailedCounts.
         filter_ : Filter, optional
             An optional filter object.
 
@@ -284,7 +388,7 @@ class Evaluator:
             )
             for iou_idx in range(average_precision.shape[0])
             for label_idx in range(average_precision.shape[1])
-            if int(label_metadata[label_idx][0]) > 0
+            if int(label_metadata[label_idx, 0]) > 0
         ]
 
         metrics[MetricType.mAP] = [
@@ -304,7 +408,7 @@ class Evaluator:
                 label=self.index_to_label[label_idx],
             )
             for label_idx in range(self.n_labels)
-            if int(label_metadata[label_idx][0]) > 0
+            if int(label_metadata[label_idx, 0]) > 0
         ]
 
         metrics[MetricType.mAPAveragedOverIOUs] = [
@@ -327,7 +431,7 @@ class Evaluator:
             )
             for score_idx in range(average_recall.shape[0])
             for label_idx in range(average_recall.shape[1])
-            if int(label_metadata[label_idx][0]) > 0
+            if int(label_metadata[label_idx, 0]) > 0
         ]
 
         metrics[MetricType.mAR] = [
@@ -349,7 +453,7 @@ class Evaluator:
                 label=self.index_to_label[label_idx],
             )
             for label_idx in range(self.n_labels)
-            if int(label_metadata[label_idx][0]) > 0
+            if int(label_metadata[label_idx, 0]) > 0
         ]
 
         metrics[MetricType.mARAveragedOverScores] = [
@@ -372,16 +476,17 @@ class Evaluator:
             )
             for iou_idx, iou_threshold in enumerate(iou_thresholds)
             for label_idx, label in self.index_to_label.items()
-            if int(label_metadata[label_idx][0]) > 0
+            if int(label_metadata[label_idx, 0]) > 0
         ]
 
         for label_idx, label in self.index_to_label.items():
+
+            if label_metadata[label_idx, 0] == 0:
+                continue
+
             for score_idx, score_threshold in enumerate(score_thresholds):
                 for iou_idx, iou_threshold in enumerate(iou_thresholds):
 
-                    if label_metadata[label_idx, 0] == 0:
-                        continue
-
                     row = precision_recall[iou_idx][score_idx][label_idx]
                     kwargs = {
                         "label": label,
@@ -422,16 +527,27 @@ class Evaluator:
                         )
                     )
 
+        if MetricType.DetailedCounts in metrics_to_return:
+            metrics[MetricType.DetailedCounts] = self._compute_detailed_counts(
+                iou_thresholds=iou_thresholds,
+                score_thresholds=score_thresholds,
+                n_samples=number_of_examples,
+            )
+
+        for metric in set(metrics.keys()):
+            if metric not in metrics_to_return:
+                del metrics[metric]
+
         return metrics
 
-    def compute_detailed_counts(
+    def _compute_detailed_counts(
         self,
         iou_thresholds: list[float] = [0.5],
         score_thresholds: list[float] = [
             score / 10.0 for score in range(1, 11)
         ],
         n_samples: int = 0,
-    ) -> list[list[DetailedCounts]]:
+    ) -> list[DetailedCounts]:
         """
         Computes detailed counting metrics.
 
@@ -454,7 +570,7 @@ class Evaluator:
             return list()
 
         metrics = compute_detailed_counts(
-            self._detailed_pairs,
+            data=self._detailed_pairs,
             label_metadata=self._label_metadata,
             iou_thresholds=np.array(iou_thresholds),
             score_thresholds=np.array(score_thresholds),
@@ -462,95 +578,111 @@ class Evaluator:
         )
 
         tp_idx = 0
-        fp_misclf_idx = tp_idx + n_samples + 1
-        fp_halluc_idx = fp_misclf_idx + n_samples + 1
-        fn_misclf_idx = fp_halluc_idx + n_samples + 1
-        fn_misprd_idx = fn_misclf_idx + n_samples + 1
+        fp_misclf_idx = 2 * n_samples + 1
+        fp_halluc_idx = 4 * n_samples + 2
+        fn_misclf_idx = 6 * n_samples + 3
+        fn_misprd_idx = 8 * n_samples + 4
+
+        def _unpack_examples(
+            iou_idx: int,
+            label_idx: int,
+            type_idx: int,
+            example_source: dict[int, NDArray[np.float16]],
+        ) -> list[list[tuple[str, tuple[float, float, float, float]]]]:
+            """
+            Unpacks metric examples from computation.
+            """
+            type_idx += 1
+
+            results = list()
+            for score_idx in range(n_scores):
+                examples = list()
+                for example_idx in range(n_samples):
+                    datum_idx = metrics[
+                        iou_idx,
+                        score_idx,
+                        label_idx,
+                        type_idx + example_idx * 2,
+                    ]
+                    annotation_idx = metrics[
+                        iou_idx,
+                        score_idx,
+                        label_idx,
+                        type_idx + example_idx * 2 + 1,
+                    ]
+                    if datum_idx >= 0:
+                        examples.append(
+                            (
+                                self.index_to_uid[datum_idx],
+                                tuple(
+                                    example_source[datum_idx][
+                                        annotation_idx
+                                    ].tolist()
+                                ),
+                            )
+                        )
+                results.append(examples)
+
+            return results
 
         n_ious, n_scores, n_labels, _ = metrics.shape
         return [
-            [
-                DetailedCounts(
-                    iou_threshold=iou_thresholds[iou_idx],
-                    label=self.index_to_label[label_idx],
-                    score_thresholds=score_thresholds,
-                    tp=metrics[iou_idx, :, label_idx, tp_idx]
-                    .astype(int)
-                    .tolist(),
-                    tp_examples=[
-                        [
-                            self.index_to_uid[int(datum_idx)]
-                            for datum_idx in metrics[iou_idx][score_idx][
-                                label_idx
-                            ][tp_idx + 1 : fp_misclf_idx]
-                            if int(datum_idx) >= 0
-                        ]
-                        for score_idx in range(n_scores)
-                    ],
-                    fp_misclassification=metrics[
-                        iou_idx, :, label_idx, fp_misclf_idx
-                    ]
-                    .astype(int)
-                    .tolist(),
-                    fp_misclassification_examples=[
-                        [
-                            self.index_to_uid[int(datum_idx)]
-                            for datum_idx in metrics[iou_idx][score_idx][
-                                label_idx
-                            ][fp_misclf_idx + 1 : fp_halluc_idx]
-                            if int(datum_idx) >= 0
-                        ]
-                        for score_idx in range(n_scores)
-                    ],
-                    fp_hallucination=metrics[
-                        iou_idx, :, label_idx, fp_halluc_idx
-                    ]
-                    .astype(int)
-                    .tolist(),
-                    fp_hallucination_examples=[
-                        [
-                            self.index_to_uid[int(datum_idx)]
-                            for datum_idx in metrics[iou_idx][score_idx][
-                                label_idx
-                            ][fp_halluc_idx + 1 : fn_misclf_idx]
-                            if int(datum_idx) >= 0
-                        ]
-                        for score_idx in range(n_scores)
-                    ],
-                    fn_misclassification=metrics[
-                        iou_idx, :, label_idx, fn_misclf_idx
-                    ]
-                    .astype(int)
-                    .tolist(),
-                    fn_misclassification_examples=[
-                        [
-                            self.index_to_uid[int(datum_idx)]
-                            for datum_idx in metrics[iou_idx][score_idx][
-                                label_idx
-                            ][fn_misclf_idx + 1 : fn_misprd_idx]
-                            if int(datum_idx) >= 0
-                        ]
-                        for score_idx in range(n_scores)
-                    ],
-                    fn_missing_prediction=metrics[
-                        iou_idx, :, label_idx, fn_misprd_idx
-                    ]
-                    .astype(int)
-                    .tolist(),
-                    fn_missing_prediction_examples=[
-                        [
-                            self.index_to_uid[int(datum_idx)]
-                            for datum_idx in metrics[iou_idx][score_idx][
-                                label_idx
-                            ][fn_misprd_idx + 1 :]
-                            if int(datum_idx) >= 0
-                        ]
-                        for score_idx in range(n_scores)
-                    ],
-                )
-                for iou_idx in range(n_ious)
-            ]
+            DetailedCounts(
+                iou_threshold=iou_thresholds[iou_idx],
+                label=self.index_to_label[label_idx],
+                score_thresholds=score_thresholds,
+                tp=metrics[iou_idx, :, label_idx, tp_idx].astype(int).tolist(),
+                fp_misclassification=metrics[
+                    iou_idx, :, label_idx, fp_misclf_idx
+                ]
+                .astype(int)
+                .tolist(),
+                fp_hallucination=metrics[iou_idx, :, label_idx, fp_halluc_idx]
+                .astype(int)
+                .tolist(),
+                fn_misclassification=metrics[
+                    iou_idx, :, label_idx, fn_misclf_idx
+                ]
+                .astype(int)
+                .tolist(),
+                fn_missing_prediction=metrics[
+                    iou_idx, :, label_idx, fn_misprd_idx
+                ]
+                .astype(int)
+                .tolist(),
+                tp_examples=_unpack_examples(
+                    iou_idx=iou_idx,
+                    label_idx=label_idx,
+                    type_idx=tp_idx,
+                    example_source=self.prediction_examples,
+                ),
+                fp_misclassification_examples=_unpack_examples(
+                    iou_idx=iou_idx,
+                    label_idx=label_idx,
+                    type_idx=fp_misclf_idx,
+                    example_source=self.prediction_examples,
+                ),
+                fp_hallucination_examples=_unpack_examples(
+                    iou_idx=iou_idx,
+                    label_idx=label_idx,
+                    type_idx=fp_halluc_idx,
+                    example_source=self.prediction_examples,
+                ),
+                fn_misclassification_examples=_unpack_examples(
+                    iou_idx=iou_idx,
+                    label_idx=label_idx,
+                    type_idx=fn_misclf_idx,
+                    example_source=self.groundtruth_examples,
+                ),
+                fn_missing_prediction_examples=_unpack_examples(
+                    iou_idx=iou_idx,
+                    label_idx=label_idx,
+                    type_idx=fn_misprd_idx,
+                    example_source=self.groundtruth_examples,
+                ),
+            )
             for label_idx in range(n_labels)
+            for iou_idx in range(n_ious)
         ]
 
 
@@ -561,7 +693,7 @@ class DataLoader:
 
     def __init__(self):
         self._evaluator = Evaluator()
-        self.pairs = list()
+        self.pairs: list[NDArray[np.floating]] = list()
         self.groundtruth_count = defaultdict(lambda: defaultdict(int))
         self.prediction_count = defaultdict(lambda: defaultdict(int))
 
@@ -624,9 +756,143 @@ class DataLoader:
             self._evaluator.label_key_to_index[label[0]],
         )
 
-    def add_data(
+    def _compute_ious_and_cache_pairs(
+        self,
+        uid_index: int,
+        keyed_groundtruths: dict,
+        keyed_predictions: dict,
+        annotation_type: type[BoundingBox] | type[Polygon] | type[Bitmask],
+    ) -> None:
+        """
+        Compute IOUs between groundtruths and preditions before storing as pairs.
+
+        Parameters
+        ----------
+        uid_index: int
+            The index of the detection.
+        keyed_groundtruths: dict
+            A dictionary of groundtruths.
+        keyed_predictions: dict
+            A dictionary of predictions.
+        annotation_type: type[BoundingBox] | type[Polygon] | type[Bitmask]
+            The type of annotation to compute IOUs for.
+        """
+        gt_keys = set(keyed_groundtruths.keys())
+        pd_keys = set(keyed_predictions.keys())
+        joint_keys = gt_keys.intersection(pd_keys)
+        gt_unique_keys = gt_keys - pd_keys
+        pd_unique_keys = pd_keys - gt_keys
+
+        pairs = list()
+        for key in joint_keys:
+            n_predictions = len(keyed_predictions[key])
+            n_groundtruths = len(keyed_groundtruths[key])
+            data = _get_annotation_data(
+                keyed_groundtruths=keyed_groundtruths,
+                keyed_predictions=keyed_predictions,
+                key=key,
+                annotation_type=annotation_type,
+            )
+            ious = compute_iou(data=data, annotation_type=annotation_type)
+            mask_nonzero_iou = (ious > 1e-9).reshape(
+                (n_predictions, n_groundtruths)
+            )
+            mask_ious_halluc = ~(mask_nonzero_iou.any(axis=1))
+            mask_ious_misprd = ~(mask_nonzero_iou.any(axis=0))
+
+            pairs.extend(
+                [
+                    np.array(
+                        [
+                            float(uid_index),
+                            float(gidx),
+                            float(pidx),
+                            ious[pidx * len(keyed_groundtruths[key]) + gidx],
+                            float(glabel),
+                            float(plabel),
+                            float(score),
+                        ]
+                    )
+                    for pidx, plabel, score, _ in keyed_predictions[key]
+                    for gidx, glabel, _ in keyed_groundtruths[key]
+                    if ious[pidx * len(keyed_groundtruths[key]) + gidx] > 1e-9
+                ]
+            )
+            pairs.extend(
+                [
+                    np.array(
+                        [
+                            float(uid_index),
+                            -1.0,
+                            float(pidx),
+                            0.0,
+                            -1.0,
+                            float(plabel),
+                            float(score),
+                        ]
+                    )
+                    for pidx, plabel, score, _ in keyed_predictions[key]
+                    if mask_ious_halluc[pidx]
+                ]
+            )
+            pairs.extend(
+                [
+                    np.array(
+                        [
+                            float(uid_index),
+                            float(gidx),
+                            -1.0,
+                            0.0,
+                            float(glabel),
+                            -1.0,
+                            -1.0,
+                        ]
+                    )
+                    for gidx, glabel, _ in keyed_groundtruths[key]
+                    if mask_ious_misprd[gidx]
+                ]
+            )
+        for key in gt_unique_keys:
+            pairs.extend(
+                [
+                    np.array(
+                        [
+                            float(uid_index),
+                            float(gidx),
+                            -1.0,
+                            0.0,
+                            float(glabel),
+                            -1.0,
+                            -1.0,
+                        ]
+                    )
+                    for gidx, glabel, _ in keyed_groundtruths[key]
+                ]
+            )
+        for key in pd_unique_keys:
+            pairs.extend(
+                [
+                    np.array(
+                        [
+                            float(uid_index),
+                            -1.0,
+                            float(pidx),
+                            0.0,
+                            -1.0,
+                            float(plabel),
+                            float(score),
+                        ]
+                    )
+                    for pidx, plabel, score, _ in keyed_predictions[key]
+                ]
+            )
+
+        self.pairs.append(np.array(pairs))
+
+    def _add_data(
         self,
         detections: list[Detection],
+        annotation_type: type[Bitmask] | type[BoundingBox] | type[Polygon],
         show_progress: bool = False,
     ):
         """
@@ -636,6 +902,8 @@ class DataLoader:
         ----------
         detections : list[Detection]
             A list of Detection objects.
+        annotation_type : type[Bitmask] | type[BoundingBox] | type[Polygon]
+            The annotation type to process.
         show_progress : bool, default=False
             Toggle for tqdm progress bar.
         """
@@ -650,108 +918,157 @@ class DataLoader:
             # update datum uid index
             uid_index = self._add_datum(uid=detection.uid)
 
+            # initialize bounding box examples
+            self._evaluator.groundtruth_examples[uid_index] = np.zeros(
+                (len(detection.groundtruths), 4), dtype=np.float16
+            )
+            self._evaluator.prediction_examples[uid_index] = np.zeros(
+                (len(detection.predictions), 4), dtype=np.float16
+            )
+
             # cache labels and annotations
             keyed_groundtruths = defaultdict(list)
             keyed_predictions = defaultdict(list)
+
+            representation_property = _get_annotation_representation(
+                annotation_type=annotation_type
+            )
+
             for gidx, gann in enumerate(detection.groundtruths):
+                if not isinstance(gann, annotation_type):
+                    raise ValueError(
+                        f"Expected {annotation_type}, but annotation is of type {type(gann)}."
+                    )
+
+                if isinstance(gann, BoundingBox):
+                    self._evaluator.groundtruth_examples[uid_index][
+                        gidx
+                    ] = getattr(gann, representation_property)
+                else:
+                    converted_box = gann.to_box()
+                    self._evaluator.groundtruth_examples[uid_index][gidx] = (
+                        getattr(converted_box, "extrema")
+                        if converted_box is not None
+                        else None
+                    )
                 for glabel in gann.labels:
                     label_idx, label_key_idx = self._add_label(glabel)
                     self.groundtruth_count[label_idx][uid_index] += 1
+                    representation = getattr(gann, representation_property)
                     keyed_groundtruths[label_key_idx].append(
                         (
                             gidx,
                             label_idx,
-                            gann.extrema,
+                            representation,
                         )
                     )
+
             for pidx, pann in enumerate(detection.predictions):
+                if not isinstance(pann, annotation_type):
+                    raise ValueError(
+                        f"Expected {annotation_type}, but annotation is of type {type(pann)}."
+                    )
+
+                if isinstance(pann, BoundingBox):
+                    self._evaluator.prediction_examples[uid_index][
+                        pidx
+                    ] = getattr(pann, representation_property)
+                else:
+                    converted_box = pann.to_box()
+                    self._evaluator.prediction_examples[uid_index][pidx] = (
+                        getattr(converted_box, "extrema")
+                        if converted_box is not None
+                        else None
+                    )
                 for plabel, pscore in zip(pann.labels, pann.scores):
                     label_idx, label_key_idx = self._add_label(plabel)
                     self.prediction_count[label_idx][uid_index] += 1
+                    representation = representation = getattr(
+                        pann, representation_property
+                    )
                     keyed_predictions[label_key_idx].append(
                         (
                             pidx,
                             label_idx,
                             pscore,
-                            pann.extrema,
+                            representation,
                         )
                     )
 
-            gt_keys = set(keyed_groundtruths.keys())
-            pd_keys = set(keyed_predictions.keys())
-            joint_keys = gt_keys.intersection(pd_keys)
-            gt_unique_keys = gt_keys - pd_keys
-            pd_unique_keys = pd_keys - gt_keys
-
-            pairs = list()
-            for key in joint_keys:
-                boxes = np.array(
-                    [
-                        np.array([*gextrema, *pextrema])
-                        for _, _, _, pextrema in keyed_predictions[key]
-                        for _, _, gextrema in keyed_groundtruths[key]
-                    ]
-                )
-                ious = compute_iou(boxes)
-                pairs.extend(
-                    [
-                        np.array(
-                            [
-                                float(uid_index),
-                                float(gidx),
-                                float(pidx),
-                                ious[
-                                    pidx * len(keyed_groundtruths[key]) + gidx
-                                ],
-                                float(glabel),
-                                float(plabel),
-                                float(score),
-                            ]
-                        )
-                        for pidx, plabel, score, _ in keyed_predictions[key]
-                        for gidx, glabel, _ in keyed_groundtruths[key]
-                    ]
-                )
-            for key in gt_unique_keys:
-                pairs.extend(
-                    [
-                        np.array(
-                            [
-                                float(uid_index),
-                                float(gidx),
-                                -1.0,
-                                0.0,
-                                float(glabel),
-                                -1.0,
-                                -1.0,
-                            ]
-                        )
-                        for gidx, glabel, _ in keyed_groundtruths[key]
-                    ]
-                )
-            for key in pd_unique_keys:
-                pairs.extend(
-                    [
-                        np.array(
-                            [
-                                float(uid_index),
-                                -1.0,
-                                float(pidx),
-                                0.0,
-                                -1.0,
-                                float(plabel),
-                                float(score),
-                            ]
-                        )
-                        for pidx, plabel, score, _ in keyed_predictions[key]
-                    ]
-                )
+            self._compute_ious_and_cache_pairs(
+                uid_index=uid_index,
+                keyed_groundtruths=keyed_groundtruths,
+                keyed_predictions=keyed_predictions,
+                annotation_type=annotation_type,
+            )
+
+    def add_bounding_boxes(
+        self,
+        detections: list[Detection],
+        show_progress: bool = False,
+    ):
+        """
+        Adds bounding box detections to the cache.
+
+        Parameters
+        ----------
+        detections : list[Detection]
+            A list of Detection objects.
+        show_progress : bool, default=False
+            Toggle for tqdm progress bar.
+        """
+        return self._add_data(
+            detections=detections,
+            show_progress=show_progress,
+            annotation_type=BoundingBox,
+        )
+
+    def add_polygons(
+        self,
+        detections: list[Detection],
+        show_progress: bool = False,
+    ):
+        """
+        Adds polygon detections to the cache.
+
+        Parameters
+        ----------
+        detections : list[Detection]
+            A list of Detection objects.
+        show_progress : bool, default=False
+            Toggle for tqdm progress bar.
+        """
+        return self._add_data(
+            detections=detections,
+            show_progress=show_progress,
+            annotation_type=Polygon,
+        )
+
+    def add_bitmasks(
+        self,
+        detections: list[Detection],
+        show_progress: bool = False,
+    ):
+        """
+        Adds bitmask detections to the cache.
 
-            self.pairs.append(np.array(pairs))
+        Parameters
+        ----------
+        detections : list[Detection]
+            A list of Detection objects.
+        show_progress : bool, default=False
+            Toggle for tqdm progress bar.
+        """
+        return self._add_data(
+            detections=detections,
+            show_progress=show_progress,
+            annotation_type=Bitmask,
+        )
 
-    def add_data_from_valor_dict(
+    def _add_data_from_valor_dict(
         self,
         detections: list[tuple[dict, dict]],
+        annotation_type: type[Bitmask] | type[BoundingBox] | type[Polygon],
         show_progress: bool = False,
     ):
         """
@@ -761,20 +1078,14 @@ class DataLoader:
         ----------
         detections : list[tuple[dict, dict]]
             A list of groundtruth, prediction pairs in Valor-format dictionaries.
+        annotation_type : type[Bitmask] | type[BoundingBox] | type[Polygon]
+            The annotation type to process.
         show_progress : bool, default=False
             Toggle for tqdm progress bar.
         """
 
-        def _get_bbox_extrema(
-            data: list[list[list[float]]],
-        ) -> tuple[float, float, float, float]:
-            x = [point[0] for shape in data for point in shape]
-            y = [point[1] for shape in data for point in shape]
-            return (min(x), max(x), min(y), max(y))
-
         disable_tqdm = not show_progress
         for groundtruth, prediction in tqdm(detections, disable=disable_tqdm):
-
             # update metadata
             self._evaluator.n_datums += 1
             self._evaluator.n_groundtruths += len(groundtruth["annotations"])
@@ -783,10 +1094,45 @@ class DataLoader:
             # update datum uid index
             uid_index = self._add_datum(uid=groundtruth["datum"]["uid"])
 
+            # initialize bounding box examples
+            self._evaluator.groundtruth_examples[uid_index] = np.zeros(
+                (len(groundtruth["annotations"]), 4), dtype=np.float16
+            )
+            self._evaluator.prediction_examples[uid_index] = np.zeros(
+                (len(prediction["annotations"]), 4), dtype=np.float16
+            )
+
             # cache labels and annotations
             keyed_groundtruths = defaultdict(list)
             keyed_predictions = defaultdict(list)
+
+            annotation_key = _get_valor_dict_annotation_key(
+                annotation_type=annotation_type
+            )
+            invalid_keys = list(
+                filter(
+                    lambda x: x != annotation_key,
+                    ["bounding_box", "raster", "polygon"],
+                )
+            )
+
             for gidx, gann in enumerate(groundtruth["annotations"]):
+                if (gann[annotation_key] is None) or any(
+                    [gann[k] is not None for k in invalid_keys]
+                ):
+                    raise ValueError(
+                        f"Input JSON doesn't contain {annotation_type} data, or contains data for multiple annotation types."
+                    )
+                if annotation_type == BoundingBox:
+                    self._evaluator.groundtruth_examples[uid_index][
+                        gidx
+                    ] = np.array(
+                        _get_annotation_representation_from_valor_dict(
+                            gann[annotation_key],
+                            annotation_type=annotation_type,
+                        ),
+                    )
+
                 for valor_label in gann["labels"]:
                     glabel = (valor_label["key"], valor_label["value"])
                     label_idx, label_key_idx = self._add_label(glabel)
@@ -795,10 +1141,29 @@ class DataLoader:
                         (
                             gidx,
                             label_idx,
-                            _get_bbox_extrema(gann["bounding_box"]),
+                            _get_annotation_representation_from_valor_dict(
+                                gann[annotation_key],
+                                annotation_type=annotation_type,
+                            ),
                         )
                     )
             for pidx, pann in enumerate(prediction["annotations"]):
+                if (pann[annotation_key] is None) or any(
+                    [pann[k] is not None for k in invalid_keys]
+                ):
+                    raise ValueError(
+                        f"Input JSON doesn't contain {annotation_type} data, or contains data for multiple annotation types."
+                    )
+
+                if annotation_type == BoundingBox:
+                    self._evaluator.prediction_examples[uid_index][
+                        pidx
+                    ] = np.array(
+                        _get_annotation_representation_from_valor_dict(
+                            pann[annotation_key],
+                            annotation_type=annotation_type,
+                        )
+                    )
                 for valor_label in pann["labels"]:
                     plabel = (valor_label["key"], valor_label["value"])
                     pscore = valor_label["score"]
@@ -809,81 +1174,40 @@ class DataLoader:
                             pidx,
                             label_idx,
                             pscore,
-                            _get_bbox_extrema(pann["bounding_box"]),
+                            _get_annotation_representation_from_valor_dict(
+                                pann[annotation_key],
+                                annotation_type=annotation_type,
+                            ),
                         )
                     )
 
-            gt_keys = set(keyed_groundtruths.keys())
-            pd_keys = set(keyed_predictions.keys())
-            joint_keys = gt_keys.intersection(pd_keys)
-            gt_unique_keys = gt_keys - pd_keys
-            pd_unique_keys = pd_keys - gt_keys
-
-            pairs = list()
-            for key in joint_keys:
-                boxes = np.array(
-                    [
-                        np.array([*gextrema, *pextrema])
-                        for _, _, _, pextrema in keyed_predictions[key]
-                        for _, _, gextrema in keyed_groundtruths[key]
-                    ]
-                )
-                ious = compute_iou(boxes)
-                pairs.extend(
-                    [
-                        np.array(
-                            [
-                                float(uid_index),
-                                float(gidx),
-                                float(pidx),
-                                ious[
-                                    pidx * len(keyed_groundtruths[key]) + gidx
-                                ],
-                                float(glabel),
-                                float(plabel),
-                                float(score),
-                            ]
-                        )
-                        for pidx, plabel, score, _ in keyed_predictions[key]
-                        for gidx, glabel, _ in keyed_groundtruths[key]
-                    ]
-                )
-            for key in gt_unique_keys:
-                pairs.extend(
-                    [
-                        np.array(
-                            [
-                                float(uid_index),
-                                float(gidx),
-                                -1.0,
-                                0.0,
-                                float(glabel),
-                                -1.0,
-                                -1.0,
-                            ]
-                        )
-                        for gidx, glabel, _ in keyed_groundtruths[key]
-                    ]
-                )
-            for key in pd_unique_keys:
-                pairs.extend(
-                    [
-                        np.array(
-                            [
-                                float(uid_index),
-                                -1.0,
-                                float(pidx),
-                                0.0,
-                                -1.0,
-                                float(plabel),
-                                float(score),
-                            ]
-                        )
-                        for pidx, plabel, score, _ in keyed_predictions[key]
-                    ]
-                )
+            self._compute_ious_and_cache_pairs(
+                uid_index=uid_index,
+                keyed_groundtruths=keyed_groundtruths,
+                keyed_predictions=keyed_predictions,
+                annotation_type=annotation_type,
+            )
 
-            self.pairs.append(np.array(pairs))
+    def add_bounding_boxes_from_valor_dict(
+        self,
+        detections: list[tuple[dict, dict]],
+        show_progress: bool = False,
+    ):
+        """
+        Adds Valor-format bounding box detections to the cache.
+
+        Parameters
+        ----------
+        detections : list[tuple[dict, dict]]
+            A list of groundtruth, prediction pairs in Valor-format dictionaries.
+        show_progress : bool, default=False
+            Toggle for tqdm progress bar.
+        """
+        return self._add_data_from_valor_dict(
+            detections=detections,
+            show_progress=show_progress,
+            annotation_type=BoundingBox,
+        )
 
     def finalize(self) -> Evaluator:
         """