diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 75784ee8..dc88e9cc 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import astuple, dataclass -from typing import Any, Iterator, List, Optional, Tuple +from typing import Any, Iterator, List, Optional, Tuple, Union import numpy as np @@ -469,23 +469,43 @@ class Detections: raise ValueError(f"{anchor} is not supported.") - def __getitem__(self, index: np.ndarray) -> Detections: - if isinstance(index, np.ndarray) and ( - index.dtype == bool or index.dtype == int - ): - return Detections( - xyxy=self.xyxy[index], - mask=self.mask[index] if self.mask is not None else None, - 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, - ) - raise TypeError( - f"Detections.__getitem__ not supported for index of type {type(index)}." + def __getitem__( + self, index: Union[int, slice, List[int], np.ndarray] + ) -> Detections: + """ + Get a subset of the Detections object. + + Args: + index (Union[int, slice, List[int], np.ndarray]): The index or indices of the subset of the Detections + + Returns: + (Detections): A subset of the Detections object. + + Example: + ```python + >>> import supervision as sv + + >>> detections = sv.Detections(...) + + >>> first_detection = detections[0] + + >>> first_10_detections = detections[0:10] + + >>> some_detections = detections[[0, 2, 4]] + + >>> class_0_detections = detections[detections.class_id == 0] + + >>> high_confidence_detections = detections[detections.confidence > 0.5] + ``` + """ + if isinstance(index, int): + index = [index] + return Detections( + xyxy=self.xyxy[index], + mask=self.mask[index] if self.mask is not None else None, + 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, ) @property diff --git a/test/detection/test_core.py b/test/detection/test_core.py index e24d802f..0ced48a0 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -28,13 +28,6 @@ DETECTIONS = Detections( ) -PREDICTIONS = np.array([ - [ 2254, 906, 2447, 1353, 0.90538, 0], - [ 2049, 1133, 2226, 1371, 0.59002, 56], - [ 727, 1224, 838, 1601, 0.51119, 39], - ], dtype=np.float32) - - def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None) -> Detections: return Detections( xyxy = np.array(xyxy, dtype=np.float32), @@ -101,11 +94,108 @@ def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None) ), DoesNotRaise() ), # take no detections + ( + DETECTIONS, + [0, 2], + Detections( + xyxy=np.array([ + [ 2254, 906, 2447, 1353], + [ 727, 1224, 838, 1601] + ], dtype=np.float32), + confidence=np.array([ + 0.90538, + 0.51119 + ], dtype=np.float32), + class_id=np.array([ + 0, + 39 + ], dtype=int) + ), + DoesNotRaise() + ), # take only first and third detection using List[int] index + ( + DETECTIONS, + np.array([0, 2]), + Detections( + xyxy=np.array([ + [ 2254, 906, 2447, 1353], + [ 727, 1224, 838, 1601] + ], dtype=np.float32), + confidence=np.array([ + 0.90538, + 0.51119 + ], dtype=np.float32), + class_id=np.array([ + 0, + 39 + ], dtype=int) + ), + DoesNotRaise() + ), # take only first and third detection using np.ndarray index + ( + DETECTIONS, + 0, + Detections( + xyxy=np.array([ + [ 2254, 906, 2447, 1353] + ], dtype=np.float32), + confidence=np.array([ + 0.90538 + ], dtype=np.float32), + class_id=np.array([ + 0 + ], dtype=int) + ), + DoesNotRaise() + ), # take only first detection by index + ( + DETECTIONS, + slice(1, 3), + Detections( + xyxy=np.array([ + [ 2049, 1133, 2226, 1371], + [ 727, 1224, 838, 1601] + ], dtype=np.float32), + confidence=np.array([ + 0.59002, + 0.51119 + ], dtype=np.float32), + class_id=np.array([ + 56, + 39 + ], dtype=int) + ), + 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, + np.array([0, 2, 10]), + None, + pytest.raises(IndexError) + ), + ( + DETECTIONS, + np.array([True, True, True, True, True, True, True, True, True, True, True]), + None, + pytest.raises(IndexError) + ) ] ) def test_getitem( detections: Detections, - index: Union[int, slice, np.ndarray], + index: Union[int, slice, List[int], np.ndarray], expected_result: Optional[Detections], exception: Exception ) -> None: