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 7aa4dab0..578c358a 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -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( diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 40a8b7c0..6b6b9837 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -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, diff --git a/test/annotators/test_utils.py b/test/annotators/test_utils.py index 7e883a73..98d9be6b 100644 --- a/test/annotators/test_utils.py +++ b/test/annotators/test_utils.py @@ -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