From e1e455f99cf3a52b4dc5cd99ac53cf3a5515e7bf Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 12 Oct 2023 19:10:19 +0200 Subject: [PATCH 1/2] 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 fa6e6044b..31d013140 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -18,7 +18,7 @@ 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 895b89331..603cec5d1 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -6,10 +6,9 @@ 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 @@ def __init__( 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 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: """ 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 @@ def __init__( 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 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: 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 @@ def __init__( 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 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: 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 @@ def __init__( 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 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: ) 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 @@ def __init__( 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 @@ def __init__( 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 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: """ 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 @@ def __init__( 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 @@ def __init__( 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 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: """ 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 @@ def annotate( 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 @@ def annotate( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def annotate( 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 @@ def annotate( 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 @@ def annotate( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: 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 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: 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 40a8b7c00..5c8e594af 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -8,18 +8,19 @@ 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 7e883a735..e119477f8 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, + 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, + 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, - 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, + np.array([1, 0]), + 1, + DoesNotRaise(), + ), # multiple detections; custom lookup; correct length ( - 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, + 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 From 476a31fa790df5ec4e0fe6b042062e9b49eb4bab Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 12 Oct 2023 17:11:58 +0000 Subject: [PATCH 2/2] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/annotators/core.py | 69 +++++++++++++----------- supervision/annotators/utils.py | 5 +- test/annotators/test_utils.py | 94 ++++++++++++++++----------------- 3 files changed, 86 insertions(+), 82 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 603cec5d1..8005f49b8 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -5,11 +5,7 @@ import numpy as np from supervision.annotators.base import BaseAnnotator -from supervision.annotators.utils import ( - ColorLookup, - Trace, - resolve_color, -) +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 @@ -42,7 +38,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with bounding boxes based on the provided detections. @@ -79,8 +75,9 @@ def annotate( color=self.color, detections=detections, detection_idx=detection_idx, - color_lookup=self.color_lookup if - custom_color_lookup is None else custom_color_lookup + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, ) cv2.rectangle( img=scene, @@ -119,7 +116,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with masks based on the provided detections. @@ -158,8 +155,9 @@ def annotate( color=self.color, detections=detections, detection_idx=detection_idx, - color_lookup=self.color_lookup if - custom_color_lookup is None else custom_color_lookup + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, ) mask = detections.mask[detection_idx] colored_mask = np.zeros_like(scene, dtype=np.uint8) @@ -198,7 +196,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with box masks based on the provided detections. @@ -236,8 +234,9 @@ def annotate( color=self.color, detections=detections, detection_idx=detection_idx, - color_lookup=self.color_lookup if - custom_color_lookup is None else custom_color_lookup + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, ) cv2.rectangle( img=scene, @@ -283,7 +282,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with halos based on the provided detections. @@ -326,8 +325,9 @@ def annotate( color=self.color, detections=detections, detection_idx=detection_idx, - color_lookup=self.color_lookup if - custom_color_lookup is None else custom_color_lookup + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, ) mask = detections.mask[detection_idx] fmask = np.logical_or(fmask, mask) @@ -376,7 +376,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with ellipses based on the provided detections. @@ -413,8 +413,9 @@ def annotate( color=self.color, detections=detections, detection_idx=detection_idx, - color_lookup=self.color_lookup if - custom_color_lookup is None else custom_color_lookup + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, ) center = (int((x1 + x2) / 2), y2) width = x2 - x1 @@ -462,7 +463,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with box corners based on the provided detections. @@ -499,8 +500,9 @@ def annotate( color=self.color, detections=detections, detection_idx=detection_idx, - color_lookup=self.color_lookup if - custom_color_lookup is None else custom_color_lookup + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, ) corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)] @@ -545,7 +547,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with circles based on the provided detections. @@ -585,8 +587,9 @@ def annotate( color=self.color, detections=detections, detection_idx=detection_idx, - color_lookup=self.color_lookup if - custom_color_lookup is None else custom_color_lookup + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, ) cv2.circle( img=scene, @@ -682,7 +685,7 @@ def annotate( scene: np.ndarray, detections: Detections, labels: List[str] = None, - custom_color_lookup: Optional[np.ndarray] = None + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with labels based on the provided detections. @@ -721,8 +724,9 @@ def annotate( color=self.color, detections=detections, detection_idx=detection_idx, - color_lookup=self.color_lookup if - custom_color_lookup is None else custom_color_lookup + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, ) text = ( f"{detections.class_id[detection_idx]}" @@ -861,7 +865,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Draws trace paths on the frame based on the detection coordinates provided. @@ -901,8 +905,9 @@ def annotate( color=self.color, detections=detections, detection_idx=detection_idx, - color_lookup=self.color_lookup if - custom_color_lookup is None else custom_color_lookup + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, ) xy = self.trace.get(tracker_id=tracker_id) if len(xy) > 1: diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 5c8e594af..6b6b98373 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -12,6 +12,7 @@ class ColorLookup(Enum): """ Enum for annotator color lookup. """ + INDEX = "index" CLASS = "class" TRACK = "track" @@ -20,7 +21,7 @@ class ColorLookup(Enum): def resolve_color_idx( detections: Detections, detection_idx: int, - color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS + color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS, ) -> int: if detection_idx >= len(detections): raise ValueError( @@ -63,7 +64,7 @@ def resolve_color( color: Union[Color, ColorPalette], detections: Detections, detection_idx: int, - color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS + color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS, ) -> Color: idx = resolve_color_idx( detections=detections, diff --git a/test/annotators/test_utils.py b/test/annotators/test_utils.py index e119477f8..98d9be6b2 100644 --- a/test/annotators/test_utils.py +++ b/test/annotators/test_utils.py @@ -1,10 +1,8 @@ from contextlib import ExitStack as DoesNotRaise - -import numpy as np - from test.utils import mock_detections from typing import Optional +import numpy as np import pytest from supervision.annotators.utils import ColorLookup, resolve_color_idx @@ -15,83 +13,83 @@ "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, - ColorLookup.INDEX, - 0, - DoesNotRaise(), + 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, - ColorLookup.CLASS, - 5, - DoesNotRaise(), + 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, - ColorLookup.TRACK, - 2, - DoesNotRaise(), + 0, + ColorLookup.TRACK, + 2, + DoesNotRaise(), ), # multiple detections; track lookup ( - Detections.empty(), - 0, - ColorLookup.INDEX, - None, - pytest.raises(ValueError), + 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, - ColorLookup.INDEX, - None, - pytest.raises(ValueError), + 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, - ColorLookup.CLASS, - None, - pytest.raises(ValueError), + 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, - ColorLookup.TRACK, - None, - pytest.raises(ValueError), + 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(), + 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), + 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 ], )