diff --git a/supervision/detection/core.py b/supervision/detection/core.py index d66772f8..70eeac22 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -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 diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 8ff34f4d..bc47d9cc 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -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 diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 7cd87fb7..675680d7 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -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}"