xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge branch 'roboflow:develop' into develop

This commit is contained in:
Rajarshi Misra 2023-10-14 13:55:55 +05:30 committed by GitHub
parent 3dbbcf79ae 9ded70c8da
commit e71232f6ff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 209 additions and 111 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
supervision/__init__.py
View File

@ -18,7 +18,7 @@ from supervision.annotators.core import (
MaskAnnotator,
TraceAnnotator,
)
from supervision.annotators.utils import ColorMap
from supervision.annotators.utils import ColorLookup
from supervision.classification.core import Classifications
from supervision.dataset.core import (
BaseDataset,

230
supervision/annotators/core.py
View File

@ -5,12 +5,7 @@ import cv2
import numpy as np
from supervision.annotators.base import BaseAnnotator
from supervision.annotators.utils import (
ColorMap,
Trace,
resolve_color,
resolve_color_idx,
)
from supervision.annotators.utils import ColorLookup, Trace, resolve_color
from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette
from supervision.geometry.core import Position
@ -25,27 +20,34 @@ class BoundingBoxAnnotator(BaseAnnotator):
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 2,
color_map: str = "class",
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
color (Union[Color, ColorPalette]): The color or color palette to use for
annotating detections.
thickness (int): Thickness of the bounding box lines.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.color_map: ColorMap = ColorMap(color_map)
self.color_lookup: ColorLookup = color_lookup
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
def annotate(
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with bounding boxes based on the provided detections.
Args:
scene (np.ndarray): The image where bounding boxes will be drawn.
detections (Detections): Object detections to annotate.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
np.ndarray: The annotated image.
@ -69,12 +71,14 @@ class BoundingBoxAnnotator(BaseAnnotator):
"""
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
idx = resolve_color_idx(
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
color = resolve_color(color=self.color, idx=idx)
cv2.rectangle(
img=scene,
pt1=(x1, y1),
@ -94,27 +98,34 @@ class MaskAnnotator(BaseAnnotator):
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
opacity: float = 0.5,
color_map: str = "class",
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
color (Union[Color, ColorPalette]): The color or color palette to use for
annotating detections.
opacity (float): Opacity of the overlay mask. Must be between `0` and `1`.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.opacity = opacity
self.color_map: ColorMap = ColorMap(color_map)
self.color_lookup: ColorLookup = color_lookup
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
def annotate(
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with masks based on the provided detections.
Args:
scene (np.ndarray): The image where masks will be drawn.
detections (Detections): Object detections to annotate.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
np.ndarray: The annotated image.
@ -140,12 +151,14 @@ class MaskAnnotator(BaseAnnotator):
return scene
for detection_idx in np.flip(np.argsort(detections.area)):
idx = resolve_color_idx(
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
color = resolve_color(color=self.color, idx=idx)
mask = detections.mask[detection_idx]
colored_mask = np.zeros_like(scene, dtype=np.uint8)
colored_mask[:] = color.as_bgr()
@ -165,27 +178,34 @@ class BoxMaskAnnotator(BaseAnnotator):
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
opacity: float = 0.5,
color_map: str = "class",
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
color (Union[Color, ColorPalette]): The color or color palette to use for
annotating detections.
opacity (float): Opacity of the overlay mask. Must be between `0` and `1`.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.color_map: ColorMap = ColorMap(color_map)
self.color_lookup: ColorLookup = color_lookup
self.opacity = opacity
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
def annotate(
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with box masks based on the provided detections.
Args:
scene (np.ndarray): The image where bounding boxes will be drawn.
detections (Detections): Object detections to annotate.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
np.ndarray: The annotated image.
@ -210,12 +230,14 @@ class BoxMaskAnnotator(BaseAnnotator):
mask_image = scene.copy()
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
idx = resolve_color_idx(
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
color = resolve_color(color=self.color, idx=idx)
cv2.rectangle(
img=scene,
pt1=(x1, y1),
@ -239,30 +261,37 @@ class HaloAnnotator(BaseAnnotator):
color: Union[Color, ColorPalette] = ColorPalette.default(),
opacity: float = 0.8,
kernel_size: int = 40,
color_map: str = "class",
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
color (Union[Color, ColorPalette]): The color or color palette to use for
annotating detections.
opacity (float): Opacity of the overlay mask. Must be between `0` and `1`.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
kernel_size (int): The size of the average pooling kernel used for creating
the halo.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.opacity = opacity
self.color_map: ColorMap = ColorMap(color_map)
self.color_lookup: ColorLookup = color_lookup
self.kernel_size: int = kernel_size
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
def annotate(
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with halos based on the provided detections.
Args:
scene (np.ndarray): The image where masks will be drawn.
detections (Detections): Object detections to annotate.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
np.ndarray: The annotated image.
@ -292,12 +321,14 @@ class HaloAnnotator(BaseAnnotator):
)
for detection_idx in np.flip(np.argsort(detections.area)):
idx = resolve_color_idx(
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
color = resolve_color(color=self.color, idx=idx)
mask = detections.mask[detection_idx]
fmask = np.logical_or(fmask, mask)
color_bgr = color.as_bgr()
@ -323,7 +354,7 @@ class EllipseAnnotator(BaseAnnotator):
thickness: int = 2,
start_angle: int = -45,
end_angle: int = 235,
color_map: str = "class",
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
@ -332,22 +363,29 @@ class EllipseAnnotator(BaseAnnotator):
thickness (int): Thickness of the ellipse lines.
start_angle (int): Starting angle of the ellipse.
end_angle (int): Ending angle of the ellipse.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.start_angle: int = start_angle
self.end_angle: int = end_angle
self.color_map: ColorMap = ColorMap(color_map)
self.color_lookup: ColorLookup = color_lookup
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
def annotate(
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with ellipses based on the provided detections.
Args:
scene (np.ndarray): The image where ellipses will be drawn.
detections (Detections): Object detections to annotate.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
np.ndarray: The annotated image.
@ -371,13 +409,14 @@ class EllipseAnnotator(BaseAnnotator):
"""
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
idx = resolve_color_idx(
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
color = resolve_color(color=self.color, idx=idx)
center = (int((x1 + x2) / 2), y2)
width = x2 - x1
cv2.ellipse(
@ -404,7 +443,7 @@ class BoxCornerAnnotator(BaseAnnotator):
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 4,
corner_length: int = 15,
color_map: str = "class",
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
@ -412,21 +451,28 @@ class BoxCornerAnnotator(BaseAnnotator):
annotating detections.
thickness (int): Thickness of the corner lines.
corner_length (int): Length of each corner line.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.corner_length: int = corner_length
self.color_map: ColorMap = ColorMap(color_map)
self.color_lookup: ColorLookup = color_lookup
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
def annotate(
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with box corners based on the provided detections.
Args:
scene (np.ndarray): The image where box corners will be drawn.
detections (Detections): Object detections to annotate.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
np.ndarray: The annotated image.
@ -450,12 +496,14 @@ class BoxCornerAnnotator(BaseAnnotator):
"""
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
idx = resolve_color_idx(
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
color = resolve_color(color=self.color, idx=idx)
corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
for x, y in corners:
@ -479,26 +527,27 @@ class CircleAnnotator(BaseAnnotator):
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 4,
color_map: str = "class",
thickness: int = 2,
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
color (Union[Color, ColorPalette]): The color or color palette to use for
annotating detections.
thickness (int): Thickness of the circle line.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.color_map: ColorMap = ColorMap(color_map)
self.color_lookup: ColorLookup = color_lookup
def annotate(
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with circles based on the provided detections.
@ -506,6 +555,8 @@ class CircleAnnotator(BaseAnnotator):
Args:
scene (np.ndarray): The image where box corners will be drawn.
detections (Detections): Object detections to annotate.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
np.ndarray: The annotated image.
@ -532,19 +583,14 @@ class CircleAnnotator(BaseAnnotator):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
center = ((x1 + x2) // 2, (y1 + y2) // 2)
distance = sqrt((x1 - center[0]) ** 2 + (y1 - center[1]) ** 2)
idx = resolve_color_idx(
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
color = (
self.color.by_idx(idx)
if isinstance(self.color, ColorPalette)
else self.color
)
cv2.circle(
img=scene,
center=center,
@ -569,7 +615,7 @@ class LabelAnnotator:
text_thickness: int = 1,
text_padding: int = 10,
text_position: Position = Position.TOP_LEFT,
color_map: str = "class",
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
@ -581,8 +627,8 @@ class LabelAnnotator:
text_padding (int): Padding around the text within its background box.
text_position (Position): Position of the text relative to the detection.
Possible values are defined in the `Position` enum.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.text_color: Color = text_color
@ -590,7 +636,7 @@ class LabelAnnotator:
self.text_thickness: int = text_thickness
self.text_padding: int = text_padding
self.text_position: Position = text_position
self.color_map: ColorMap = ColorMap(color_map)
self.color_lookup: ColorLookup = color_lookup
@staticmethod
def resolve_text_background_xyxy(
@ -639,6 +685,7 @@ class LabelAnnotator:
scene: np.ndarray,
detections: Detections,
labels: List[str] = None,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with labels based on the provided detections.
@ -647,6 +694,8 @@ class LabelAnnotator:
scene (np.ndarray): The image where labels will be drawn.
detections (Detections): Object detections to annotate.
labels (List[str]): Optional. Custom labels for each detection.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
np.ndarray: The annotated image.
@ -671,12 +720,14 @@ class LabelAnnotator:
font = cv2.FONT_HERSHEY_SIMPLEX
for detection_idx in range(len(detections)):
detection_xyxy = detections.xyxy[detection_idx].astype(int)
idx = resolve_color_idx(
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
color = resolve_color(color=self.color, idx=idx)
text = (
f"{detections.class_id[detection_idx]}"
if (labels is None or len(detections) != len(labels))
@ -790,7 +841,7 @@ class TraceAnnotator:
position: Optional[Position] = Position.CENTER,
trace_length: int = 30,
thickness: int = 2,
color_map: str = "class",
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
@ -801,16 +852,21 @@ class TraceAnnotator:
trace_length (int): The maximum length of the trace in terms of historical
points. Defaults to `30`.
thickness (int): The thickness of the trace lines. Defaults to `2`.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.position = position
self.trace = Trace(max_size=trace_length)
self.thickness = thickness
self.color_map: ColorMap = ColorMap(color_map)
self.color_lookup: ColorLookup = color_lookup
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
def annotate(
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Draws trace paths on the frame based on the detection coordinates provided.
@ -818,6 +874,8 @@ class TraceAnnotator:
scene (np.ndarray): The image on which the traces will be drawn.
detections (Detections): The detections which include coordinates for
which the traces will be drawn.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
np.ndarray: The image with the trace paths drawn on it.
@ -843,12 +901,14 @@ class TraceAnnotator:
for detection_idx in range(len(detections)):
tracker_id = int(detections.tracker_id[detection_idx])
idx = resolve_color_idx(
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
color = resolve_color(color=self.color, idx=idx)
xy = self.trace.get(tracker_id=tracker_id)
if len(xy) > 1:
scene = cv2.polylines(

37
supervision/annotators/utils.py
View File

@ -8,9 +8,9 @@ from supervision.draw.color import Color, ColorPalette
from supervision.geometry.core import Position
class ColorMap(Enum):
class ColorLookup(Enum):
"""
Enum for annotator color mapping.
Enum for annotator color lookup.
"""
INDEX = "index"
@ -19,7 +19,9 @@ class ColorMap(Enum):
def resolve_color_idx(
detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS
detections: Detections,
detection_idx: int,
color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS,
) -> int:
if detection_idx >= len(detections):
raise ValueError(
@ -27,16 +29,23 @@ def resolve_color_idx(
f"is out of bounds for detections of length {len(detections)}"
)
if color_map == ColorMap.INDEX:
if isinstance(color_lookup, np.ndarray):
if len(color_lookup) != len(detections):
raise ValueError(
f"Length of color lookup {len(color_lookup)}"
f"does not match length of detections {len(detections)}"
)
return color_lookup[detection_idx]
elif color_lookup == ColorLookup.INDEX:
return detection_idx
elif color_map == ColorMap.CLASS:
elif color_lookup == ColorLookup.CLASS:
if detections.class_id is None:
raise ValueError(
"Could not resolve color by class because"
"Detections do not have class_id"
)
return detections.class_id[detection_idx]
elif color_map == ColorMap.TRACK:
elif color_lookup == ColorLookup.TRACK:
if detections.tracker_id is None:
raise ValueError(
"Could not resolve color by track because"
@ -45,12 +54,26 @@ def resolve_color_idx(
return detections.tracker_id[detection_idx]
def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color:
def get_color_by_index(color: Union[Color, ColorPalette], idx: int) -> Color:
if isinstance(color, ColorPalette):
return color.by_idx(idx)
return color
def resolve_color(
color: Union[Color, ColorPalette],
detections: Detections,
detection_idx: int,
color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS,
) -> Color:
idx = resolve_color_idx(
detections=detections,
detection_idx=detection_idx,
color_lookup=color_lookup,
)
return get_color_by_index(color=color, idx=idx)
class Trace:
def __init__(
self,

51
test/annotators/test_utils.py
View File

@ -2,14 +2,15 @@ from contextlib import ExitStack as DoesNotRaise
from test.utils import mock_detections
from typing import Optional
import numpy as np
import pytest
from supervision.annotators.utils import ColorMap, resolve_color_idx
from supervision.annotators.utils import ColorLookup, resolve_color_idx
from supervision.detection.core import Detections
@pytest.mark.parametrize(
"detections, detection_idx, color_map, expected_result, exception",
"detections, detection_idx, color_lookup, expected_result, exception",
[
(
mock_detections(
@ -18,10 +19,10 @@ from supervision.detection.core import Detections
tracker_id=[2, 6],
),
0,
ColorMap.INDEX,
ColorLookup.INDEX,
0,
DoesNotRaise(),
), # multiple detections; index mapping
), # multiple detections; index lookup
(
mock_detections(
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
@ -29,10 +30,10 @@ from supervision.detection.core import Detections
tracker_id=[2, 6],
),
0,
ColorMap.CLASS,
ColorLookup.CLASS,
5,
DoesNotRaise(),
), # multiple detections; class mapping
), # multiple detections; class lookup
(
mock_detections(
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
@ -40,17 +41,17 @@ from supervision.detection.core import Detections
tracker_id=[2, 6],
),
0,
ColorMap.TRACK,
ColorLookup.TRACK,
2,
DoesNotRaise(),
), # multiple detections; track mapping
), # multiple detections; track lookup
(
Detections.empty(),
0,
ColorMap.INDEX,
ColorLookup.INDEX,
None,
pytest.raises(ValueError),
), # no detections; index mapping; out of bounds
), # no detections; index lookup; out of bounds
(
mock_detections(
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
@ -58,30 +59,44 @@ from supervision.detection.core import Detections
tracker_id=[2, 6],
),
2,
ColorMap.INDEX,
ColorLookup.INDEX,
None,
pytest.raises(ValueError),
), # multiple detections; index mapping; out of bounds
), # multiple detections; index lookup; out of bounds
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
ColorMap.CLASS,
ColorLookup.CLASS,
None,
pytest.raises(ValueError),
), # multiple detections; class mapping; no class_id
), # multiple detections; class lookup; no class_id
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
ColorMap.TRACK,
ColorLookup.TRACK,
None,
pytest.raises(ValueError),
), # multiple detections; class mapping; no track_id
), # multiple detections; class lookup; no track_id
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
np.array([1, 0]),
1,
DoesNotRaise(),
), # multiple detections; custom lookup; correct length
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
np.array([1]),
None,
pytest.raises(ValueError),
), # multiple detections; custom lookup; wrong length
],
)
def test_resolve_color_idx(
detections: Detections,
detection_idx: int,
color_map: ColorMap,
color_lookup: ColorLookup,
expected_result: Optional[int],
exception: Exception,
) -> None:
@ -89,6 +104,6 @@ def test_resolve_color_idx(
result = resolve_color_idx(
detections=detections,
detection_idx=detection_idx,
color_map=color_map,
color_lookup=color_lookup,
)
assert result == expected_result