From e1e455f99cf3a52b4dc5cd99ac53cf3a5515e7bf Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 12 Oct 2023 19:10:19 +0200 Subject: [PATCH] Update color mapping system in detection annotations Adjusted the way colors are applied to detection annotations to allow support for custom color schemes. Modified 'ColorMap' to 'ColorLookup' across multiple files, adding ability to use a color lookup array. This refinement will enable users to assign custom colors to different detection annotations, offering greater flexibility in visual representation of detection data. Updated related tests to reflect these changes. --- supervision/__init__.py | 2 +- supervision/annotators/core.py | 217 ++++++++++++++++++++------------ supervision/annotators/utils.py | 38 ++++-- test/annotators/test_utils.py | 109 +++++++++------- 4 files changed, 230 insertions(+), 136 deletions(-) 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