okooooooo/supervision
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #700 from roboflow/feature/extend-detections-with-data-field-to-support-custom-payload

Browse Source 
feature/extend-detections-with-data-field-to-support-custom-payload
This commit is contained in:
Piotr Skalski authored and GitHub committed 2024-01-02 18:04:39 +01:00
2 parents b300c59f78 e3d92a01c4
commit 8528a8a738
8 files changed
+715 -133
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

supervision/dataset/formats/pascal_voc.py
+6 -6
View File
@@ -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
supervision/detection/core.py
+140 -103
View File
@@ -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:
"""
supervision/detection/tools/inference_slicer.py
+1 -2
View File
@@ -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:
"""
supervision/detection/utils.py
+203 -2
View File
@@ -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
test/dataset/formats/test_pascal_voc.py
+7 -6
View File
@@ -122,7 +122,7 @@ NO_DETECTIONS = """<annotation></annotation>"""
["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 = """<annotation></annotation>"""
["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 = """<annotation></annotation>"""
(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 = """<annotation></annotation>"""
[],
(100, 100),
False,
mock_detections(np.empty((0, 4)), None, []),
mock_detections(xyxy=np.empty((0, 4)), class_id=[]),
DoesNotRaise(),
),
],
test/detection/test_core.py
+66 -2
View File
@@ -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
test/detection/test_utils.py
+274 -1
View File
@@ -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}"
test/test_utils.py
+18 -11
View File
@@ -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 {},
)