diff --git a/supervision/__init__.py b/supervision/__init__.py index fa6e6044..31d01314 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -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, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 895b8933..603cec5d 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -6,10 +6,9 @@ import numpy as np from supervision.annotators.base import BaseAnnotator from supervision.annotators.utils import ( - ColorMap, + ColorLookup, Trace, resolve_color, - resolve_color_idx, ) from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette @@ -25,27 +24,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 +75,13 @@ 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 +101,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 +154,13 @@ 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 +180,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 +232,13 @@ 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 +262,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 +322,13 @@ 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,13 @@ 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 +442,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 +450,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 +495,13 @@ 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 +525,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 +553,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 +581,13 @@ 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 +612,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 +624,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 +633,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 +682,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 +691,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 +717,13 @@ 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 +837,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 +848,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 +870,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 +897,13 @@ 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( diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 40a8b7c0..5c8e594a 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -8,18 +8,19 @@ 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" CLASS = "class" TRACK = "track" 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 +28,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 +53,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, diff --git a/test/annotators/test_utils.py b/test/annotators/test_utils.py index 7e883a73..e119477f 100644 --- a/test/annotators/test_utils.py +++ b/test/annotators/test_utils.py @@ -1,87 +1,104 @@ from contextlib import ExitStack as DoesNotRaise + +import numpy as np + from test.utils import mock_detections from typing import Optional 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( + mock_detections( xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[5, 3], tracker_id=[2, 6], ), - 0, - ColorMap.INDEX, - 0, - DoesNotRaise(), - ), # multiple detections; index mapping + 0, + ColorLookup.INDEX, + 0, + DoesNotRaise(), + ), # multiple detections; index lookup ( - mock_detections( + mock_detections( xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[5, 3], tracker_id=[2, 6], ), - 0, - ColorMap.CLASS, - 5, - DoesNotRaise(), - ), # multiple detections; class mapping + 0, + ColorLookup.CLASS, + 5, + DoesNotRaise(), + ), # multiple detections; class lookup ( - mock_detections( + mock_detections( xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[5, 3], tracker_id=[2, 6], ), - 0, - ColorMap.TRACK, - 2, - DoesNotRaise(), - ), # multiple detections; track mapping + 0, + ColorLookup.TRACK, + 2, + DoesNotRaise(), + ), # multiple detections; track lookup ( - Detections.empty(), - 0, - ColorMap.INDEX, - None, - pytest.raises(ValueError), - ), # no detections; index mapping; out of bounds + Detections.empty(), + 0, + ColorLookup.INDEX, + None, + pytest.raises(ValueError), + ), # no detections; index lookup; out of bounds ( - mock_detections( + mock_detections( xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[5, 3], tracker_id=[2, 6], ), - 2, - ColorMap.INDEX, - None, - pytest.raises(ValueError), - ), # multiple detections; index mapping; out of bounds + 2, + ColorLookup.INDEX, + None, + pytest.raises(ValueError), + ), # multiple detections; index lookup; out of bounds ( - mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), - 0, - ColorMap.CLASS, - None, - pytest.raises(ValueError), - ), # multiple detections; class mapping; no class_id + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + 0, + ColorLookup.CLASS, + None, + pytest.raises(ValueError), + ), # multiple detections; class lookup; no class_id ( - mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), - 0, - ColorMap.TRACK, - None, - pytest.raises(ValueError), - ), # multiple detections; class mapping; no track_id + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + 0, + ColorLookup.TRACK, + None, + pytest.raises(ValueError), + ), # 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 +106,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