From f3f545d8938489aa9c6e94ae7a3dd0795c3ec46f Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 28 Dec 2023 15:58:51 +0100 Subject: [PATCH 01/21] simplifying `Detections.__eq__` method --- supervision/detection/core.py | 28 ++++------------------------ test/detection/test_core.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 24 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index ec5d9f72..339ff4c9 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -126,30 +126,10 @@ class 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), ] ) diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 0d17e09b..2f2c020e 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -273,3 +273,35 @@ 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]] + ), + False + ), # detection with xyxy field + detection with xyxy, confidence fields + ] +) +def test_equal( + detections_a: Detections, + detections_b: Detections, + expected_result: bool +) -> None: + assert (detections_a == detections_b) == expected_result From f2031d53d6097ddef75a434ea053b499de0556ad Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 28 Dec 2023 15:03:46 +0000 Subject: [PATCH 02/21] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto?= =?UTF-8?q?=20format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- test/detection/test_core.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 2f2c020e..d3e17a83 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -289,19 +289,13 @@ def test_get_anchor_coordinates( True, ), # detections with xyxy field ( - mock_detections( - xyxy=[[10, 10, 20, 20]], confidence=[0.5] - ), - mock_detections( - xyxy=[[10, 10, 20, 20]] - ), - False + 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 - ] + ], ) def test_equal( - detections_a: Detections, - detections_b: Detections, - expected_result: bool + detections_a: Detections, detections_b: Detections, expected_result: bool ) -> None: assert (detections_a == detections_b) == expected_result From ea6439b8fd4445c6e87d8337b5608d15f1402aae Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 28 Dec 2023 17:52:25 +0100 Subject: [PATCH 03/21] `Detections.__eq__` updated and tested --- supervision/detection/core.py | 23 ++++- test/dataset/formats/test_pascal_voc.py | 96 ++++++++--------- test/detection/test_core.py | 131 ++++++++++++++---------- test/test_utils.py | 42 ++++++-- 4 files changed, 179 insertions(+), 113 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 339ff4c9..97072dd6 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import astuple, dataclass +from dataclasses import astuple, dataclass, field from typing import Any, Dict, Iterator, List, Optional, Tuple, Union import numpy as np @@ -60,6 +60,25 @@ def _validate_tracker_id(tracker_id: Any, n: int) -> None: raise ValueError("tracker_id must be None or 1d np.ndarray with (n,) shape") +def _data_payload_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) + ) + + @dataclass class Detections: """ @@ -83,6 +102,7 @@ class Detections: confidence: Optional[np.ndarray] = None class_id: Optional[np.ndarray] = None tracker_id: Optional[np.ndarray] = None + data: Dict[str, np.ndarray] = field(default_factory=dict) def __post_init__(self): n = len(self.xyxy) @@ -130,6 +150,7 @@ class Detections: 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), + _data_payload_equal(self.data, other.data), ] ) diff --git a/test/dataset/formats/test_pascal_voc.py b/test/dataset/formats/test_pascal_voc.py index 269f41c7..8f113c43 100644 --- a/test/dataset/formats/test_pascal_voc.py +++ b/test/dataset/formats/test_pascal_voc.py @@ -114,51 +114,51 @@ N_CLASS_N_BBOX = """test1 NO_DETECTIONS = """""" -@pytest.mark.parametrize( - "xml_string, classes, resolution_wh, force_masks, expected_result, exception", - [ - ( - ONE_CLASS_ONE_BBOX, - ["test"], - (100, 100), - False, - mock_detections(np.array([[0, 0, 10, 10]]), None, [0]), - DoesNotRaise(), - ), - ( - ONE_CLASS_N_BBOX, - ["test"], - (100, 100), - False, - mock_detections(np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), None, [0, 0]), - DoesNotRaise(), - ), - ( - N_CLASS_N_BBOX, - ["test", "test2"], - (100, 100), - False, - mock_detections( - np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]), - None, - [0, 0, 1], - ), - DoesNotRaise(), - ), - ( - NO_DETECTIONS, - [], - (100, 100), - False, - mock_detections(np.empty((0, 4)), None, []), - DoesNotRaise(), - ), - ], -) -def test_detections_from_xml_obj( - xml_string, classes, resolution_wh, force_masks, expected_result, exception -): - with exception: - root = ET.fromstring(xml_string) - result, _ = detections_from_xml_obj(root, classes, resolution_wh, force_masks) - assert result == expected_result +# @pytest.mark.parametrize( +# "xml_string, classes, resolution_wh, force_masks, expected_result, exception", +# [ +# ( +# ONE_CLASS_ONE_BBOX, +# ["test"], +# (100, 100), +# False, +# mock_detections(np.array([[0, 0, 10, 10]]), None, [0]), +# DoesNotRaise(), +# ), +# ( +# ONE_CLASS_N_BBOX, +# ["test"], +# (100, 100), +# False, +# mock_detections(np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), None, [0, 0]), +# DoesNotRaise(), +# ), +# ( +# N_CLASS_N_BBOX, +# ["test", "test2"], +# (100, 100), +# False, +# mock_detections( +# np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]), +# None, +# [0, 0, 1], +# ), +# DoesNotRaise(), +# ), +# ( +# NO_DETECTIONS, +# [], +# (100, 100), +# False, +# mock_detections(np.empty((0, 4)), None, []), +# DoesNotRaise(), +# ), +# ], +# ) +# def test_detections_from_xml_obj( +# xml_string, classes, resolution_wh, force_masks, expected_result, exception +# ): +# with exception: +# root = ET.fromstring(xml_string) +# result, _ = detections_from_xml_obj(root, classes, resolution_wh, force_masks) +# assert result == expected_result diff --git a/test/detection/test_core.py b/test/detection/test_core.py index d3e17a83..6a383876 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -138,59 +138,59 @@ def test_getitem( assert result == expected_result -@pytest.mark.parametrize( - "detections_list, expected_result, exception", - [ - ([], Detections.empty(), DoesNotRaise()), # empty detections list - ( - [Detections.empty()], - Detections.empty(), - DoesNotRaise(), - ), # single empty detections - ( - [mock_detections(xyxy=[[10, 10, 20, 20]])], - mock_detections(xyxy=[[10, 10, 20, 20]]), - DoesNotRaise(), - ), # single detection with xyxy field - ( - [mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()], - mock_detections(xyxy=[[10, 10, 20, 20]]), - DoesNotRaise(), - ), # single detection with xyxy field + empty detection - ( - [ - mock_detections(xyxy=[[10, 10, 20, 20]]), - mock_detections(xyxy=[[20, 20, 30, 30]]), - ], - mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), - DoesNotRaise(), - ), # two detections with xyxy field - ( - [ - mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]), - mock_detections(xyxy=[[20, 20, 30, 30]]), - ], - mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), - DoesNotRaise(), - ), # detection with xyxy, class_id fields + detection with xyxy field - ( - [ - mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]), - mock_detections(xyxy=[[20, 20, 30, 30]], class_id=[1]), - ], - mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[0, 1]), - DoesNotRaise(), - ), # two detections with xyxy, class_id fields - ], -) -def test_merge( - detections_list: List[Detections], - expected_result: Optional[Detections], - exception: Exception, -) -> None: - with exception: - result = Detections.merge(detections_list=detections_list) - assert result == expected_result +# @pytest.mark.parametrize( +# "detections_list, expected_result, exception", +# [ +# ([], Detections.empty(), DoesNotRaise()), # empty detections list +# ( +# [Detections.empty()], +# Detections.empty(), +# DoesNotRaise(), +# ), # single empty detections +# ( +# [mock_detections(xyxy=[[10, 10, 20, 20]])], +# mock_detections(xyxy=[[10, 10, 20, 20]]), +# DoesNotRaise(), +# ), # single detection with xyxy field +# ( +# [mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()], +# mock_detections(xyxy=[[10, 10, 20, 20]]), +# DoesNotRaise(), +# ), # single detection with xyxy field + empty detection +# ( +# [ +# mock_detections(xyxy=[[10, 10, 20, 20]]), +# mock_detections(xyxy=[[20, 20, 30, 30]]), +# ], +# mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), +# DoesNotRaise(), +# ), # two detections with xyxy field +# ( +# [ +# mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]), +# mock_detections(xyxy=[[20, 20, 30, 30]]), +# ], +# mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), +# DoesNotRaise(), +# ), # detection with xyxy, class_id fields + detection with xyxy field +# ( +# [ +# mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]), +# mock_detections(xyxy=[[20, 20, 30, 30]], class_id=[1]), +# ], +# mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[0, 1]), +# DoesNotRaise(), +# ), # two detections with xyxy, class_id fields +# ], +# ) +# def test_merge( +# detections_list: List[Detections], +# expected_result: Optional[Detections], +# exception: Exception, +# ) -> None: +# with exception: +# result = Detections.merge(detections_list=detections_list) +# assert result == expected_result @pytest.mark.parametrize( @@ -288,11 +288,36 @@ def test_get_anchor_coordinates( 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, 2, 3]}), + mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1, 2, 3]}), + True, + ), # detections with xyxy, data fields + ( + mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1, 2, 3]}), + 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, 2, 3]}), + mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_2": [1, 2, 3]}), + False, + ), # detections with xyxy, and different data field names + ( + mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [1, 2, 3]}), + mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [3, 2, 1]}), + False, + ), # detections with xyxy, and different data field values ], ) def test_equal( diff --git a/test/test_utils.py b/test/test_utils.py index a1d6907e..d307387d 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Optional, Dict, Any import numpy as np @@ -7,19 +7,39 @@ 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 {} ) From b2d98fe457e07c0b2e5fae9f578ff93ba2f0620a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 28 Dec 2023 16:52:42 +0000 Subject: [PATCH 04/21] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto?= =?UTF-8?q?=20format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 8 +++----- test/dataset/formats/test_pascal_voc.py | 2 -- test/test_utils.py | 25 ++++++------------------- 3 files changed, 9 insertions(+), 26 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 97072dd6..abcd6c21 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -61,8 +61,7 @@ def _validate_tracker_id(tracker_id: Any, n: int) -> None: def _data_payload_equal( - data_a: Dict[str, np.ndarray], - data_b: Dict[str, np.ndarray] + data_a: Dict[str, np.ndarray], data_b: Dict[str, np.ndarray] ) -> bool: """ Compares the data payloads of two Detections instances. @@ -73,9 +72,8 @@ def _data_payload_equal( 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) + 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 ) diff --git a/test/dataset/formats/test_pascal_voc.py b/test/dataset/formats/test_pascal_voc.py index 8f113c43..238967dc 100644 --- a/test/dataset/formats/test_pascal_voc.py +++ b/test/dataset/formats/test_pascal_voc.py @@ -1,5 +1,4 @@ from contextlib import ExitStack as DoesNotRaise -from test.test_utils import mock_detections from typing import List, Optional import defusedxml.ElementTree as ET @@ -7,7 +6,6 @@ import numpy as np import pytest from supervision.dataset.formats.pascal_voc import ( - detections_from_xml_obj, object_to_pascal_voc, parse_polygon_points, ) diff --git a/test/test_utils.py b/test/test_utils.py index d307387d..b8ef9231 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Dict, Any +from typing import Any, Dict, List, Optional import numpy as np @@ -13,33 +13,20 @@ def mock_detections( 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), - mask=( - mask - if mask is None - else np.array(mask, dtype=bool) - ), + 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) + 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) + tracker_id if tracker_id is None else np.array(tracker_id, dtype=int) ), - data=convert_data(data) if data else {} + data=convert_data(data) if data else {}, ) From efed23f64e1db5ca4bd8b5b3829ce30cb7f9e49e Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 29 Dec 2023 15:05:07 +0100 Subject: [PATCH 05/21] `Detections.merge` updated and tested --- supervision/detection/core.py | 74 ++++++++++-- supervision/detection/utils.py | 1 - test/detection/test_core.py | 214 ++++++++++++++++++++++++--------- 3 files changed, 222 insertions(+), 67 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index abcd6c21..2a57eb82 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import astuple, dataclass, field +from itertools import chain from typing import Any, Dict, Iterator, List, Optional, Tuple, Union import numpy as np @@ -60,7 +61,7 @@ def _validate_tracker_id(tracker_id: Any, n: int) -> None: raise ValueError("tracker_id must be None or 1d np.ndarray with (n,) shape") -def _data_payload_equal( +def is_data_equal( data_a: Dict[str, np.ndarray], data_b: Dict[str, np.ndarray] ) -> bool: """ @@ -77,10 +78,56 @@ def _data_payload_equal( ) +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). + """ + 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.") + + 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 + + @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 @@ -93,6 +140,9 @@ 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. """ xyxy: np.ndarray @@ -100,7 +150,7 @@ class Detections: confidence: Optional[np.ndarray] = None class_id: Optional[np.ndarray] = None tracker_id: Optional[np.ndarray] = None - data: Dict[str, np.ndarray] = field(default_factory=dict) + data: Dict[str, Union[np.ndarray, List]] = field(default_factory=dict) def __post_init__(self): n = len(self.xyxy) @@ -148,7 +198,7 @@ class Detections: 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), - _data_payload_equal(self.data, other.data), + is_data_equal(self.data, other.data), ] ) @@ -712,18 +762,19 @@ class Detections: 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) + xyxy, mask, confidence, class_id, tracker_id, data = [ + list(field_values) for field_values in zip(*detections_tuples_list) ] - def __all_not_none(item_list: List[Any]): + def all_not_none(item_list: List[Any]): return all(x is not None for x in item_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 + 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 + data = merge_data(data) return cls( xyxy=xyxy, @@ -731,6 +782,7 @@ class Detections: confidence=confidence, class_id=class_id, tracker_id=tracker_id, + data=data, ) def get_anchors_coordinates(self, anchor: Position) -> np.ndarray: diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 86226371..a26f72ff 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -299,7 +299,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: diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 6a383876..41102d57 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -1,11 +1,11 @@ from contextlib import ExitStack as DoesNotRaise from test.test_utils import mock_detections -from typing import List, Optional, Union +from typing import List, Optional, Union, Dict, Any import numpy as np import pytest -from supervision.detection.core import Detections +from supervision.detection.core import Detections, merge_data from supervision.geometry.core import Position PREDICTIONS = np.array( @@ -138,59 +138,59 @@ def test_getitem( assert result == expected_result -# @pytest.mark.parametrize( -# "detections_list, expected_result, exception", -# [ -# ([], Detections.empty(), DoesNotRaise()), # empty detections list -# ( -# [Detections.empty()], -# Detections.empty(), -# DoesNotRaise(), -# ), # single empty detections -# ( -# [mock_detections(xyxy=[[10, 10, 20, 20]])], -# mock_detections(xyxy=[[10, 10, 20, 20]]), -# DoesNotRaise(), -# ), # single detection with xyxy field -# ( -# [mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()], -# mock_detections(xyxy=[[10, 10, 20, 20]]), -# DoesNotRaise(), -# ), # single detection with xyxy field + empty detection -# ( -# [ -# mock_detections(xyxy=[[10, 10, 20, 20]]), -# mock_detections(xyxy=[[20, 20, 30, 30]]), -# ], -# mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), -# DoesNotRaise(), -# ), # two detections with xyxy field -# ( -# [ -# mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]), -# mock_detections(xyxy=[[20, 20, 30, 30]]), -# ], -# mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), -# DoesNotRaise(), -# ), # detection with xyxy, class_id fields + detection with xyxy field -# ( -# [ -# mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]), -# mock_detections(xyxy=[[20, 20, 30, 30]], class_id=[1]), -# ], -# mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[0, 1]), -# DoesNotRaise(), -# ), # two detections with xyxy, class_id fields -# ], -# ) -# def test_merge( -# detections_list: List[Detections], -# expected_result: Optional[Detections], -# exception: Exception, -# ) -> None: -# with exception: -# result = Detections.merge(detections_list=detections_list) -# assert result == expected_result +@pytest.mark.parametrize( + "detections_list, expected_result, exception", + [ + ([], Detections.empty(), DoesNotRaise()), # empty detections list + ( + [Detections.empty()], + Detections.empty(), + DoesNotRaise(), + ), # single empty detections + ( + [mock_detections(xyxy=[[10, 10, 20, 20]])], + mock_detections(xyxy=[[10, 10, 20, 20]]), + DoesNotRaise(), + ), # single detection with xyxy field + ( + [mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()], + mock_detections(xyxy=[[10, 10, 20, 20]]), + DoesNotRaise(), + ), # single detection with xyxy field + empty detection + ( + [ + mock_detections(xyxy=[[10, 10, 20, 20]]), + mock_detections(xyxy=[[20, 20, 30, 30]]), + ], + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + DoesNotRaise(), + ), # two detections with xyxy field + ( + [ + mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]), + mock_detections(xyxy=[[20, 20, 30, 30]]), + ], + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + DoesNotRaise(), + ), # detection with xyxy, class_id fields + detection with xyxy field + ( + [ + mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]), + mock_detections(xyxy=[[20, 20, 30, 30]], class_id=[1]), + ], + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[0, 1]), + DoesNotRaise(), + ), # two detections with xyxy, class_id fields + ], +) +def test_merge( + detections_list: List[Detections], + expected_result: Optional[Detections], + exception: Exception, +) -> None: + with exception: + result = Detections.merge(detections_list=detections_list) + assert result == expected_result @pytest.mark.parametrize( @@ -324,3 +324,107 @@ def test_equal( detections_a: Detections, detections_b: Detections, expected_result: bool ) -> None: assert (detections_a == detections_b) == 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": [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 + ] +) +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}" From 27eb823809aff2cda60f9a854d2a0e2df7acaa6b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 29 Dec 2023 14:05:20 +0000 Subject: [PATCH 06/21] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto?= =?UTF-8?q?=20format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 6 ++---- test/detection/test_core.py | 37 ++++++++++++++++------------------- 2 files changed, 19 insertions(+), 24 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 2a57eb82..704bbc73 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -61,9 +61,7 @@ def _validate_tracker_id(tracker_id: Any, n: int) -> None: raise ValueError("tracker_id must be None or 1d np.ndarray with (n,) shape") -def is_data_equal( - data_a: Dict[str, np.ndarray], data_b: Dict[str, np.ndarray] -) -> bool: +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. @@ -79,7 +77,7 @@ def is_data_equal( def merge_data( - data_list: List[Dict[str, Union[np.ndarray, List]]] + 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. diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 41102d57..b21423e9 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -1,6 +1,6 @@ from contextlib import ExitStack as DoesNotRaise from test.test_utils import mock_detections -from typing import List, Optional, Union, Dict, Any +from typing import Any, Dict, List, Optional, Union import numpy as np import pytest @@ -335,17 +335,12 @@ def test_equal( DoesNotRaise(), ), # empty data list ( - [ - {} - ], + [{}], {}, DoesNotRaise(), ), # single empty data dict ( - [ - {}, - {} - ], + [{}, {}], {}, DoesNotRaise(), ), # two empty data dicts @@ -392,39 +387,41 @@ def test_equal( ), # 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]), "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']) + "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], "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']) + "test_2": np.array(["a", "b", "c", "c", "b", "a"]), }, DoesNotRaise(), ), # two data dicts with the same field names and mixed values - ] + ], ) def test_merge_data( data_list: List[Dict[str, Any]], expected_result: Optional[Dict[str, Any]], - exception: Exception + 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}" + 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}" + assert ( + result[key] == expected_result[key] + ), f"Mismatch in non-array data for key {key}" From 701b0c0b97d757ae21625e2fd753c71b3d781a92 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 29 Dec 2023 15:15:38 +0100 Subject: [PATCH 07/21] fix `test_detections_from_xml_obj` --- supervision/dataset/formats/pascal_voc.py | 11 ++- test/dataset/formats/test_pascal_voc.py | 103 ++++++++++++---------- 2 files changed, 59 insertions(+), 55 deletions(-) diff --git a/supervision/dataset/formats/pascal_voc.py b/supervision/dataset/formats/pascal_voc.py index 74f6be70..32ba7a74 100644 --- a/supervision/dataset/formats/pascal_voc.py +++ b/supervision/dataset/formats/pascal_voc.py @@ -268,12 +268,11 @@ 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/test/dataset/formats/test_pascal_voc.py b/test/dataset/formats/test_pascal_voc.py index 238967dc..fa58f257 100644 --- a/test/dataset/formats/test_pascal_voc.py +++ b/test/dataset/formats/test_pascal_voc.py @@ -7,8 +7,9 @@ import pytest from supervision.dataset.formats.pascal_voc import ( object_to_pascal_voc, - parse_polygon_points, + parse_polygon_points, detections_from_xml_obj, ) +from test.test_utils import mock_detections def are_xml_elements_equal(elem1, elem2): @@ -112,51 +113,55 @@ N_CLASS_N_BBOX = """test1 NO_DETECTIONS = """""" -# @pytest.mark.parametrize( -# "xml_string, classes, resolution_wh, force_masks, expected_result, exception", -# [ -# ( -# ONE_CLASS_ONE_BBOX, -# ["test"], -# (100, 100), -# False, -# mock_detections(np.array([[0, 0, 10, 10]]), None, [0]), -# DoesNotRaise(), -# ), -# ( -# ONE_CLASS_N_BBOX, -# ["test"], -# (100, 100), -# False, -# mock_detections(np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), None, [0, 0]), -# DoesNotRaise(), -# ), -# ( -# N_CLASS_N_BBOX, -# ["test", "test2"], -# (100, 100), -# False, -# mock_detections( -# np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]), -# None, -# [0, 0, 1], -# ), -# DoesNotRaise(), -# ), -# ( -# NO_DETECTIONS, -# [], -# (100, 100), -# False, -# mock_detections(np.empty((0, 4)), None, []), -# DoesNotRaise(), -# ), -# ], -# ) -# def test_detections_from_xml_obj( -# xml_string, classes, resolution_wh, force_masks, expected_result, exception -# ): -# with exception: -# root = ET.fromstring(xml_string) -# result, _ = detections_from_xml_obj(root, classes, resolution_wh, force_masks) -# assert result == expected_result +@pytest.mark.parametrize( + "xml_string, classes, resolution_wh, force_masks, expected_result, exception", + [ + ( + ONE_CLASS_ONE_BBOX, + ["test"], + (100, 100), + False, + mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]), + DoesNotRaise(), + ), + ( + ONE_CLASS_N_BBOX, + ["test"], + (100, 100), + False, + mock_detections( + xyxy=np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), + class_id=[0, 0] + ), + DoesNotRaise(), + ), + ( + N_CLASS_N_BBOX, + ["test", "test2"], + (100, 100), + False, + mock_detections( + xyxy=np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]), + class_id=[0, 0, 1], + ), + DoesNotRaise(), + ), + ( + NO_DETECTIONS, + [], + (100, 100), + False, + mock_detections(xyxy=np.empty((0, 4)), class_id=[]), + DoesNotRaise(), + ), + ], +) +def test_detections_from_xml_obj( + xml_string, classes, resolution_wh, force_masks, expected_result, exception +): + with exception: + root = ET.fromstring(xml_string) + result, _ = detections_from_xml_obj(root, classes, resolution_wh, force_masks) + print(result) + print(expected_result) + assert result == expected_result From 14761cb20a596e275252699ae5112dab37f13454 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 29 Dec 2023 15:15:57 +0100 Subject: [PATCH 08/21] fix `test_detections_from_xml_obj` --- test/dataset/formats/test_pascal_voc.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/test/dataset/formats/test_pascal_voc.py b/test/dataset/formats/test_pascal_voc.py index fa58f257..4aee7d74 100644 --- a/test/dataset/formats/test_pascal_voc.py +++ b/test/dataset/formats/test_pascal_voc.py @@ -162,6 +162,4 @@ def test_detections_from_xml_obj( with exception: root = ET.fromstring(xml_string) result, _ = detections_from_xml_obj(root, classes, resolution_wh, force_masks) - print(result) - print(expected_result) assert result == expected_result From 63ccc98d0b683adff58c5f70509dbf8f2f8afbae Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 29 Dec 2023 14:15:59 +0000 Subject: [PATCH 09/21] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto?= =?UTF-8?q?=20format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/dataset/formats/pascal_voc.py | 3 ++- test/dataset/formats/test_pascal_voc.py | 8 ++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/supervision/dataset/formats/pascal_voc.py b/supervision/dataset/formats/pascal_voc.py index 32ba7a74..ba77ddf8 100644 --- a/supervision/dataset/formats/pascal_voc.py +++ b/supervision/dataset/formats/pascal_voc.py @@ -271,7 +271,8 @@ def detections_from_xml_obj( annotation = Detections( xyxy=xyxy.astype(np.float32), mask=np.array(masks).astype(bool) if with_masks else None, - class_id=class_id) + class_id=class_id, + ) return annotation, extended_classes diff --git a/test/dataset/formats/test_pascal_voc.py b/test/dataset/formats/test_pascal_voc.py index fa58f257..23aa65a4 100644 --- a/test/dataset/formats/test_pascal_voc.py +++ b/test/dataset/formats/test_pascal_voc.py @@ -1,4 +1,5 @@ from contextlib import ExitStack as DoesNotRaise +from test.test_utils import mock_detections from typing import List, Optional import defusedxml.ElementTree as ET @@ -6,10 +7,10 @@ import numpy as np import pytest from supervision.dataset.formats.pascal_voc import ( + detections_from_xml_obj, object_to_pascal_voc, - parse_polygon_points, detections_from_xml_obj, + parse_polygon_points, ) -from test.test_utils import mock_detections def are_xml_elements_equal(elem1, elem2): @@ -130,8 +131,7 @@ NO_DETECTIONS = """""" (100, 100), False, mock_detections( - xyxy=np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), - class_id=[0, 0] + xyxy=np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), class_id=[0, 0] ), DoesNotRaise(), ), From 9be71da8886fc05783cdf93f8d688e3df8c70989 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 29 Dec 2023 16:15:36 +0100 Subject: [PATCH 10/21] handle one more edge case in `Detections.merge` logic --- supervision/detection/core.py | 12 +++++++++- test/detection/test_core.py | 45 +++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 704bbc73..b972afaa 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -87,7 +87,11 @@ def merge_data( Returns: A single data payload containing the merged data, preserving the original data - types (list or np.ndarray). + 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 {} @@ -96,6 +100,12 @@ def merge_data( 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: diff --git a/test/detection/test_core.py b/test/detection/test_core.py index b21423e9..3b9a53f1 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -344,6 +344,35 @@ def test_equal( {}, 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]}, @@ -407,6 +436,22 @@ def test_equal( }, 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( From db245bbc34f6049f7193c0d0a7ee72464706b85d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 29 Dec 2023 15:15:48 +0000 Subject: [PATCH 11/21] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto?= =?UTF-8?q?=20format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index b972afaa..214b5f42 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -104,7 +104,8 @@ def merge_data( 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.") + "All data values within a single object must have equal length." + ) merged_data = {key: [] for key in all_keys_sets[0]} From 95e2ad6d5225de927fa44d760549108d59735a6d Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 29 Dec 2023 16:36:54 +0100 Subject: [PATCH 12/21] refactor of `core.py`; moving utils to `utils.py` --- supervision/detection/core.py | 131 +------------- .../detection/tools/inference_slicer.py | 3 +- supervision/detection/utils.py | 127 +++++++++++++- test/detection/test_core.py | 164 +----------------- test/detection/test_utils.py | 150 +++++++++++++++- 5 files changed, 292 insertions(+), 283 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index b972afaa..c4995e30 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -1,7 +1,6 @@ from __future__ import annotations from dataclasses import astuple, dataclass, field -from itertools import chain from typing import Any, Dict, Iterator, List, Optional, Tuple, Union import numpy as np @@ -11,127 +10,12 @@ from supervision.detection.utils import ( extract_ultralytics_masks, non_max_suppression, process_roboflow_result, - xywh_to_xyxy, + xywh_to_xyxy, validate_xyxy, validate_mask, validate_class_id, validate_confidence, + validate_tracker_id, validate_data, is_data_equal, merge_data, ) 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") - - -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 - - @dataclass class Detections: """ @@ -162,11 +46,12 @@ class Detections: 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_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_data(data=self.data, n=n) def __len__(self): """ 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 a26f72ff..0499044b 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 List, Optional, Tuple, Dict, Union, Any import cv2 import numpy as np @@ -469,3 +470,127 @@ 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, 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_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") + + +def validate_data(data: Dict[str, Union[np.ndarray, List]], 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 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 diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 3b9a53f1..e71cdbb7 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -1,11 +1,11 @@ from contextlib import ExitStack as DoesNotRaise from test.test_utils import mock_detections -from typing import Any, Dict, List, Optional, Union +from typing import List, Optional, Union import numpy as np import pytest -from supervision.detection.core import Detections, merge_data +from supervision.detection.core import Detections from supervision.geometry.core import Position PREDICTIONS = np.array( @@ -299,23 +299,23 @@ def test_get_anchor_coordinates( False, ), # detection with xyxy field + detection with xyxy, confidence fields ( - mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1, 2, 3]}), - mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1, 2, 3]}), + 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, 2, 3]}), + 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, 2, 3]}), - mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_2": [1, 2, 3]}), + 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, 2, 3]}), - mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [3, 2, 1]}), + 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 ], @@ -324,149 +324,3 @@ def test_equal( detections_a: Detections, detections_b: Detections, expected_result: bool ) -> None: assert (detections_a == detections_b) == 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}" diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 83b71b85..139a02cf 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 List, Optional, Tuple, Dict, Any import numpy as np import pytest @@ -11,7 +11,7 @@ from supervision.detection.utils import ( move_boxes, non_max_suppression, process_roboflow_result, - scale_boxes, + scale_boxes, merge_data, ) TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool) @@ -640,3 +640,149 @@ 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}" From ea21a75135d24c6d9f0f171f637bc712d6aeb9e8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 29 Dec 2023 15:37:48 +0000 Subject: [PATCH 13/21] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto?= =?UTF-8?q?=20format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 11 +++++++++-- supervision/detection/utils.py | 8 +++++--- test/detection/test_utils.py | 5 +++-- 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index c4995e30..ea1a7e31 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -8,10 +8,17 @@ import numpy as np from supervision.detection.utils import ( calculate_masks_centroids, extract_ultralytics_masks, + is_data_equal, + merge_data, non_max_suppression, process_roboflow_result, - xywh_to_xyxy, validate_xyxy, validate_mask, validate_class_id, validate_confidence, - validate_tracker_id, validate_data, is_data_equal, merge_data, + validate_class_id, + validate_confidence, + validate_data, + validate_mask, + validate_tracker_id, + validate_xyxy, + xywh_to_xyxy, ) from supervision.geometry.core import Position diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 0499044b..5cccdb3f 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1,5 +1,5 @@ from itertools import chain -from typing import List, Optional, Tuple, Dict, Union, Any +from typing import Any, Dict, List, Optional, Tuple, Union import cv2 import numpy as np @@ -520,7 +520,8 @@ def validate_data(data: Dict[str, Union[np.ndarray, List]], n: int) -> None: 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}") + 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") @@ -568,7 +569,8 @@ def merge_data( 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.") + "All data values within a single object must have equal length." + ) merged_data = {key: [] for key in all_keys_sets[0]} diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 139a02cf..7cd87fb7 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, Dict, Any +from typing import Any, Dict, List, Optional, Tuple import numpy as np import pytest @@ -8,10 +8,11 @@ from supervision.detection.utils import ( calculate_masks_centroids, clip_boxes, filter_polygons_by_area, + merge_data, move_boxes, non_max_suppression, process_roboflow_result, - scale_boxes, merge_data, + scale_boxes, ) TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool) From 099d36011d4b37539c46c3f2deb92ee98d0eea2c Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 29 Dec 2023 16:53:41 +0100 Subject: [PATCH 14/21] better `Detections.merge` docs + more `Detections.merge` tests --- supervision/detection/core.py | 27 ++++++++++++++++++++++++--- test/detection/test_core.py | 10 ++++++++++ 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index ea1a7e31..8269571a 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -650,12 +650,33 @@ 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: diff --git a/test/detection/test_core.py b/test/detection/test_core.py index e71cdbb7..40ab22a6 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -181,6 +181,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( From 5984d3374cd46cded2e7c8455bf9749076fe2bf2 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 29 Dec 2023 18:24:24 +0100 Subject: [PATCH 15/21] better `Detections.merge` validation --- supervision/detection/core.py | 63 +++++++++++++++++++--------------- supervision/detection/utils.py | 53 +++++++++++++++++++++++----- test/detection/test_core.py | 15 +++++--- 3 files changed, 92 insertions(+), 39 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 8269571a..1ce5599f 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -1,7 +1,7 @@ from __future__ import annotations -from dataclasses import astuple, dataclass, field -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 @@ -12,13 +12,8 @@ from supervision.detection.utils import ( merge_data, non_max_suppression, process_roboflow_result, - validate_class_id, - validate_confidence, - validate_data, - validate_mask, - validate_tracker_id, - validate_xyxy, xywh_to_xyxy, + validate_detections_fields, ) from supervision.geometry.core import Position @@ -52,13 +47,14 @@ class Detections: 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_data(data=self.data, 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): """ @@ -682,20 +678,33 @@ class Detections: 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, data = [ - list(field_values) for field_values 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 - data = merge_data(data) + 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, diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 5cccdb3f..81f5983f 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -472,45 +472,65 @@ def calculate_masks_centroids(masks: np.ndarray) -> np.ndarray: return np.column_stack((centroid_x, centroid_y)).astype(int) -def validate_xyxy(xyxy: Any, n: int) -> None: - is_valid = isinstance(xyxy, np.ndarray) and xyxy.shape == (n, 4) +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("xyxy must be 2d np.ndarray with (n, 4) shape") + 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("mask must be 3d np.ndarray with (n, H, W) shape") + 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("class_id must be None or 1d np.ndarray with (n,) shape") + 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("confidence must be None or 1d np.ndarray with (n,) shape") + 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("tracker_id must be None or 1d np.ndarray with (n,) shape") + 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, Union[np.ndarray, List]], n: int) -> None: +def validate_data(data: Dict[str, Any], n: int) -> None: for key, value in data.items(): if isinstance(value, list): if len(value) != n: @@ -526,6 +546,23 @@ def validate_data(data: Dict[str, Union[np.ndarray, List]], n: int) -> None: 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. diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 40ab22a6..84cf8a98 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -143,17 +143,24 @@ def test_getitem( [ ([], Detections.empty(), DoesNotRaise()), # empty detections list ( - [Detections.empty()], + [ + Detections.empty() + ], Detections.empty(), DoesNotRaise(), ), # single empty detections ( - [mock_detections(xyxy=[[10, 10, 20, 20]])], + [ + mock_detections(xyxy=[[10, 10, 20, 20]]) + ], mock_detections(xyxy=[[10, 10, 20, 20]]), 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 +178,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 ( [ From e75e52792559a78358feedd6b9ca44b8d6f280b1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 29 Dec 2023 17:24:36 +0000 Subject: [PATCH 16/21] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto?= =?UTF-8?q?=20format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 22 ++++++++++++---------- supervision/detection/utils.py | 27 ++++++++++++++++----------- test/detection/test_core.py | 8 ++------ 3 files changed, 30 insertions(+), 27 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 1ce5599f..d66772f8 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -12,8 +12,8 @@ from supervision.detection.utils import ( merge_data, non_max_suppression, process_roboflow_result, - xywh_to_xyxy, validate_detections_fields, + xywh_to_xyxy, ) from supervision.geometry.core import Position @@ -53,7 +53,7 @@ class Detections: confidence=self.confidence, class_id=self.class_id, tracker_id=self.tracker_id, - data=self.data + data=self.data, ) def __len__(self): @@ -685,7 +685,7 @@ class Detections: confidence=detections.confidence, class_id=detections.class_id, tracker_id=detections.tracker_id, - data=detections.data + data=detections.data, ) xyxy = np.vstack([d.xyxy for d in detections_list]) @@ -695,14 +695,16 @@ class Detections: 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]) + 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') + 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]) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 81f5983f..8ff34f4d 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -474,60 +474,65 @@ def calculate_masks_centroids(masks: np.ndarray) -> np.ndarray: def validate_xyxy(xyxy: Any) -> None: expected_shape = "(_, 4)" - actual_shape = str(getattr(xyxy, 'shape', None)) + 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}") + f"{actual_shape}" + ) def validate_mask(mask: Any, n: int) -> None: expected_shape = f"({n}, H, W)" - actual_shape = str(getattr(mask, 'shape', None)) + 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}") + 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)) + 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}") + f"shape {actual_shape}" + ) def validate_confidence(confidence: Any, n: int) -> None: expected_shape = f"({n},)" - actual_shape = str(getattr(confidence, 'shape', None)) + 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}") + 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)) + 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}") + f"shape {actual_shape}" + ) def validate_data(data: Dict[str, Any], n: int) -> None: @@ -552,7 +557,7 @@ def validate_detections_fields( class_id: Any, confidence: Any, tracker_id: Any, - data: Dict[str, Any] + data: Dict[str, Any], ) -> None: validate_xyxy(xyxy) n = len(xyxy) diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 84cf8a98..f3b739e8 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -143,16 +143,12 @@ def test_getitem( [ ([], Detections.empty(), DoesNotRaise()), # empty detections list ( - [ - Detections.empty() - ], + [Detections.empty()], Detections.empty(), DoesNotRaise(), ), # single empty detections ( - [ - mock_detections(xyxy=[[10, 10, 20, 20]]) - ], + [mock_detections(xyxy=[[10, 10, 20, 20]])], mock_detections(xyxy=[[10, 10, 20, 20]]), DoesNotRaise(), ), # single detection with xyxy field From ced9f35b6ee36fdd2b4441e48a6b45d99168e332 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 2 Jan 2024 15:05:20 +0100 Subject: [PATCH 17/21] initial `Detections.__getitem__` implementation --- supervision/detection/core.py | 3 +- supervision/detection/utils.py | 33 +++++++++ test/detection/test_utils.py | 131 +++++++++++++++++++++++++++++++++ 3 files changed, 166 insertions(+), 1 deletion(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index d66772f8..70eeac22 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -13,7 +13,7 @@ from supervision.detection.utils import ( non_max_suppression, process_roboflow_result, validate_detections_fields, - xywh_to_xyxy, + xywh_to_xyxy, get_data_item, ) from supervision.geometry.core import Position @@ -821,6 +821,7 @@ 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), ) @property diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 8ff34f4d..bc47d9cc 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -638,3 +638,36 @@ def merge_data( ) 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/detection/test_utils.py b/test/detection/test_utils.py index 7cd87fb7..675680d7 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -13,6 +13,7 @@ from supervision.detection.utils import ( non_max_suppression, process_roboflow_result, scale_boxes, + get_data_item, ) TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool) @@ -787,3 +788,133 @@ def test_merge_data( 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}" From 51694819e3ac85883e9575bece6af978548db793 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 2 Jan 2024 14:05:34 +0000 Subject: [PATCH 18/21] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto?= =?UTF-8?q?=20format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 3 ++- supervision/detection/utils.py | 2 +- test/detection/test_utils.py | 33 ++++++++++++++------------------- 3 files changed, 17 insertions(+), 21 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 70eeac22..8d319302 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -8,12 +8,13 @@ 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, get_data_item, + xywh_to_xyxy, ) from supervision.geometry.core import Position diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index bc47d9cc..2df89a84 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -642,7 +642,7 @@ def merge_data( def get_data_item( data: Dict[str, Union[np.ndarray, List]], - index: Union[int, slice, List[int], np.ndarray] + 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. diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 675680d7..a24119ce 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -8,12 +8,12 @@ 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, scale_boxes, - get_data_item, ) TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool) @@ -793,12 +793,7 @@ def test_merge_data( @pytest.mark.parametrize( "data, index, expected_result, exception", [ - ( - {}, - 0, - {}, - DoesNotRaise() - ), # empty data dict + ({}, 0, {}, DoesNotRaise()), # empty data dict ( { "test_1": [1, 2, 3], @@ -807,7 +802,7 @@ def test_merge_data( { "test_1": [1], }, - DoesNotRaise() + DoesNotRaise(), ), # single data dict with a single field name and list values ( { @@ -817,7 +812,7 @@ def test_merge_data( { "test_1": np.array([1]), }, - DoesNotRaise() + DoesNotRaise(), ), # single data dict with a single field name and np.array values as 1D arrays ( { @@ -827,7 +822,7 @@ def test_merge_data( { "test_1": [1, 2], }, - DoesNotRaise() + DoesNotRaise(), ), # single data dict with a single field name and list values ( { @@ -837,7 +832,7 @@ def test_merge_data( { "test_1": np.array([1, 2]), }, - DoesNotRaise() + DoesNotRaise(), ), # single data dict with a single field name and np.array values as 1D arrays ( { @@ -847,7 +842,7 @@ def test_merge_data( { "test_1": [3], }, - DoesNotRaise() + DoesNotRaise(), ), # single data dict with a single field name and list values ( { @@ -857,7 +852,7 @@ def test_merge_data( { "test_1": np.array([3]), }, - DoesNotRaise() + DoesNotRaise(), ), # single data dict with a single field name and np.array values as 1D arrays ( { @@ -867,7 +862,7 @@ def test_merge_data( { "test_1": [1, 3], }, - DoesNotRaise() + DoesNotRaise(), ), # single data dict with a single field name and list values ( { @@ -877,7 +872,7 @@ def test_merge_data( { "test_1": np.array([1, 3]), }, - DoesNotRaise() + DoesNotRaise(), ), # single data dict with a single field name and np.array values as 1D arrays ( { @@ -887,7 +882,7 @@ def test_merge_data( { "test_1": [1, 3], }, - DoesNotRaise() + DoesNotRaise(), ), # single data dict with a single field name and list values ( { @@ -897,9 +892,9 @@ def test_merge_data( { "test_1": np.array([1, 3]), }, - DoesNotRaise() + DoesNotRaise(), ), - ] + ], ) def test_get_data_item( data: Dict[str, Any], @@ -916,5 +911,5 @@ def test_get_data_item( ), f"Mismatch in arrays for key {key}" else: assert ( - result[key] == expected_result[key] + result[key] == expected_result[key] ), f"Mismatch in non-array data for key {key}" From a10b9840afa6ee1cd6b960b260d57bb22c8a2026 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 2 Jan 2024 15:58:44 +0100 Subject: [PATCH 19/21] initial `Detections.__iter__` implementation --- supervision/detection/core.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 70eeac22..8ff5340c 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -71,11 +71,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 ( @@ -84,6 +85,7 @@ 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): From 0f6396669b1e2f5d1861dc8aab5d90e943ded504 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 2 Jan 2024 17:37:26 +0100 Subject: [PATCH 20/21] initial `Detections.__setitem__` implementation + `Detections.__getitem__` changes --- supervision/detection/core.py | 41 +++++++++++++++++++++++++---------- 1 file changed, 30 insertions(+), 11 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 5216bb84..0385bddf 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -787,35 +787,41 @@ 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 + ``` >>> import supervision as sv >>> 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( @@ -827,6 +833,19 @@ class Detections: 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. + """ + if not isinstance(value, (np.ndarray, list)): + raise TypeError("Value must be a np.ndarray or a list") + + self.data[key] = value + @property def area(self) -> np.ndarray: """ From e3d92a01c4eff23ffc2f3358d8ae144c61a6105b Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 2 Jan 2024 17:58:08 +0100 Subject: [PATCH 21/21] ready for colab tests --- supervision/detection/core.py | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 0385bddf..201b4e9a 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -38,6 +38,14 @@ class 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 @@ -806,7 +814,7 @@ class Detections: the data field. Example: - ``` + ```python >>> import supervision as sv >>> detections = sv.Detections(...) @@ -840,10 +848,33 @@ class Detections: 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