diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 31998225..0ba9e4f4 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -834,9 +834,10 @@ class Detections: This method takes a list of Detections objects and combines their respective fields (`xyxy`, `mask`, `confidence`, `class_id`, and `tracker_id`) - into a single Detections object. If all elements in a field are not - `None`, the corresponding field will be stacked. - Otherwise, the field will be set to `None`. + into a single Detections object. + + For example, if merging Detections with 3 and 4 detected objects, this method + will return a Detections with 7 objects (7 entries in `xyxy`, `mask`, etc). Args: detections_list (List[Detections]): A list of Detections objects to merge. @@ -894,13 +895,12 @@ class Detections: 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]) - ) + stack_list = [ + d.__getattribute__(name) + for d in detections_list + if d.__getattribute__(name) is not None + ] + return np.vstack(stack_list) if name == "mask" else np.hstack(stack_list) mask = stack_or_none("mask") confidence = stack_or_none("confidence") diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 3eeba5b4..6b378042 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -678,7 +678,9 @@ def merge_data( Merges the data payloads of a list of Detections instances. Args: - data_list: The data payloads of the instances. + data_list: The data payloads of the Detections instances. Each data payload + is a dictionary with the same keys, and the values are either lists or + np.ndarray. Returns: A single data payload containing the merged data, preserving the original data @@ -691,10 +693,6 @@ def merge_data( 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: @@ -702,10 +700,23 @@ def merge_data( "All data values within a single object must have equal length." ) - merged_data = {key: [] for key in all_keys_sets[0]} + keys_by_data = [set(data.keys()) for data in data_list] + keys_by_data = [keys for keys in keys_by_data if len(keys) > 0] + if not keys_by_data: + return {} + common_keys = set.intersection(*keys_by_data) + all_keys = set.union(*keys_by_data) + if common_keys != all_keys: + raise ValueError( + f"All sv.Detections.data dictionaries must have the same keys. Common " + f"keys: {common_keys}, but some dictionaries have additional keys: " + f"{all_keys.difference(common_keys)}." + ) + + merged_data = {key: [] for key in all_keys} for data in data_list: - for key in merged_data: + for key in data: merged_data[key].append(data[key]) for key in merged_data: diff --git a/test/detection/test_core.py b/test/detection/test_core.py index f3b739e8..12f3de28 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -30,6 +30,84 @@ DETECTIONS = Detections( ) +# Merge test +TEST_MASK = np.zeros((1000, 1000), dtype=bool) +TEST_MASK[300:351, 200:251] = True +TEST_DET_1 = Detections( + xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40], [50, 50, 60, 60]]), + mask=np.array([TEST_MASK, TEST_MASK, TEST_MASK]), + confidence=np.array([0.1, 0.2, 0.3]), + class_id=np.array([1, 2, 3]), + tracker_id=np.array([1, 2, 3]), + data={ + "some_key": [1, 2, 3], + "other_key": [["1", "2"], ["3", "4"], ["5", "6"]], + }, +) +TEST_DET_2 = Detections( + xyxy=np.array([[70, 70, 80, 80], [90, 90, 100, 100]]), + mask=np.array([TEST_MASK, TEST_MASK]), + confidence=np.array([0.4, 0.5]), + class_id=np.array([4, 5]), + tracker_id=np.array([4, 5]), + data={ + "some_key": [4, 5], + "other_key": [["7", "8"], ["9", "10"]], + }, +) +TEST_DET_1_2 = Detections( + xyxy=np.array( + [ + [10, 10, 20, 20], + [30, 30, 40, 40], + [50, 50, 60, 60], + [70, 70, 80, 80], + [90, 90, 100, 100], + ] + ), + mask=np.array([TEST_MASK, TEST_MASK, TEST_MASK, TEST_MASK, TEST_MASK]), + confidence=np.array([0.1, 0.2, 0.3, 0.4, 0.5]), + class_id=np.array([1, 2, 3, 4, 5]), + tracker_id=np.array([1, 2, 3, 4, 5]), + data={ + "some_key": [1, 2, 3, 4, 5], + "other_key": [["1", "2"], ["3", "4"], ["5", "6"], ["7", "8"], ["9", "10"]], + }, +) +TEST_DET_ZERO_LENGTH = Detections( + xyxy=np.empty((0, 4), dtype=np.float32), + mask=np.empty((0, *TEST_MASK.shape), dtype=bool), + confidence=np.empty((0,)), + class_id=np.empty((0,)), + tracker_id=np.empty((0,)), + data={ + "some_key": [], + "other_key": [], + }, +) +TEST_DET_NONE = Detections( + xyxy=np.empty((0, 4), dtype=np.float32), +) +TEST_DET_DIFFERENT_FIELDS = Detections( + xyxy=np.array([[88, 88, 99, 99]]), + mask=np.array([np.logical_not(TEST_MASK)]), + confidence=None, + class_id=None, + tracker_id=np.array([9]), + data={"some_key": [9], "other_key": [["11", "12"]]}, +) +TEST_DET_DIFFERENT_DATA = Detections( + xyxy=np.array([[88, 88, 99, 99]]), + mask=np.array([np.logical_not(TEST_MASK)]), + confidence=np.array([0.9]), + class_id=np.array([9]), + tracker_id=np.array([9]), + data={ + "never_seen_key": [9], + }, +) + + @pytest.mark.parametrize( "detections, index, expected_result, exception", [ @@ -148,52 +226,58 @@ def test_getitem( DoesNotRaise(), ), # single empty detections ( - [mock_detections(xyxy=[[10, 10, 20, 20]])], - mock_detections(xyxy=[[10, 10, 20, 20]]), + [Detections.empty(), Detections.empty()], + Detections.empty(), DoesNotRaise(), - ), # single detection with xyxy field + ), # two empty detections + ( + [TEST_DET_1], + TEST_DET_1, + DoesNotRaise(), + ), # single detection with fields + ( + [TEST_DET_NONE], + TEST_DET_NONE, + DoesNotRaise(), + ), # Single weakly-defined detection + ( + [TEST_DET_1, TEST_DET_2], + TEST_DET_1_2, + DoesNotRaise(), + ), # Fields with same keys + # Fields and empty + ( + [TEST_DET_1, Detections.empty()], + TEST_DET_1, + DoesNotRaise(), + ), # single detection with fields ( [ - mock_detections(xyxy=[[10, 10, 20, 20]]), - mock_detections(xyxy=np.empty((0, 4), dtype=np.float32)), + TEST_DET_1, + TEST_DET_ZERO_LENGTH, ], - mock_detections(xyxy=[[10, 10, 20, 20]]), + TEST_DET_1, DoesNotRaise(), - ), # single detection with xyxy field + empty detection + ), # Single detection and empty-array fields ( [ - mock_detections(xyxy=[[10, 10, 20, 20]]), - mock_detections(xyxy=[[20, 20, 30, 30]]), + TEST_DET_1, + TEST_DET_NONE, ], - mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + TEST_DET_1, DoesNotRaise(), - ), # two detections with xyxy field + ), # Single detection and None fields (+ missing Dict keys) + # Errors: Non-zero-length differently defined keys & data ( - [ - 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]]), + [TEST_DET_1, TEST_DET_DIFFERENT_FIELDS], + None, pytest.raises(ValueError), - ), # detection with xyxy, class_id fields + detection with xyxy field + ), # Non-empty detections with different fields ( - [ - 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 - ( - [ - 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 + [TEST_DET_1, TEST_DET_DIFFERENT_DATA], + None, + pytest.raises(ValueError), + ), # Non-empty detections with different data keys ], ) def test_merge( diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 1c4a1d34..097c5c6e 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -911,6 +911,14 @@ def test_calculate_masks_centroids( {"test_1": []}, DoesNotRaise(), ), # single data dict with a single field name and empty list values + ( + [ + {"test_1": []}, + {"test_1": []}, + ], + {"test_1": []}, + DoesNotRaise(), + ), # two data dicts with the same field name and empty list values ( [ {"test_1": np.array([])}, @@ -918,6 +926,14 @@ def test_calculate_masks_centroids( {"test_1": np.array([])}, DoesNotRaise(), ), # single data dict with a single field name and empty np.array values + ( + [ + {"test_1": np.array([])}, + {"test_1": np.array([])}, + ], + {"test_1": np.array([])}, + DoesNotRaise(), + ), # two data dicts with the same field name and empty np.array values ( [ {"test_1": [1, 2, 3]}, @@ -932,7 +948,7 @@ def test_calculate_masks_centroids( ], {"test_1": [3, 2, 1]}, DoesNotRaise(), - ), # two data dicts with the same field name and empty and list values + ), # two data dicts with the same field name; one of with empty list as value ( [ {"test_1": [1, 2, 3]}, @@ -1012,6 +1028,49 @@ def test_calculate_masks_centroids( None, pytest.raises(ValueError), ), # two data dicts with the same field name and different length arrays values + ( + [{}, {"test_1": [1, 2, 3]}], + {"test_1": [1, 2, 3]}, + DoesNotRaise(), + ), # two data dicts; one empty and one non-empty dict + ( + [{"test_1": [], "test_2": []}, {"test_1": [1, 2, 3], "test_2": [1, 2, 3]}], + {"test_1": [1, 2, 3], "test_2": [1, 2, 3]}, + DoesNotRaise(), + ), # two data dicts; one empty and one non-empty dict; same keys + ( + [{"test_1": []}, {"test_1": [1, 2, 3], "test_2": [4, 5, 6]}], + None, + pytest.raises(ValueError), + ), # two data dicts; one empty and one non-empty dict; different keys + ( + [ + { + "test_1": [1, 2, 3], + "test_2": [4, 5, 6], + "test_3": [7, 8, 9], + }, + {"test_1": [1, 2, 3], "test_2": [4, 5, 6]}, + ], + None, + pytest.raises(ValueError), + ), # two data dicts; one with three keys, one with two keys + ( + [ + {"test_1": [1, 2, 3]}, + {"test_1": [1, 2, 3], "test_2": [1, 2, 3]}, + ], + None, + pytest.raises(ValueError), + ), # some keys missing in one dict + ( + [ + {"test_1": [1, 2, 3], "test_2": ["a", "b"]}, + {"test_1": [4, 5], "test_2": ["c", "d", "e"]}, + ], + None, + pytest.raises(ValueError), + ), # different value lengths for the same key ], ) def test_merge_data( @@ -1021,6 +1080,9 @@ def test_merge_data( ): with exception: result = merge_data(data_list=data_list) + if expected_result is None: + assert False, f"Expected an error, but got result {result}" + for key in result: if isinstance(result[key], np.ndarray): assert np.array_equal(