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

Retain type info by not excluding empty detections

This commit is contained in:
Linas Kondrackis 2024-05-09 18:19:51 +03:00
parent 4cd1fbcef9
commit 01f7eb5be6
2 changed files with 42 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
supervision/detection/utils.py
View File

@ -713,27 +713,10 @@ def merge_data(
merged_data = {key: [] for key in all_keys}
for data in data_list:
for key in data:
if len(data[key]) == 0:
continue
merged_data[key].append(data[key])
for key in merged_data:
if len(merged_data[key]) == 0:
for data in data_list:
if key not in data:
continue
data_type = type(data[key])
break
if data_type == np.ndarray:
merged_data[key] = np.array(merged_data[key])
elif data_type == list:
merged_data[key] = list(merged_data[key])
else:
raise ValueError(
f"Inconsistent data types for key '{key}'. Only np.ndarray and list "
f"types are allowed."
)
elif all(isinstance(item, list) for item in merged_data[key]):
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

80
test/detection/test_core.py
View File

@ -33,73 +33,75 @@ DETECTIONS = Detections(
# Merge test
TEST_MASK = np.zeros((1000, 1000), dtype=bool)
TEST_MASK[300:351, 200:251] = True
TEST_DET_1 = mock_detections(
xyxy=[[10, 10, 20, 20], [30, 30, 40, 40], [50, 50, 60, 60]],
mask=[TEST_MASK, TEST_MASK, TEST_MASK],
confidence=[0.1, 0.2, 0.3],
class_id=[1, 2, 3],
tracker_id=[1, 2, 3],
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 = mock_detections(
xyxy=[[70, 70, 80, 80], [90, 90, 100, 100]],
mask=[TEST_MASK, TEST_MASK],
confidence=[0.4, 0.5],
class_id=[4, 5],
tracker_id=[4, 5],
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 = mock_detections(
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],
tracker_id=[1, 2, 3, 4, 5],
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 = mock_detections(
TEST_DET_ZERO_LENGTH = Detections(
xyxy=np.empty((0, 4), dtype=np.float32),
mask=np.empty((0, *TEST_MASK.shape), dtype=bool),
confidence=[],
class_id=[],
tracker_id=[],
confidence=np.empty((0,)),
class_id=np.empty((0,)),
tracker_id=np.empty((0,)),
data={
"some_key": [],
"other_key": [],
},
)
TEST_DET_NONE = mock_detections(
TEST_DET_NONE = Detections(
xyxy=np.empty((0, 4), dtype=np.float32),
)
TEST_DET_DIFFERENT_FIELDS = mock_detections(
xyxy=[[88, 88, 99, 99]],
mask=[np.logical_not(TEST_MASK)],
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=[9],
tracker_id=np.array([9]),
data={"some_key": [9], "other_key": [["11", "12"]]},
)
TEST_DET_DIFFERENT_DATA = mock_detections(
xyxy=[[88, 88, 99, 99]],
mask=[np.logical_not(TEST_MASK)],
confidence=[0.9],
class_id=[9],
tracker_id=[9],
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],
},