From efed23f64e1db5ca4bd8b5b3829ce30cb7f9e49e Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 29 Dec 2023 15:05:07 +0100 Subject: [PATCH] `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}"