diff --git a/supervision/dataset/formats/pascal_voc.py b/supervision/dataset/formats/pascal_voc.py
index 74f6be70..ba77ddf8 100644
--- a/supervision/dataset/formats/pascal_voc.py
+++ b/supervision/dataset/formats/pascal_voc.py
@@ -268,12 +268,12 @@ def detections_from_xml_obj(
[extended_classes.index(class_name) for class_name in class_names]
)
- if with_masks:
- annotation = Detections(
- xyxy=xyxy, mask=np.array(masks).astype(bool), class_id=class_id
- )
- else:
- annotation = Detections(xyxy=xyxy, class_id=class_id)
+ annotation = Detections(
+ xyxy=xyxy.astype(np.float32),
+ mask=np.array(masks).astype(bool) if with_masks else None,
+ class_id=class_id,
+ )
+
return annotation, extended_classes
diff --git a/supervision/detection/core.py b/supervision/detection/core.py
index ec5d9f72..201b4e9a 100644
--- a/supervision/detection/core.py
+++ b/supervision/detection/core.py
@@ -1,69 +1,28 @@
from __future__ import annotations
-from dataclasses import astuple, dataclass
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+from dataclasses import dataclass, field
+from typing import Dict, Iterator, List, Optional, Tuple, Union
import numpy as np
from supervision.detection.utils import (
calculate_masks_centroids,
extract_ultralytics_masks,
+ get_data_item,
+ is_data_equal,
+ merge_data,
non_max_suppression,
process_roboflow_result,
+ validate_detections_fields,
xywh_to_xyxy,
)
from supervision.geometry.core import Position
-def _validate_xyxy(xyxy: Any, n: int) -> None:
- is_valid = isinstance(xyxy, np.ndarray) and xyxy.shape == (n, 4)
- if not is_valid:
- raise ValueError("xyxy must be 2d np.ndarray with (n, 4) shape")
-
-
-def _validate_mask(mask: Any, n: int) -> None:
- is_valid = mask is None or (
- isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n
- )
- if not is_valid:
- raise ValueError("mask must be 3d np.ndarray with (n, H, W) shape")
-
-
-def validate_inference_callback(callback) -> None:
- tmp_img = np.zeros((256, 256, 3), dtype=np.uint8)
- res = callback(tmp_img)
- if not isinstance(res, Detections):
- raise ValueError("Callback function must return sv.Detection type")
-
-
-def _validate_class_id(class_id: Any, n: int) -> None:
- is_valid = class_id is None or (
- isinstance(class_id, np.ndarray) and class_id.shape == (n,)
- )
- if not is_valid:
- raise ValueError("class_id must be None or 1d np.ndarray with (n,) shape")
-
-
-def _validate_confidence(confidence: Any, n: int) -> None:
- is_valid = confidence is None or (
- isinstance(confidence, np.ndarray) and confidence.shape == (n,)
- )
- if not is_valid:
- raise ValueError("confidence must be None or 1d np.ndarray with (n,) shape")
-
-
-def _validate_tracker_id(tracker_id: Any, n: int) -> None:
- is_valid = tracker_id is None or (
- isinstance(tracker_id, np.ndarray) and tracker_id.shape == (n,)
- )
- if not is_valid:
- raise ValueError("tracker_id must be None or 1d np.ndarray with (n,) shape")
-
-
@dataclass
class Detections:
"""
- Data class containing information about the detections in a video frame.
+ A dataclass representing detection results.
Attributes:
xyxy (np.ndarray): An array of shape `(n, 4)` containing
@@ -76,6 +35,17 @@ class Detections:
`(n,)` containing the class ids of the detections.
tracker_id (Optional[np.ndarray]): An array of shape
`(n,)` containing the tracker ids of the detections.
+ data (Dict[str, Union[np.ndarray, List]]): A dictionary containing additional
+ data where each key is a string representing the data type, and the value
+ is either a NumPy array or a list of corresponding data.
+
+ !!! warning
+
+ The `data` field in the `sv.Detections` class is currently in an experimental
+ phase. Please be aware that its API and functionality are subject to change in
+ future updates as we continue to refine and improve its capabilities.
+ We encourage users to experiment with this feature and provide feedback, but
+ also to be prepared for potential modifications in upcoming releases.
"""
xyxy: np.ndarray
@@ -83,14 +53,17 @@ class Detections:
confidence: Optional[np.ndarray] = None
class_id: Optional[np.ndarray] = None
tracker_id: Optional[np.ndarray] = None
+ data: Dict[str, Union[np.ndarray, List]] = field(default_factory=dict)
def __post_init__(self):
- n = len(self.xyxy)
- _validate_xyxy(xyxy=self.xyxy, n=n)
- _validate_mask(mask=self.mask, n=n)
- _validate_class_id(class_id=self.class_id, n=n)
- _validate_confidence(confidence=self.confidence, n=n)
- _validate_tracker_id(tracker_id=self.tracker_id, n=n)
+ validate_detections_fields(
+ xyxy=self.xyxy,
+ mask=self.mask,
+ confidence=self.confidence,
+ class_id=self.class_id,
+ tracker_id=self.tracker_id,
+ data=self.data,
+ )
def __len__(self):
"""
@@ -107,11 +80,12 @@ class Detections:
Optional[float],
Optional[int],
Optional[int],
+ Dict[str, Union[np.ndarray, List]],
]
]:
"""
Iterates over the Detections object and yield a tuple of
- `(xyxy, mask, confidence, class_id, tracker_id)` for each detection.
+ `(xyxy, mask, confidence, class_id, tracker_id, data)` for each detection.
"""
for i in range(len(self.xyxy)):
yield (
@@ -120,36 +94,18 @@ class Detections:
self.confidence[i] if self.confidence is not None else None,
self.class_id[i] if self.class_id is not None else None,
self.tracker_id[i] if self.tracker_id is not None else None,
+ get_data_item(self.data, i),
)
def __eq__(self, other: Detections):
return all(
[
np.array_equal(self.xyxy, other.xyxy),
- any(
- [
- self.mask is None and other.mask is None,
- np.array_equal(self.mask, other.mask),
- ]
- ),
- any(
- [
- self.class_id is None and other.class_id is None,
- np.array_equal(self.class_id, other.class_id),
- ]
- ),
- any(
- [
- self.confidence is None and other.confidence is None,
- np.array_equal(self.confidence, other.confidence),
- ]
- ),
- any(
- [
- self.tracker_id is None and other.tracker_id is None,
- np.array_equal(self.tracker_id, other.tracker_id),
- ]
- ),
+ np.array_equal(self.mask, other.mask),
+ np.array_equal(self.class_id, other.class_id),
+ np.array_equal(self.confidence, other.confidence),
+ np.array_equal(self.tracker_id, other.tracker_id),
+ is_data_equal(self.data, other.data),
]
)
@@ -701,30 +657,67 @@ class Detections:
Example:
```python
- >>> from supervision import Detections
+ import numpy as np
+ import supervision as sv
- >>> detections_1 = Detections(...)
- >>> detections_2 = Detections(...)
+ >>> detections_1 = sv.Detections(
+ ... xyxy=np.array([[15, 15, 100, 100], [200, 200, 300, 300]]),
+ ... class_id=np.array([1, 2]),
+ ... data={'feature_vector': np.array([0.1, 0.2)])}
+ ... )
+
+ >>> detections_2 = sv.Detections(
+ ... xyxy=np.array([[30, 30, 120, 120]]),
+ ... class_id=np.array([1]),
+ ... data={'feature_vector': [np.array([0.3])]}
+ ... )
>>> merged_detections = Detections.merge([detections_1, detections_2])
+
+ >>> merged_detections.xyxy
+ array([[ 15, 15, 100, 100],
+ [200, 200, 300, 300],
+ [ 30, 30, 120, 120]])
+
+ >>> merged_detections.class_id
+ array([1, 2, 1])
+
+ >>> merged_detections.data['feature_vector']
+ array([0.1, 0.2, 0.3])
```
"""
if len(detections_list) == 0:
return Detections.empty()
- detections_tuples_list = [astuple(detection) for detection in detections_list]
- xyxy, mask, confidence, class_id, tracker_id = [
- list(field) for field in zip(*detections_tuples_list)
- ]
+ for detections in detections_list:
+ validate_detections_fields(
+ xyxy=detections.xyxy,
+ mask=detections.mask,
+ confidence=detections.confidence,
+ class_id=detections.class_id,
+ tracker_id=detections.tracker_id,
+ data=detections.data,
+ )
- def __all_not_none(item_list: List[Any]):
- return all(x is not None for x in item_list)
+ xyxy = np.vstack([d.xyxy for d in detections_list])
- xyxy = np.vstack(xyxy)
- mask = np.vstack(mask) if __all_not_none(mask) else None
- confidence = np.hstack(confidence) if __all_not_none(confidence) else None
- class_id = np.hstack(class_id) if __all_not_none(class_id) else None
- tracker_id = np.hstack(tracker_id) if __all_not_none(tracker_id) else None
+ def stack_or_none(name: str):
+ if all(d.__getattribute__(name) is None for d in detections_list):
+ return None
+ if any(d.__getattribute__(name) is None for d in detections_list):
+ raise ValueError(f"All or none of the '{name}' fields must be None")
+ return (
+ np.vstack([d.__getattribute__(name) for d in detections_list])
+ if name == "mask"
+ else np.hstack([d.__getattribute__(name) for d in detections_list])
+ )
+
+ mask = stack_or_none("mask")
+ confidence = stack_or_none("confidence")
+ class_id = stack_or_none("class_id")
+ tracker_id = stack_or_none("tracker_id")
+
+ data = merge_data([d.data for d in detections_list])
return cls(
xyxy=xyxy,
@@ -732,6 +725,7 @@ class Detections:
confidence=confidence,
class_id=class_id,
tracker_id=tracker_id,
+ data=data,
)
def get_anchors_coordinates(self, anchor: Position) -> np.ndarray:
@@ -801,17 +795,23 @@ class Detections:
raise ValueError(f"{anchor} is not supported.")
def __getitem__(
- self, index: Union[int, slice, List[int], np.ndarray]
- ) -> Detections:
+ self, index: Union[int, slice, List[int], np.ndarray, str]
+ ) -> Union[Detections, List, np.ndarray, None]:
"""
- Get a subset of the Detections object.
+ Get a subset of the Detections object or access an item from its data field.
+
+ When provided with an integer, slice, list of integers, or a numpy array, this
+ method returns a new Detections object that represents a subset of the original
+ detections. When provided with a string, it accesses the corresponding item in
+ the data dictionary.
Args:
- index (Union[int, slice, List[int], np.ndarray]):
- The index or indices of the subset of the Detections
+ index (Union[int, slice, List[int], np.ndarray, str]): The index, indices,
+ or key to access a subset of the Detections or an item from the data.
Returns:
- (Detections): A subset of the Detections object.
+ Union[Detections, Any]: A subset of the Detections object or an item from
+ the data field.
Example:
```python
@@ -820,16 +820,16 @@ class Detections:
>>> detections = sv.Detections(...)
>>> first_detection = detections[0]
-
>>> first_10_detections = detections[0:10]
-
>>> some_detections = detections[[0, 2, 4]]
-
>>> class_0_detections = detections[detections.class_id == 0]
-
>>> high_confidence_detections = detections[detections.confidence > 0.5]
+
+ >>> feature_vector = detections['feature_vector']
```
"""
+ if isinstance(index, str):
+ return self.data.get(index)
if isinstance(index, int):
index = [index]
return Detections(
@@ -838,8 +838,45 @@ class Detections:
confidence=self.confidence[index] if self.confidence is not None else None,
class_id=self.class_id[index] if self.class_id is not None else None,
tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
+ data=get_data_item(self.data, index),
)
+ def __setitem__(self, key: str, value: Union[np.ndarray, List]):
+ """
+ Set a value in the data dictionary of the Detections object.
+
+ Args:
+ key (str): The key in the data dictionary to set.
+ value (Union[np.ndarray, List]): The value to set for the key.
+
+ Example:
+ ```python
+ >>> import cv2
+ >>> from ultralytics import YOLO
+ >>> import supervision as sv
+
+ >>> model = YOLO('yolov8s.pt')
+
+ >>> image = cv2.imread(SOURCE_IMAGE_PATH)
+
+ >>> result = model(image)[0]
+ >>> detections = sv.Detections.from_ultralytics(result)
+
+ >>> detections['names'] = [
+ ... model.model.names[class_id]
+ ... for class_id
+ ... in detections.class_id
+ ... ]
+ ```
+ """
+ if not isinstance(value, (np.ndarray, list)):
+ raise TypeError("Value must be a np.ndarray or a list")
+
+ if isinstance(value, list):
+ value = np.array(value)
+
+ self.data[key] = value
+
@property
def area(self) -> np.ndarray:
"""
diff --git a/supervision/detection/tools/inference_slicer.py b/supervision/detection/tools/inference_slicer.py
index 5f6fb391..2c0c48a6 100644
--- a/supervision/detection/tools/inference_slicer.py
+++ b/supervision/detection/tools/inference_slicer.py
@@ -3,7 +3,7 @@ from typing import Callable, Optional, Tuple
import numpy as np
-from supervision.detection.core import Detections, validate_inference_callback
+from supervision.detection.core import Detections
from supervision.detection.utils import move_boxes
from supervision.utils.image import crop_image
@@ -60,7 +60,6 @@ class InferenceSlicer:
self.iou_threshold = iou_threshold
self.callback = callback
self.thread_workers = thread_workers
- validate_inference_callback(callback=callback)
def __call__(self, image: np.ndarray) -> Detections:
"""
diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py
index 86226371..2df89a84 100644
--- a/supervision/detection/utils.py
+++ b/supervision/detection/utils.py
@@ -1,4 +1,5 @@
-from typing import List, Optional, Tuple
+from itertools import chain
+from typing import Any, Dict, List, Optional, Tuple, Union
import cv2
import numpy as np
@@ -299,7 +300,6 @@ def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]:
orig_shape = yolov8_results.orig_shape
inference_shape = tuple(yolov8_results.masks.data.shape[1:])
- gain = 0
pad = (0, 0)
if inference_shape != orig_shape:
@@ -470,3 +470,204 @@ def calculate_masks_centroids(masks: np.ndarray) -> np.ndarray:
centroid_y = sum_over_mask(vertical_indices, aggregation_axis) / total_pixels
return np.column_stack((centroid_x, centroid_y)).astype(int)
+
+
+def validate_xyxy(xyxy: Any) -> None:
+ expected_shape = "(_, 4)"
+ actual_shape = str(getattr(xyxy, "shape", None))
+ is_valid = isinstance(xyxy, np.ndarray) and xyxy.ndim == 2 and xyxy.shape[1] == 4
+ if not is_valid:
+ raise ValueError(
+ f"xyxy must be a 2D np.ndarray with shape {expected_shape}, but got shape "
+ f"{actual_shape}"
+ )
+
+
+def validate_mask(mask: Any, n: int) -> None:
+ expected_shape = f"({n}, H, W)"
+ actual_shape = str(getattr(mask, "shape", None))
+ is_valid = mask is None or (
+ isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n
+ )
+ if not is_valid:
+ raise ValueError(
+ f"mask must be a 3D np.ndarray with shape {expected_shape}, but got shape "
+ f"{actual_shape}"
+ )
+
+
+def validate_class_id(class_id: Any, n: int) -> None:
+ expected_shape = f"({n},)"
+ actual_shape = str(getattr(class_id, "shape", None))
+ is_valid = class_id is None or (
+ isinstance(class_id, np.ndarray) and class_id.shape == (n,)
+ )
+ if not is_valid:
+ raise ValueError(
+ f"class_id must be a 1D np.ndarray with shape {expected_shape}, but got "
+ f"shape {actual_shape}"
+ )
+
+
+def validate_confidence(confidence: Any, n: int) -> None:
+ expected_shape = f"({n},)"
+ actual_shape = str(getattr(confidence, "shape", None))
+ is_valid = confidence is None or (
+ isinstance(confidence, np.ndarray) and confidence.shape == (n,)
+ )
+ if not is_valid:
+ raise ValueError(
+ f"confidence must be a 1D np.ndarray with shape {expected_shape}, but got "
+ f"shape {actual_shape}"
+ )
+
+
+def validate_tracker_id(tracker_id: Any, n: int) -> None:
+ expected_shape = f"({n},)"
+ actual_shape = str(getattr(tracker_id, "shape", None))
+ is_valid = tracker_id is None or (
+ isinstance(tracker_id, np.ndarray) and tracker_id.shape == (n,)
+ )
+ if not is_valid:
+ raise ValueError(
+ f"tracker_id must be a 1D np.ndarray with shape {expected_shape}, but got "
+ f"shape {actual_shape}"
+ )
+
+
+def validate_data(data: Dict[str, Any], n: int) -> None:
+ for key, value in data.items():
+ if isinstance(value, list):
+ if len(value) != n:
+ raise ValueError(f"Length of list for key '{key}' must be {n}")
+ elif isinstance(value, np.ndarray):
+ if value.ndim == 1 and value.shape[0] != n:
+ raise ValueError(f"Shape of np.ndarray for key '{key}' must be ({n},)")
+ elif value.ndim > 1 and value.shape[0] != n:
+ raise ValueError(
+ f"First dimension of np.ndarray for key '{key}' must have size {n}"
+ )
+ else:
+ raise ValueError(f"Value for key '{key}' must be a list or np.ndarray")
+
+
+def validate_detections_fields(
+ xyxy: Any,
+ mask: Any,
+ class_id: Any,
+ confidence: Any,
+ tracker_id: Any,
+ data: Dict[str, Any],
+) -> None:
+ validate_xyxy(xyxy)
+ n = len(xyxy)
+ validate_mask(mask, n)
+ validate_class_id(class_id, n)
+ validate_confidence(confidence, n)
+ validate_tracker_id(tracker_id, n)
+ validate_data(data, n)
+
+
+def is_data_equal(data_a: Dict[str, np.ndarray], data_b: Dict[str, np.ndarray]) -> bool:
+ """
+ Compares the data payloads of two Detections instances.
+
+ Args:
+ data_a, data_b: The data payloads of the instances.
+
+ Returns:
+ True if the data payloads are equal, False otherwise.
+ """
+ return set(data_a.keys()) == set(data_b.keys()) and all(
+ np.array_equal(data_a[key], data_b[key]) for key in data_a
+ )
+
+
+def merge_data(
+ data_list: List[Dict[str, Union[np.ndarray, List]]],
+) -> Dict[str, Union[np.ndarray, List]]:
+ """
+ Merges the data payloads of a list of Detections instances.
+
+ Args:
+ data_list: The data payloads of the instances.
+
+ Returns:
+ A single data payload containing the merged data, preserving the original data
+ types (list or np.ndarray).
+
+ Raises:
+ ValueError: If data values within a single object have different lengths or if
+ dictionaries have different keys.
+ """
+ if not data_list:
+ return {}
+
+ all_keys_sets = [set(data.keys()) for data in data_list]
+ if not all(keys_set == all_keys_sets[0] for keys_set in all_keys_sets):
+ raise ValueError("All data dictionaries must have the same keys to merge.")
+
+ for data in data_list:
+ lengths = [len(value) for value in data.values()]
+ if len(set(lengths)) > 1:
+ raise ValueError(
+ "All data values within a single object must have equal length."
+ )
+
+ merged_data = {key: [] for key in all_keys_sets[0]}
+
+ for data in data_list:
+ for key in merged_data:
+ merged_data[key].append(data[key])
+
+ for key in merged_data:
+ if all(isinstance(item, list) for item in merged_data[key]):
+ merged_data[key] = list(chain.from_iterable(merged_data[key]))
+ elif all(isinstance(item, np.ndarray) for item in merged_data[key]):
+ ndim = merged_data[key][0].ndim
+ if ndim == 1:
+ merged_data[key] = np.hstack(merged_data[key])
+ elif ndim > 1:
+ merged_data[key] = np.vstack(merged_data[key])
+ else:
+ raise ValueError(f"Unexpected array dimension for key '{key}'.")
+ else:
+ raise ValueError(
+ f"Inconsistent data types for key '{key}'. Only np.ndarray and list "
+ f"types are allowed."
+ )
+
+ return merged_data
+
+
+def get_data_item(
+ data: Dict[str, Union[np.ndarray, List]],
+ index: Union[int, slice, List[int], np.ndarray],
+) -> Dict[str, Union[np.ndarray, List]]:
+ """
+ Retrieve a subset of the data dictionary based on the given index.
+
+ Args:
+ data: The data dictionary of the Detections object.
+ index: The index or indices specifying the subset to retrieve.
+
+ Returns:
+ A subset of the data dictionary corresponding to the specified index.
+ """
+ subset_data = {}
+ for key, value in data.items():
+ if isinstance(value, np.ndarray):
+ subset_data[key] = value[index]
+ elif isinstance(value, list):
+ if isinstance(index, slice):
+ subset_data[key] = value[index]
+ elif isinstance(index, (list, np.ndarray)):
+ subset_data[key] = [value[i] for i in index]
+ elif isinstance(index, int):
+ subset_data[key] = [value[index]]
+ else:
+ raise TypeError(f"Unsupported index type: {type(index)}")
+ else:
+ raise TypeError(f"Unsupported data type for key '{key}': {type(value)}")
+
+ return subset_data
diff --git a/test/dataset/formats/test_pascal_voc.py b/test/dataset/formats/test_pascal_voc.py
index 269f41c7..bce29271 100644
--- a/test/dataset/formats/test_pascal_voc.py
+++ b/test/dataset/formats/test_pascal_voc.py
@@ -122,7 +122,7 @@ NO_DETECTIONS = """"""
["test"],
(100, 100),
False,
- mock_detections(np.array([[0, 0, 10, 10]]), None, [0]),
+ mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
DoesNotRaise(),
),
(
@@ -130,7 +130,9 @@ NO_DETECTIONS = """"""
["test"],
(100, 100),
False,
- mock_detections(np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), None, [0, 0]),
+ mock_detections(
+ xyxy=np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), class_id=[0, 0]
+ ),
DoesNotRaise(),
),
(
@@ -139,9 +141,8 @@ NO_DETECTIONS = """"""
(100, 100),
False,
mock_detections(
- np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]),
- None,
- [0, 0, 1],
+ xyxy=np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]),
+ class_id=[0, 0, 1],
),
DoesNotRaise(),
),
@@ -150,7 +151,7 @@ NO_DETECTIONS = """"""
[],
(100, 100),
False,
- mock_detections(np.empty((0, 4)), None, []),
+ mock_detections(xyxy=np.empty((0, 4)), class_id=[]),
DoesNotRaise(),
),
],
diff --git a/test/detection/test_core.py b/test/detection/test_core.py
index 0d17e09b..f3b739e8 100644
--- a/test/detection/test_core.py
+++ b/test/detection/test_core.py
@@ -153,7 +153,10 @@ def test_getitem(
DoesNotRaise(),
), # single detection with xyxy field
(
- [mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()],
+ [
+ mock_detections(xyxy=[[10, 10, 20, 20]]),
+ mock_detections(xyxy=np.empty((0, 4), dtype=np.float32)),
+ ],
mock_detections(xyxy=[[10, 10, 20, 20]]),
DoesNotRaise(),
), # single detection with xyxy field + empty detection
@@ -171,7 +174,7 @@ def test_getitem(
mock_detections(xyxy=[[20, 20, 30, 30]]),
],
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
- DoesNotRaise(),
+ pytest.raises(ValueError),
), # detection with xyxy, class_id fields + detection with xyxy field
(
[
@@ -181,6 +184,16 @@ def test_getitem(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[0, 1]),
DoesNotRaise(),
), # two detections with xyxy, class_id fields
+ (
+ [
+ mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
+ mock_detections(xyxy=[[20, 20, 30, 30]], data={"test": [2]}),
+ ],
+ mock_detections(
+ xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], data={"test": [1, 2]}
+ ),
+ DoesNotRaise(),
+ ), # two detections with xyxy, data fields
],
)
def test_merge(
@@ -273,3 +286,54 @@ def test_get_anchor_coordinates(
result = detections.get_anchors_coordinates(anchor)
with exception:
assert np.array_equal(result, expected_result)
+
+
+@pytest.mark.parametrize(
+ "detections_a, detections_b, expected_result",
+ [
+ (
+ Detections.empty(),
+ Detections.empty(),
+ True,
+ ), # empty detections
+ (
+ mock_detections(xyxy=[[10, 10, 20, 20]]),
+ mock_detections(xyxy=[[10, 10, 20, 20]]),
+ True,
+ ), # detections with xyxy field
+ (
+ mock_detections(xyxy=[[10, 10, 20, 20]], confidence=[0.5]),
+ mock_detections(xyxy=[[10, 10, 20, 20]], confidence=[0.5]),
+ True,
+ ), # detections with xyxy, confidence fields
+ (
+ mock_detections(xyxy=[[10, 10, 20, 20]], confidence=[0.5]),
+ mock_detections(xyxy=[[10, 10, 20, 20]]),
+ False,
+ ), # detection with xyxy field + detection with xyxy, confidence fields
+ (
+ mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
+ mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
+ True,
+ ), # detections with xyxy, data fields
+ (
+ mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
+ mock_detections(xyxy=[[10, 10, 20, 20]]),
+ False,
+ ), # detection with xyxy field + detection with xyxy, data fields
+ (
+ mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [1]}),
+ mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_2": [1]}),
+ False,
+ ), # detections with xyxy, and different data field names
+ (
+ mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [1]}),
+ mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [3]}),
+ False,
+ ), # detections with xyxy, and different data field values
+ ],
+)
+def test_equal(
+ detections_a: Detections, detections_b: Detections, expected_result: bool
+) -> None:
+ assert (detections_a == detections_b) == expected_result
diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py
index 83b71b85..a24119ce 100644
--- a/test/detection/test_utils.py
+++ b/test/detection/test_utils.py
@@ -1,5 +1,5 @@
from contextlib import ExitStack as DoesNotRaise
-from typing import List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import pytest
@@ -8,6 +8,8 @@ from supervision.detection.utils import (
calculate_masks_centroids,
clip_boxes,
filter_polygons_by_area,
+ get_data_item,
+ merge_data,
move_boxes,
non_max_suppression,
process_roboflow_result,
@@ -640,3 +642,274 @@ def test_calculate_masks_centroids(
with exception:
result = calculate_masks_centroids(masks=masks)
assert np.array_equal(result, expected_result)
+
+
+@pytest.mark.parametrize(
+ "data_list, expected_result, exception",
+ [
+ (
+ [],
+ {},
+ DoesNotRaise(),
+ ), # empty data list
+ (
+ [{}],
+ {},
+ DoesNotRaise(),
+ ), # single empty data dict
+ (
+ [{}, {}],
+ {},
+ DoesNotRaise(),
+ ), # two empty data dicts
+ (
+ [
+ {"test_1": []},
+ ],
+ {"test_1": []},
+ DoesNotRaise(),
+ ), # single data dict with a single field name and empty list values
+ (
+ [
+ {"test_1": np.array([])},
+ ],
+ {"test_1": np.array([])},
+ DoesNotRaise(),
+ ), # single data dict with a single field name and empty np.array values
+ (
+ [
+ {"test_1": [1, 2, 3]},
+ ],
+ {"test_1": [1, 2, 3]},
+ DoesNotRaise(),
+ ), # single data dict with a single field name and list values
+ (
+ [
+ {"test_1": []},
+ {"test_1": [3, 2, 1]},
+ ],
+ {"test_1": [3, 2, 1]},
+ DoesNotRaise(),
+ ), # two data dicts with the same field name and empty and list values
+ (
+ [
+ {"test_1": [1, 2, 3]},
+ {"test_1": [3, 2, 1]},
+ ],
+ {"test_1": [1, 2, 3, 3, 2, 1]},
+ DoesNotRaise(),
+ ), # two data dicts with the same field name and list values
+ (
+ [
+ {"test_1": [1, 2, 3]},
+ {"test_1": [3, 2, 1]},
+ {"test_1": [1, 2, 3]},
+ ],
+ {"test_1": [1, 2, 3, 3, 2, 1, 1, 2, 3]},
+ DoesNotRaise(),
+ ), # three data dicts with the same field name and list values
+ (
+ [
+ {"test_1": [1, 2, 3]},
+ {"test_2": [3, 2, 1]},
+ ],
+ None,
+ pytest.raises(ValueError),
+ ), # two data dicts with different field names
+ (
+ [
+ {"test_1": np.array([1, 2, 3])},
+ {"test_1": np.array([3, 2, 1])},
+ ],
+ {"test_1": np.array([1, 2, 3, 3, 2, 1])},
+ DoesNotRaise(),
+ ), # two data dicts with the same field name and np.array values as 1D arrays
+ (
+ [
+ {"test_1": np.array([[1, 2, 3]])},
+ {"test_1": np.array([[3, 2, 1]])},
+ ],
+ {"test_1": np.array([[1, 2, 3], [3, 2, 1]])},
+ DoesNotRaise(),
+ ), # two data dicts with the same field name and np.array values as 2D arrays
+ (
+ [
+ {"test_1": np.array([1, 2, 3]), "test_2": np.array(["a", "b", "c"])},
+ {"test_1": np.array([3, 2, 1]), "test_2": np.array(["c", "b", "a"])},
+ ],
+ {
+ "test_1": np.array([1, 2, 3, 3, 2, 1]),
+ "test_2": np.array(["a", "b", "c", "c", "b", "a"]),
+ },
+ DoesNotRaise(),
+ ), # two data dicts with the same field names and np.array values
+ (
+ [
+ {"test_1": [1, 2, 3], "test_2": np.array(["a", "b", "c"])},
+ {"test_1": [3, 2, 1], "test_2": np.array(["c", "b", "a"])},
+ ],
+ {
+ "test_1": [1, 2, 3, 3, 2, 1],
+ "test_2": np.array(["a", "b", "c", "c", "b", "a"]),
+ },
+ DoesNotRaise(),
+ ), # two data dicts with the same field names and mixed values
+ (
+ [
+ {"test_1": np.array([1, 2, 3])},
+ {"test_1": np.array([[3, 2, 1]])},
+ ],
+ None,
+ pytest.raises(ValueError),
+ ), # two data dicts with the same field name and 1D and 2D arrays values
+ (
+ [
+ {"test_1": np.array([1, 2, 3]), "test_2": np.array(["a", "b"])},
+ {"test_1": np.array([3, 2, 1]), "test_2": np.array(["c", "b", "a"])},
+ ],
+ None,
+ pytest.raises(ValueError),
+ ), # two data dicts with the same field name and different length arrays values
+ ],
+)
+def test_merge_data(
+ data_list: List[Dict[str, Any]],
+ expected_result: Optional[Dict[str, Any]],
+ exception: Exception,
+):
+ with exception:
+ result = merge_data(data_list=data_list)
+ for key in result:
+ if isinstance(result[key], np.ndarray):
+ assert np.array_equal(
+ result[key], expected_result[key]
+ ), f"Mismatch in arrays for key {key}"
+ else:
+ assert (
+ result[key] == expected_result[key]
+ ), f"Mismatch in non-array data for key {key}"
+
+
+@pytest.mark.parametrize(
+ "data, index, expected_result, exception",
+ [
+ ({}, 0, {}, DoesNotRaise()), # empty data dict
+ (
+ {
+ "test_1": [1, 2, 3],
+ },
+ 0,
+ {
+ "test_1": [1],
+ },
+ DoesNotRaise(),
+ ), # single data dict with a single field name and list values
+ (
+ {
+ "test_1": np.array([1, 2, 3]),
+ },
+ 0,
+ {
+ "test_1": np.array([1]),
+ },
+ DoesNotRaise(),
+ ), # single data dict with a single field name and np.array values as 1D arrays
+ (
+ {
+ "test_1": [1, 2, 3],
+ },
+ slice(0, 2),
+ {
+ "test_1": [1, 2],
+ },
+ DoesNotRaise(),
+ ), # single data dict with a single field name and list values
+ (
+ {
+ "test_1": np.array([1, 2, 3]),
+ },
+ slice(0, 2),
+ {
+ "test_1": np.array([1, 2]),
+ },
+ DoesNotRaise(),
+ ), # single data dict with a single field name and np.array values as 1D arrays
+ (
+ {
+ "test_1": [1, 2, 3],
+ },
+ -1,
+ {
+ "test_1": [3],
+ },
+ DoesNotRaise(),
+ ), # single data dict with a single field name and list values
+ (
+ {
+ "test_1": np.array([1, 2, 3]),
+ },
+ -1,
+ {
+ "test_1": np.array([3]),
+ },
+ DoesNotRaise(),
+ ), # single data dict with a single field name and np.array values as 1D arrays
+ (
+ {
+ "test_1": [1, 2, 3],
+ },
+ [0, 2],
+ {
+ "test_1": [1, 3],
+ },
+ DoesNotRaise(),
+ ), # single data dict with a single field name and list values
+ (
+ {
+ "test_1": np.array([1, 2, 3]),
+ },
+ [0, 2],
+ {
+ "test_1": np.array([1, 3]),
+ },
+ DoesNotRaise(),
+ ), # single data dict with a single field name and np.array values as 1D arrays
+ (
+ {
+ "test_1": [1, 2, 3],
+ },
+ np.array([0, 2]),
+ {
+ "test_1": [1, 3],
+ },
+ DoesNotRaise(),
+ ), # single data dict with a single field name and list values
+ (
+ {
+ "test_1": np.array([1, 2, 3]),
+ },
+ np.array([0, 2]),
+ {
+ "test_1": np.array([1, 3]),
+ },
+ DoesNotRaise(),
+ ),
+ ],
+)
+def test_get_data_item(
+ data: Dict[str, Any],
+ index: Any,
+ expected_result: Optional[Dict[str, Any]],
+ exception: Exception,
+):
+ with exception:
+ result = get_data_item(data=data, index=index)
+ for key in result:
+ if isinstance(result[key], np.ndarray):
+ assert np.array_equal(
+ result[key], expected_result[key]
+ ), f"Mismatch in arrays for key {key}"
+ else:
+ assert (
+ result[key] == expected_result[key]
+ ), f"Mismatch in non-array data for key {key}"
diff --git a/test/test_utils.py b/test/test_utils.py
index a1d6907e..b8ef9231 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import Any, Dict, List, Optional
import numpy as np
@@ -7,19 +7,26 @@ from supervision.detection.core import Detections
def mock_detections(
xyxy: List[List[float]],
- confidence: List[float] = None,
- class_id: List[int] = None,
- tracker_id: List[int] = None,
+ mask: Optional[List[np.ndarray]] = None,
+ confidence: Optional[List[float]] = None,
+ class_id: Optional[List[int]] = None,
+ tracker_id: Optional[List[int]] = None,
+ data: Optional[Dict[str, List[Any]]] = None,
) -> Detections:
+ def convert_data(data: Dict[str, List[Any]]):
+ return {k: np.array(v) for k, v in data.items()}
+
return Detections(
xyxy=np.array(xyxy, dtype=np.float32),
- confidence=confidence
- if confidence is None
- else np.array(confidence, dtype=np.float32),
- class_id=class_id if class_id is None else np.array(class_id, dtype=int),
- tracker_id=tracker_id
- if tracker_id is None
- else np.array(tracker_id, dtype=int),
+ mask=(mask if mask is None else np.array(mask, dtype=bool)),
+ confidence=(
+ confidence if confidence is None else np.array(confidence, dtype=np.float32)
+ ),
+ class_id=(class_id if class_id is None else np.array(class_id, dtype=int)),
+ tracker_id=(
+ tracker_id if tracker_id is None else np.array(tracker_id, dtype=int)
+ ),
+ data=convert_data(data) if data else {},
)