From 7bb94ce966465bbcb9815acc8a0b1c8a2100e0a9 Mon Sep 17 00:00:00 2001 From: Linas Kondrackis Date: Tue, 7 May 2024 17:30:12 +0300 Subject: [PATCH] Detections.merge merges None and [] * Detections.merge is much friendlier now. If there's a None or an empty array, it will merge it happily rather than complaining that everything needs to be either None or []. * Data merge follows suit. --- supervision/detection/core.py | 20 +++++++------- supervision/detection/utils.py | 50 ++++++++++++++++++++++++++-------- test/detection/test_core.py | 48 ++++++++++++++------------------ test/detection/test_utils.py | 23 ++++++++++++++++ test/test_utils.py | 18 ++++-------- 5 files changed, 99 insertions(+), 60 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 1900954d..e9baef7a 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -831,9 +831,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. @@ -891,13 +892,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..8fac9f90 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1,5 +1,5 @@ from itertools import chain -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Set, Tuple, Union import cv2 import numpy as np @@ -691,23 +691,51 @@ 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: + lengths_set = [len(value) for value in data.values()] + if len(set(lengths_set)) > 1: raise ValueError( "All data values within a single object must have equal length." ) - merged_data = {key: [] for key in all_keys_sets[0]} - + all_keys: Set[str] = set() for data in data_list: - for key in merged_data: - merged_data[key].append(data[key]) + all_keys.update(data.keys()) + # Naively merging entries and then validating length comes with a problem: + # N values may come from data[0]["key_1"] and N values from data[1]["key_2"]. + # These should not be joined together. + # Here, as soon as we find data of len > 0, we lock the key set and raise + # a ValueError if later we find a value of len > 0 with an unknown key. + key_set = None + merged_data = {key: [] for key in all_keys} + for data in data_list: + data_key_set = set() + for key in data: + if len(data[key]) > 0: + if key_set is None: + data_key_set.add(key) + elif key not in key_set: + raise ValueError(f"Unknown key '{key}' found in data payload.") + merged_data[key].append(data[key]) + + if key_set is None and data_key_set: + key_set = data_key_set + + merged_data = {key: val for key, val in merged_data.items() if len(val) > 0} + + sum_lengths = {} # Validation. More useful than set for error message + for key, value in merged_data.items(): + sum_length = sum(len(item) for item in value) + sum_lengths[key] = sum_length + lengths_set = set(sum_lengths.values()) + if len(lengths_set) > 1: + raise ValueError( + f"All data fields should have the same lengths after merge." + f"Resulting lengths: {sum_lengths}" + ) + + key_set = set() 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])) diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 8912f4a6..8f156238 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -42,7 +42,7 @@ TEST_DET_1 = mock_detections( data={ "some_key": [1, 2, 3], "other_key": [["1", "2"], ["3", "4"], ["5", "6"]], - } + }, ) TEST_DET_2 = mock_detections( xyxy=[[70, 70, 80, 80], [90, 90, 100, 100]], @@ -53,11 +53,16 @@ TEST_DET_2 = mock_detections( data={ "some_key": [4, 5], "other_key": [["7", "8"], ["9", "10"]], - } + }, ) TEST_DET_1_2 = mock_detections( - xyxy=[[10, 10, 20, 20], [30, 30, 40, 40], [ - 50, 50, 60, 60], [70, 70, 80, 80], [90, 90, 100, 100]], + xyxy=[ + [10, 10, 20, 20], + [30, 30, 40, 40], + [50, 50, 60, 60], + [70, 70, 80, 80], + [90, 90, 100, 100], + ], mask=[TEST_MASK, TEST_MASK, TEST_MASK, TEST_MASK, TEST_MASK], confidence=[0.1, 0.2, 0.3, 0.4, 0.5], class_id=[1, 2, 3, 4, 5], @@ -65,7 +70,7 @@ TEST_DET_1_2 = mock_detections( data={ "some_key": [1, 2, 3, 4, 5], "other_key": [["1", "2"], ["3", "4"], ["5", "6"], ["7", "8"], ["9", "10"]], - } + }, ) TEST_DET_ZERO_LENGTH = mock_detections( xyxy=np.empty((0, 4), dtype=np.float32), @@ -76,7 +81,7 @@ TEST_DET_ZERO_LENGTH = mock_detections( data={ "some_key": [], "other_key": [], - } + }, ) TEST_DET_NONE = mock_detections( xyxy=np.empty((0, 4), dtype=np.float32), @@ -87,10 +92,7 @@ TEST_DET_DIFFERENT_FIELDS = mock_detections( confidence=None, class_id=None, tracker_id=[9], - data={ - "some_key": [9], - "other_key": [["11", "12"]] - } + data={"some_key": [9], "other_key": [["11", "12"]]}, ) TEST_DET_DIFFERENT_DATA = mock_detections( xyxy=[[88, 88, 99, 99]], @@ -100,11 +102,11 @@ TEST_DET_DIFFERENT_DATA = mock_detections( tracker_id=[9], data={ "never_seen_key": [9], - } + }, ) -@ pytest.mark.parametrize( +@pytest.mark.parametrize( "detections, index, expected_result, exception", [ ( @@ -189,8 +191,7 @@ TEST_DET_DIFFERENT_DATA = mock_detections( DoesNotRaise(), ), # take only first detection by index slice (1, 3) (DETECTIONS, 10, None, pytest.raises(IndexError)), # index out of range - (DETECTIONS, [0, 2, 10], None, pytest.raises( - IndexError)), # index out of range + (DETECTIONS, [0, 2, 10], None, pytest.raises(IndexError)), # index out of range (DETECTIONS, np.array([0, 2, 10]), None, pytest.raises(IndexError)), ( DETECTIONS, @@ -213,12 +214,11 @@ def test_getitem( assert result == expected_result -@ pytest.mark.parametrize( +@pytest.mark.parametrize( "detections_list, expected_result, exception", [ # Nothing ([], Detections.empty(), DoesNotRaise()), # empty detections list - # Single ( [Detections.empty()], @@ -235,7 +235,6 @@ def test_getitem( TEST_DET_NONE, DoesNotRaise(), ), # Single weakly-defined detection - # Similar ( [Detections.empty(), Detections.empty()], @@ -247,13 +246,9 @@ def test_getitem( TEST_DET_1_2, DoesNotRaise(), ), # Fields with same keys - # Fields and empty ( - [ - TEST_DET_1, - Detections.empty() - ], + [TEST_DET_1, Detections.empty()], TEST_DET_1, DoesNotRaise(), ), # single detection with fields @@ -273,19 +268,18 @@ def test_getitem( TEST_DET_1, DoesNotRaise(), ), # Single detection and None fields (+ missing Dict keys) - # Errors: Non-zero-length differently defined keys & data ( [TEST_DET_1, TEST_DET_DIFFERENT_FIELDS], None, - pytest.raises(ValueError) + pytest.raises(ValueError), ), # Non-empty detections with different fields ( [TEST_DET_1, TEST_DET_DIFFERENT_DATA], None, pytest.raises(ValueError), ), # Non-empty detections with different data keys - ] + ], ) def test_merge( detections_list: List[Detections], @@ -297,7 +291,7 @@ def test_merge( assert result == expected_result -@ pytest.mark.parametrize( +@pytest.mark.parametrize( "detections, anchor, expected_result, exception", [ ( @@ -379,7 +373,7 @@ def test_get_anchor_coordinates( assert np.array_equal(result, expected_result) -@ pytest.mark.parametrize( +@pytest.mark.parametrize( "detections_a, detections_b, expected_result", [ ( diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 1c4a1d34..0a48be36 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -1012,6 +1012,29 @@ 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(), + ), # No keys in one 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(), + ), # Empty values dicts + ( + [{"test_1": []}, {"test_1": [1, 2, 3], "test_2": [1, 2, 3]}], + {"test_1": [1, 2, 3], "test_2": [1, 2, 3]}, + DoesNotRaise(), + ), # Mix of missing key and empty values + ( + [ + {"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 ], ) def test_merge_data( diff --git a/test/test_utils.py b/test/test_utils.py index 37be31d3..b676cb54 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -21,14 +21,11 @@ def mock_detections( xyxy=np.array(xyxy, dtype=np.float32), 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) + 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)), + 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 {}, ) @@ -46,15 +43,12 @@ def mock_keypoints( return KeyPoints( xy=np.array(xy, dtype=np.float32), confidence=( - confidence if confidence is None else np.array( - confidence, dtype=np.float32) + 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)), + class_id=(class_id if class_id is None else np.array(class_id, dtype=int)), data=convert_data(data) if data else {}, ) def assert_almost_equal(actual, expected, tolerance=1e-5): - assert abs( - actual - expected) < tolerance, f"Expected {expected}, but got {actual}." + assert abs(actual - expected) < tolerance, f"Expected {expected}, but got {actual}."