Refactored metrics store, removed typing extensions

* This commit is in a messy state - tests were not adapted to new metrics.core
* MeanAveragePrecision is incomplete

Author: LinasKo

poetry.lock

@@ -59,7 +59,6 @@ tqdm = { version = ">=4.62.3,<=4.66.5", optional = true }
 # pandas: picked lowest major version that supports Python 3.8
 pandas = { version = ">=2.0.0", optional = true }
 pandas-stubs = { version = ">=2.0.0.230412", optional = true }
-typing-extensions = "^4.12.2"
 
 [tool.poetry.extras]
 desktop = ["opencv-python"]

pyproject.toml

@@ -4,4 +4,11 @@
     MetricTarget,
     UnsupportedMetricTargetError,
 )
-from supervision.metrics.intersection_over_union import IntersectionOverUnion
+from supervision.metrics.intersection_over_union import (
+    IntersectionOverUnion,
+    IntersectionOverUnionResult,
+)
+from supervision.metrics.mean_average_precision import (
+    MeanAveragePrecision,
+    MeanAveragePrecisionResult,
+)

supervision/metrics/__init__.py

@@ -2,17 +2,17 @@
 
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Any, Dict, Iterator, Tuple, Union
+from typing import Any, Iterator, Set, Tuple
 
 import numpy as np
 import numpy.typing as npt
-from typing_extensions import Self
 
 from supervision import config
 from supervision.detection.core import Detections
-from supervision.metrics.utils import len0_like, pad_mask
+from supervision.metrics.utils import pad_mask
 
 CLASS_ID_NONE = -1
+CONFIDENCE_NONE = -1
 """Used by metrics module as class ID, when none is present"""
 
 
@@ -22,7 +22,7 @@ class Metric(ABC):
     """
 
     @abstractmethod
-    def update(self, *args, **kwargs) -> Self:
+    def update(self, *args, **kwargs) -> "Metric":
         """
         Add data to the metric, without computing the result.
         Return the metric itself to allow method chaining.
@@ -78,171 +78,176 @@ def __init__(self, metric: Metric, target: MetricTarget):
         super().__init__(f"Metric {metric} does not support target {target}")
 
 
-class InternalMetricDataStore:
+class MetricData:
     """
-    Stores internal data of IntersectionOverUnion metric:
-    * Stores the basic data: boxes, masks, or oriented bounding boxes
-    * Validates data: ensures data types and shape are consistent
-    * Provides iteration by class
-
-    Provides a class-agnostic mode, where all data is treated as a single class.
-    Warning: numpy inputs are always considered as class-agnostic data.
-
-    Data here refers to content of Detections objects: boxes, masks,
-    or oriented bounding boxes.
+    A container for detection contents, decouple from Detections.
+    While a np.ndarray work for xyxy and obb, this approach solves
+    the mask concatenation problem.
     """
 
-    def __init__(self, metric_target: MetricTarget, class_agnostic: bool):
+    def __init__(self, metric_target: MetricTarget, class_agnostic: bool = False):
         self._metric_target = metric_target
         self._class_agnostic = class_agnostic
-        self._data_1: Dict[int, npt.NDArray]
-        self._data_2: Dict[int, npt.NDArray]
-        self._mask_shape: Tuple[int, int]
-        self.reset()
+        self.confidence = np.array([], dtype=np.float32)
+        self.class_id = np.array([], dtype=int)
+        self.data: npt.NDArray = self._get_empty_data()
 
-    def reset(self) -> None:
-        self._data_1 = {}
-        self._data_2 = {}
-        self._mask_shape = (0, 0)
-
-    def update(
-        self,
-        data_1: Union[npt.NDArray, Detections],
-        data_2: Union[npt.NDArray, Detections],
-    ) -> None:
-        """
-        Add new data to the store.
+    def update(self, detections: Detections):
+        """Add new detections to the store."""
+        new_data = self._get_content(detections)
+        self._validate_shape(new_data)
 
-        Use sv.Detections.empty() if only one set of data is available.
-        """
-        content_1 = self._get_content(data_1)
-        content_2 = self._get_content(data_2)
-        self._validate_shape(content_1)
-        self._validate_shape(content_2)
+        if self._metric_target == MetricTarget.BOXES:
+            self._append_boxes(new_data)
+        elif self._metric_target == MetricTarget.MASKS:
+            self._append_mask(new_data)
+        elif self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
+            self.data = np.vstack((self.data, new_data))
 
-        class_ids_1 = self._get_class_ids(data_1)
-        class_ids_2 = self._get_class_ids(data_2)
-        self._validate_class_ids(class_ids_1, class_ids_2)
+        confidence = self._get_confidence(detections)
+        self._append_confidence(confidence)
 
-        if self._metric_target == MetricTarget.MASKS:
-            content_1 = self._expand_mask_shape(content_1)
-            content_2 = self._expand_mask_shape(content_2)
-
-        for class_id in set(class_ids_1):
-            content_of_class = content_1[class_ids_1 == class_id]
-            stored_content_of_class = self._data_1.get(class_id, len0_like(content_1))
-            self._data_1[class_id] = np.vstack(
-                (stored_content_of_class, content_of_class)
-            )
+        class_id = self._get_class_id(detections)
+        self._append_class_id(class_id)
 
-        for class_id in set(class_ids_2):
-            content_of_class = content_2[class_ids_2 == class_id]
-            stored_content_of_class = self._data_2.get(class_id, len0_like(content_2))
-            self._data_2[class_id] = np.vstack(
-                (stored_content_of_class, content_of_class)
+        if len(self.class_id) != len(self.confidence) or len(self.class_id) != len(
+            self.data
+        ):
+            raise ValueError(
+                f"Inconsistent data length: class_id={len(class_id)},"
+                f" confidence={len(confidence)}, data={len(new_data)}"
             )
 
-    def __getitem__(self, class_id: int) -> Tuple[npt.NDArray, npt.NDArray]:
-        return (
-            self._data_1.get(class_id, self._make_empty()),
-            self._data_2.get(class_id, self._make_empty()),
-        )
+    def get_classes(self) -> Set[int]:
+        """Return all class IDs."""
+        return set(self.class_id)
 
-    def __iter__(
-        self,
-    ) -> Iterator[Tuple[int, npt.NDArray, npt.NDArray]]:
-        class_ids = sorted(set(self._data_1.keys()) | set(self._data_2.keys()))
-        for class_id in class_ids:
-            yield (
-                class_id,
-                *self[class_id],
-            )
+    def get_subset_by_class(self, class_id: int) -> MetricData:
+        """Return data, confidence and class_id for a specific class."""
+        mask = self.class_id == class_id
+        new_data_obj = MetricData(self._metric_target)
+        new_data_obj.data = self.data[mask]
+        new_data_obj.confidence = self.confidence[mask]
+        new_data_obj.class_id = self.class_id[mask]
+        return new_data_obj
 
-    def _get_content(self, data: Union[npt.NDArray, Detections]) -> npt.NDArray:
-        """Return boxes, masks or oriented bounding boxes from the data."""
-        if not isinstance(data, (Detections, np.ndarray)):
-            raise ValueError(
-                f"Invalid data type: {type(data)}."
-                f" Only Detections or np.ndarray are supported."
-            )
-        if isinstance(data, np.ndarray):
-            return data
+    def __len__(self) -> int:
+        return len(self.data)
 
+    def _get_content(self, detections: Detections) -> npt.NDArray:
+        """Return boxes, masks or oriented bounding boxes from the data."""
         if self._metric_target == MetricTarget.BOXES:
-            return data.xyxy
+            return detections.xyxy
         if self._metric_target == MetricTarget.MASKS:
             return (
-                data.mask if data.mask is not None else np.zeros((0, 0, 0), dtype=bool)
+                detections.mask
+                if detections.mask is not None
+                else self._get_empty_data()
             )
         if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
-            obb = data.data.get(
-                config.ORIENTED_BOX_COORDINATES, np.zeros((0, 8), dtype=np.float32)
+            obb = detections.data.get(
+                config.ORIENTED_BOX_COORDINATES, self._get_empty_data()
             )
-            return np.array(obb, dtype=np.float32)
+            return np.ndarray(obb, dtype=np.float32)
         raise ValueError(f"Invalid metric target: {self._metric_target}")
 
-    def _get_class_ids(
-        self, data: Union[npt.NDArray, Detections]
-    ) -> npt.NDArray[np.int_]:
-        """
-        Return an array of class IDs from the data. Guaranteed to
-        match the length of data.
-        """
-        if (
-            self._class_agnostic
-            or isinstance(data, np.ndarray)
-            or data.class_id is None
-        ):
-            return np.array([CLASS_ID_NONE] * len(data), dtype=int)
-        return data.class_id
-
-    def _validate_class_ids(
-        self, class_id_1: npt.NDArray[np.int_], class_id_2: npt.NDArray[np.int_]
-    ) -> None:
-        class_set = set(class_id_1) | set(class_id_2)
-        if len(class_set) >= 2 and CLASS_ID_NONE in class_set:
-            raise ValueError(
-                "Metrics cannot mix data with class ID and data without class ID."
-            )
+    def _get_class_id(self, detections: Detections) -> npt.NDArray[np.int_]:
+        if self._class_agnostic or detections.class_id is None:
+            return np.array([CLASS_ID_NONE] * len(detections), dtype=int)
+        return detections.class_id
+
+    def _get_confidence(self, detections: Detections) -> npt.NDArray[np.float32]:
+        if detections.confidence is None:
+            return np.full(len(detections), -1, dtype=np.float32)
+        return detections.confidence
+
+    def _append_class_id(self, new_class_id: npt.NDArray[np.int_]) -> None:
+        self.class_id = np.hstack((self.class_id, new_class_id))
+
+    def _append_confidence(self, new_confidence: npt.NDArray[np.float32]) -> None:
+        self.confidence = np.hstack((self.confidence, new_confidence))
+
+    def _append_boxes(self, new_boxes: npt.NDArray[np.float32]) -> None:
+        """Stack new xyxy or obb boxes on top of stored boxes."""
+        if self._metric_target not in [
+            MetricTarget.BOXES,
+            MetricTarget.ORIENTED_BOUNDING_BOXES,
+        ]:
+            raise ValueError("This method is only for box data.")
+        self.data = np.vstack((self.data, new_boxes))
+
+    def _append_mask(self, new_mask: npt.NDArray[np.bool_]) -> None:
+        """Stack new mask onto stored masks. Expand the shapes if necessary."""
+        if self._metric_target != MetricTarget.MASKS:
+            raise ValueError("This method is only for mask data.")
+        self._validate_mask_shape(new_mask)
+
+        new_width = max(self.data.shape[1], new_mask.shape[1])
+        new_height = max(self.data.shape[2], new_mask.shape[2])
+
+        data = pad_mask(self.data, (new_width, new_height))
+        new_mask = pad_mask(new_mask, (new_width, new_height))
+
+        self.data = np.vstack((data, new_mask))
+
+    def _get_empty_data(self) -> npt.NDArray:
+        if self._metric_target == MetricTarget.BOXES:
+            return np.empty((0, 4), dtype=np.float32)
+        if self._metric_target == MetricTarget.MASKS:
+            return np.empty((0, 0, 0), dtype=bool)
+        if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
+            return np.empty((0, 8), dtype=np.float32)
+        raise ValueError(f"Invalid metric target: {self._metric_target}")
 
     def _validate_shape(self, data: npt.NDArray) -> None:
-        shape = data.shape
         if self._metric_target == MetricTarget.BOXES:
-            if len(shape) != 2 or shape[1] != 4:
-                raise ValueError(f"Invalid xyxy shape: {shape}. Expected: (N, 4)")
+            if len(data.shape) != 2 or data.shape[1] != 4:
+                raise ValueError(f"Invalid xyxy shape: {data.shape}. Expected: (N, 4)")
         elif self._metric_target == MetricTarget.MASKS:
-            if len(shape) != 3:
-                raise ValueError(f"Invalid mask shape: {shape}. Expected: (N, H, W)")
+            if len(data.shape) != 3:
+                raise ValueError(
+                    f"Invalid mask shape: {data.shape}. Expected: (N, H, W)"
+                )
         elif self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
-            if len(shape) != 2 or shape[1] != 8:
-                raise ValueError(f"Invalid obb shape: {shape}. Expected: (N, 8)")
+            if len(data.shape) != 2 or data.shape[1] != 8:
+                raise ValueError(f"Invalid obb shape: {data.shape}. Expected: (N, 8)")
         else:
             raise ValueError(f"Invalid metric target: {self._metric_target}")
 
-    def _expand_mask_shape(self, data: npt.NDArray) -> npt.NDArray:
-        """Pad the stored and new data to the same shape."""
-        if self._metric_target != MetricTarget.MASKS:
-            return data
 
-        new_width = max(self._mask_shape[0], data.shape[1])
-        new_height = max(self._mask_shape[1], data.shape[2])
-        self._mask_shape = (new_width, new_height)
+class InternalMetricDataStore:
+    """
+    Stores internal data for metrics.
 
-        data = pad_mask(data, self._mask_shape)
+    Provides a class-agnostic way to access it.
+    """
 
-        for class_id, prev_data in self._data_1.items():
-            self._data_1[class_id] = pad_mask(prev_data, self._mask_shape)
-        for class_id, prev_data in self._data_2.items():
-            self._data_2[class_id] = pad_mask(prev_data, self._mask_shape)
+    def __init__(self, metric_target: MetricTarget, class_agnostic: bool = False):
+        self._metric_target = metric_target
+        self._class_agnostic = class_agnostic
+        self._data_1: MetricData
+        self._data_2: MetricData
+        self.reset()
 
-        return data
+    def reset(self) -> None:
+        self._data_1 = MetricData(self._metric_target, self._class_agnostic)
+        self._data_2 = MetricData(self._metric_target, self._class_agnostic)
 
-    def _make_empty(self) -> npt.NDArray:
-        """Create an empty data object with the best-known shape for the target."""
-        if self._metric_target == MetricTarget.BOXES:
-            return np.empty((0, 4), dtype=np.float32)
-        if self._metric_target == MetricTarget.MASKS:
-            return np.empty((0, *self._mask_shape), dtype=bool)
-        if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
-            return np.empty((0, 8), dtype=np.float32)
-        raise ValueError(f"Invalid metric target: {self._metric_target}")
+    def update(self, data_1: Detections, data_2: Detections) -> None:
+        """
+        Add new data to the store.
+
+        Use sv.Detections.empty() if only one set of data is available.
+        """
+        self._data_1.update(data_1)
+        self._data_2.update(data_2)
+
+    def __getitem__(self, class_id: int) -> Tuple[MetricData, MetricData]:
+        return (
+            self._data_1.get_subset_by_class(class_id),
+            self._data_2.get_subset_by_class(class_id),
+        )
+
+    def __iter__(self) -> Iterator[Tuple[int, MetricData, MetricData]]:
+        for class_id in self._data_1.get_classes():
+            yield class_id, *self[class_id]