From 97892e3138667f4efb11739695c4f6c4f5bcbe6c Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 5 Jun 2023 11:01:14 +0200 Subject: [PATCH 1/2] =?UTF-8?q?=E2=9B=8F=EF=B8=8Fpick=20`Detections`=20by?= =?UTF-8?q?=20index?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 28 ++++----- test/detection/test_core.py | 106 +++++++++++++++++++++++++++++++--- 2 files changed, 108 insertions(+), 26 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 75784ee8..d3901418 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,15 @@ 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: + 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: From 0d38b95da7d8a3ac4df1806a2cf1c83048078468 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 5 Jun 2023 11:31:52 +0200 Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=93=9Aupdate=20documentation=20+=20?= =?UTF-8?q?=F0=9F=96=A4=20make=20black=20happy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index d3901418..dc88e9cc 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -469,7 +469,35 @@ class Detections: raise ValueError(f"{anchor} is not supported.") - def __getitem__(self, index: Union[int, slice, List[int], np.ndarray]) -> Detections: + 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(