Merge pull request #118 from roboflow/feature/115_how_to_select_detections_by_index
feature/115_how_to_select_detections_by_index
This commit is contained in:
2 files changed
+136
-26
No files matched your search
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in new issue
Block a user