diff --git a/docs/annotators.md b/docs/annotators.md index 2077e58d..11bdfc17 100644 --- a/docs/annotators.md +++ b/docs/annotators.md @@ -124,6 +124,27 @@ +=== "Trace" + + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> trace_annotator = sv.TraceAnnotator() + >>> annotated_frame = trace_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + +
+ + ![trace-annotator-example](https://media.roboflow.com/supervision-annotator-examples/trace-annotator-example.png){ align=center width="800" } + +
+ ## BoundingBoxAnnotator :::supervision.annotators.core.BoundingBoxAnnotator @@ -147,3 +168,7 @@ ## LabelAnnotator :::supervision.annotators.core.LabelAnnotator + +## TraceAnnotator + +:::supervision.annotators.core.TraceAnnotator diff --git a/supervision/__init__.py b/supervision/__init__.py index 7ad76847..aab1eda8 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -13,6 +13,7 @@ from supervision.annotators.core import ( EllipseAnnotator, LabelAnnotator, MaskAnnotator, + TraceAnnotator, ) from supervision.classification.core import Classifications from supervision.dataset.core import ( @@ -20,7 +21,7 @@ from supervision.dataset.core import ( ClassificationDataset, DetectionDataset, ) -from supervision.detection.annotate import BoxAnnotator, TraceAnnotator +from supervision.detection.annotate import BoxAnnotator from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator from supervision.detection.tools.inference_slicer import InferenceSlicer diff --git a/supervision/annotators/base.py b/supervision/annotators/base.py index d9a6e29a..515ee1c8 100644 --- a/supervision/annotators/base.py +++ b/supervision/annotators/base.py @@ -1,55 +1,8 @@ from abc import ABC, abstractmethod -from enum import Enum -from typing import Union import numpy as np from supervision.detection.core import Detections -from supervision.draw.color import Color, ColorPalette - - -class ColorMap(Enum): - """ - Enum for annotator color mapping. - """ - - INDEX = "index" - CLASS = "class" - TRACK = "track" - - -def resolve_color_idx( - detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS -) -> int: - if detection_idx >= len(detections): - raise ValueError( - f"Detection index {detection_idx}" - f"is out of bounds for detections of length {len(detections)}" - ) - - if color_map == ColorMap.INDEX: - return detection_idx - elif color_map == ColorMap.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: - if detections.tracker_id is None: - raise ValueError( - "Could not resolve color by track because" - "Detections do not have tracker_id" - ) - return detections.tracker_id[detection_idx] - - -def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color: - if isinstance(color, ColorPalette): - return color.by_idx(idx) - else: - return color class BaseAnnotator(ABC): diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index d108af8e..350ae06b 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,12 +1,13 @@ from math import sqrt -from typing import List, Tuple, Union +from typing import List, Optional, Tuple, Union import cv2 import numpy as np -from supervision.annotators.base import ( - BaseAnnotator, +from supervision.annotators.base import BaseAnnotator +from supervision.annotators.utils import ( ColorMap, + Trace, resolve_color, resolve_color_idx, ) @@ -562,3 +563,88 @@ class LabelAnnotator: lineType=cv2.LINE_AA, ) return scene + + +class TraceAnnotator: + """ + A class for drawing trace paths on an image based on detection coordinates. + + !!! warning + + This annotator utilizes the `tracker_id`. Read + [here](https://supervision.roboflow.com/trackers/) to learn how to plug + tracking into your inference pipeline. + """ + + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + position: Optional[Position] = Position.CENTER, + trace_length: int = 30, + thickness: int = 2, + ): + """ + Args: + color (Union[Color, ColorPalette]): The color to draw the trace, can be + a single color or a color palette. + position (Optional[Position]): The position of the trace. + Defaults to `CENTER`. + 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`. + """ + self.color: Union[Color, ColorPalette] = color + self.position = position + self.trace = Trace(max_size=trace_length) + self.thickness = thickness + + def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: + """ + Draws trace paths on the frame based on the detection coordinates provided. + + Args: + 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. + + Returns: + np.ndarray: The image with the trace paths drawn on it. + + Example: + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> trace_annotator = sv.TraceAnnotator() + >>> annotated_frame = trace_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + + ![trace-annotator-example](https://media.roboflow.com/ + supervision-annotator-examples/trace-annotator-example.png) + """ + self.trace.put(detections) + + for i, (xyxy, mask, confidence, class_id, tracker_id) in enumerate(detections): + class_id = detections.class_id[i] if class_id is not None else None + idx = class_id if class_id is not None else i + color = ( + self.color.by_idx(idx) + if isinstance(self.color, ColorPalette) + else self.color + ) + + xy = self.trace.get(tracker_id=tracker_id) + if len(xy) > 1: + scene = cv2.polylines( + scene, + [xy.astype(np.int32)], + False, + color=color.as_bgr(), + thickness=self.thickness, + ) + return scene diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py new file mode 100644 index 00000000..40a8b7c0 --- /dev/null +++ b/supervision/annotators/utils.py @@ -0,0 +1,89 @@ +from enum import Enum +from typing import Optional, Union + +import numpy as np + +from supervision.detection.core import Detections +from supervision.draw.color import Color, ColorPalette +from supervision.geometry.core import Position + + +class ColorMap(Enum): + """ + Enum for annotator color mapping. + """ + + INDEX = "index" + CLASS = "class" + TRACK = "track" + + +def resolve_color_idx( + detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS +) -> int: + if detection_idx >= len(detections): + raise ValueError( + f"Detection index {detection_idx}" + f"is out of bounds for detections of length {len(detections)}" + ) + + if color_map == ColorMap.INDEX: + return detection_idx + elif color_map == ColorMap.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: + if detections.tracker_id is None: + raise ValueError( + "Could not resolve color by track because" + "Detections do not have tracker_id" + ) + return detections.tracker_id[detection_idx] + + +def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color: + if isinstance(color, ColorPalette): + return color.by_idx(idx) + return color + + +class Trace: + def __init__( + self, + max_size: Optional[int] = None, + start_frame_id: int = 0, + anchor: Position = Position.CENTER, + ) -> None: + self.current_frame_id = start_frame_id + self.max_size = max_size + self.anchor = anchor + + self.frame_id = np.array([], dtype=int) + self.xy = np.empty((0, 2), dtype=np.float32) + self.tracker_id = np.array([], dtype=int) + + def put(self, detections: Detections) -> None: + frame_id = np.full(len(detections), self.current_frame_id, dtype=int) + self.frame_id = np.concatenate([self.frame_id, frame_id]) + self.xy = np.concatenate( + [self.xy, detections.get_anchor_coordinates(self.anchor)] + ) + self.tracker_id = np.concatenate([self.tracker_id, detections.tracker_id]) + + unique_frame_id = np.unique(self.frame_id) + + if 0 < self.max_size < len(unique_frame_id): + max_allowed_frame_id = self.current_frame_id - self.max_size + 1 + filtering_mask = self.frame_id >= max_allowed_frame_id + self.frame_id = self.frame_id[filtering_mask] + self.xy = self.xy[filtering_mask] + self.tracker_id = self.tracker_id[filtering_mask] + + self.current_frame_id += 1 + + def get(self, tracker_id: int) -> np.ndarray: + return self.xy[self.tracker_id == tracker_id] diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index 83b6e0f8..5e9b379a 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -5,7 +5,6 @@ import numpy as np from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette -from supervision.geometry.core import Position class BoxAnnotator: @@ -145,116 +144,3 @@ class BoxAnnotator: lineType=cv2.LINE_AA, ) return scene - - -class Trace: - def __init__( - self, - max_size: Optional[int] = None, - start_frame_id: int = 0, - anchor: Position = Position.CENTER, - ) -> None: - self.current_frame_id = start_frame_id - self.max_size = max_size - self.anchor = anchor - - self.frame_id = np.array([], dtype=int) - self.xy = np.empty((0, 2), dtype=np.float32) - self.tracker_id = np.array([], dtype=int) - - def put(self, detections: Detections) -> None: - frame_id = np.full(len(detections), self.current_frame_id, dtype=int) - self.frame_id = np.concatenate([self.frame_id, frame_id]) - self.xy = np.concatenate( - [self.xy, detections.get_anchor_coordinates(self.anchor)] - ) - self.tracker_id = np.concatenate([self.tracker_id, detections.tracker_id]) - - unique_frame_id = np.unique(self.frame_id) - - if 0 < self.max_size < len(unique_frame_id): - max_allowed_frame_id = self.current_frame_id - self.max_size + 1 - filtering_mask = self.frame_id >= max_allowed_frame_id - self.frame_id = self.frame_id[filtering_mask] - self.xy = self.xy[filtering_mask] - self.tracker_id = self.tracker_id[filtering_mask] - - self.current_frame_id += 1 - - def get(self, tracker_id: int) -> np.ndarray: - return self.xy[self.tracker_id == tracker_id] - - -class TraceAnnotator: - """ - A class for drawing trace paths on an image based on detection coordinates. - - Attributes: - color (Union[Color, ColorPalette]): The color to draw the trace, can be - a single color or a color palette. - position (Optional[Position]): The position of the trace. Defaults to `CENTER`. - 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`. - - """ - - def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - position: Optional[Position] = Position.CENTER, - trace_length: int = 30, - thickness: int = 2, - ): - self.color: Union[Color, ColorPalette] = color - self.position = position - self.trace = Trace(max_size=trace_length) - self.thickness = thickness - - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - """ - Draws trace paths on the frame based on the detection coordinates provided. - - Args: - 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. - - Returns: - np.ndarray: The image with the trace paths drawn on it. - - Example: - ```python - >>> import supervision as sv - - >>> image = ... - >>> detections = sv.Detections(...) - - >>> trace_annotator = sv.TraceAnnotator() - >>> annotated_frame = trace_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) - ``` - """ - self.trace.put(detections) - - for i, (xyxy, mask, confidence, class_id, tracker_id) in enumerate(detections): - class_id = detections.class_id[i] if class_id is not None else None - idx = class_id if class_id is not None else i - color = ( - self.color.by_idx(idx) - if isinstance(self.color, ColorPalette) - else self.color - ) - - xy = self.trace.get(tracker_id=tracker_id) - if len(xy) > 1: - scene = cv2.polylines( - scene, - [xy.astype(np.int32)], - False, - color=color.as_bgr(), - thickness=self.thickness, - ) - return scene diff --git a/test/annotators/test_base.py b/test/annotators/test_utils.py similarity index 97% rename from test/annotators/test_base.py rename to test/annotators/test_utils.py index 73d5a147..7e883a73 100644 --- a/test/annotators/test_base.py +++ b/test/annotators/test_utils.py @@ -4,7 +4,7 @@ from typing import Optional import pytest -from supervision.annotators.core import ColorMap, resolve_color_idx +from supervision.annotators.utils import ColorMap, resolve_color_idx from supervision.detection.core import Detections