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

Detections.merge merges None and []

Browse Source 
* 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.
This commit is contained in:
Linas Kondrackis committed 2024-05-07 17:30:12 +03:00
1 parent fb49be8197
commit 7bb94ce966
5 files changed
+99 -60
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

supervision/detection/core.py
+10 -10
View File
@@ -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")
supervision/detection/utils.py
+39 -11
View File
@@ -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]))
test/detection/test_core.py
+21 -27
View File
@@ -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",
[
(
test/detection/test_utils.py
+23
View File
@@ -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(
test/test_utils.py
+6 -12
View File
@@ -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}."