initial `Detections.__getitem__` implementation
This commit is contained in:
SkalskiP
parent
e75e527925
commit
ced9f35b6e
|
|
@ -13,7 +13,7 @@ from supervision.detection.utils import (
|
|||
non_max_suppression,
|
||||
process_roboflow_result,
|
||||
validate_detections_fields,
|
||||
xywh_to_xyxy,
|
||||
xywh_to_xyxy, get_data_item,
|
||||
)
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
|
|
@ -821,6 +821,7 @@ class Detections:
|
|||
confidence=self.confidence[index] if self.confidence is not None else None,
|
||||
class_id=self.class_id[index] if self.class_id is not None else None,
|
||||
tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
|
||||
data=get_data_item(self.data, index),
|
||||
)
|
||||
|
||||
@property
|
||||
|
|
|
|||
|
|
@ -638,3 +638,36 @@ def merge_data(
|
|||
)
|
||||
|
||||
return merged_data
|
||||
|
||||
|
||||
def get_data_item(
|
||||
data: Dict[str, Union[np.ndarray, List]],
|
||||
index: Union[int, slice, List[int], np.ndarray]
|
||||
) -> Dict[str, Union[np.ndarray, List]]:
|
||||
"""
|
||||
Retrieve a subset of the data dictionary based on the given index.
|
||||
|
||||
Args:
|
||||
data: The data dictionary of the Detections object.
|
||||
index: The index or indices specifying the subset to retrieve.
|
||||
|
||||
Returns:
|
||||
A subset of the data dictionary corresponding to the specified index.
|
||||
"""
|
||||
subset_data = {}
|
||||
for key, value in data.items():
|
||||
if isinstance(value, np.ndarray):
|
||||
subset_data[key] = value[index]
|
||||
elif isinstance(value, list):
|
||||
if isinstance(index, slice):
|
||||
subset_data[key] = value[index]
|
||||
elif isinstance(index, (list, np.ndarray)):
|
||||
subset_data[key] = [value[i] for i in index]
|
||||
elif isinstance(index, int):
|
||||
subset_data[key] = [value[index]]
|
||||
else:
|
||||
raise TypeError(f"Unsupported index type: {type(index)}")
|
||||
else:
|
||||
raise TypeError(f"Unsupported data type for key '{key}': {type(value)}")
|
||||
|
||||
return subset_data
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from supervision.detection.utils import (
|
|||
non_max_suppression,
|
||||
process_roboflow_result,
|
||||
scale_boxes,
|
||||
get_data_item,
|
||||
)
|
||||
|
||||
TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool)
|
||||
|
|
@ -787,3 +788,133 @@ def test_merge_data(
|
|||
assert (
|
||||
result[key] == expected_result[key]
|
||||
), f"Mismatch in non-array data for key {key}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data, index, expected_result, exception",
|
||||
[
|
||||
(
|
||||
{},
|
||||
0,
|
||||
{},
|
||||
DoesNotRaise()
|
||||
), # empty data dict
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
0,
|
||||
{
|
||||
"test_1": [1],
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
0,
|
||||
{
|
||||
"test_1": np.array([1]),
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single data dict with a single field name and np.array values as 1D arrays
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
slice(0, 2),
|
||||
{
|
||||
"test_1": [1, 2],
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
slice(0, 2),
|
||||
{
|
||||
"test_1": np.array([1, 2]),
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single data dict with a single field name and np.array values as 1D arrays
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
-1,
|
||||
{
|
||||
"test_1": [3],
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
-1,
|
||||
{
|
||||
"test_1": np.array([3]),
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single data dict with a single field name and np.array values as 1D arrays
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
[0, 2],
|
||||
{
|
||||
"test_1": [1, 3],
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
[0, 2],
|
||||
{
|
||||
"test_1": np.array([1, 3]),
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single data dict with a single field name and np.array values as 1D arrays
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
np.array([0, 2]),
|
||||
{
|
||||
"test_1": [1, 3],
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
np.array([0, 2]),
|
||||
{
|
||||
"test_1": np.array([1, 3]),
|
||||
},
|
||||
DoesNotRaise()
|
||||
),
|
||||
]
|
||||
)
|
||||
def test_get_data_item(
|
||||
data: Dict[str, Any],
|
||||
index: Any,
|
||||
expected_result: Optional[Dict[str, Any]],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = get_data_item(data=data, index=index)
|
||||
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}"
|
||||
|
|
|
|||
Loading…
Reference in New Issue