diff --git a/supervision/__init__.py b/supervision/__init__.py index ab45651a..2a3a642d 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -29,6 +29,7 @@ from supervision.annotators.core import ( RichLabelAnnotator, RoundBoxAnnotator, TraceAnnotator, + SplineAnnotator, TriangleAnnotator, ) from supervision.annotators.utils import ColorLookup diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index b77d2c0b..09073728 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -7,6 +7,7 @@ import cv2 import numpy as np import numpy.typing as npt from PIL import Image, ImageDraw, ImageFont +from scipy.interpolate import splprep, splev from supervision.annotators.base import BaseAnnotator, ImageType from supervision.annotators.utils import ( @@ -1908,6 +1909,7 @@ class TraceAnnotator(BaseAnnotator): else custom_color_lookup, ) xy = self.trace.get(tracker_id=tracker_id) + if len(xy) > 1: scene = cv2.polylines( scene, @@ -1918,6 +1920,86 @@ class TraceAnnotator(BaseAnnotator): ) return scene +class SplineAnnotator(BaseAnnotator): + """ + A class for drawing trace paths on an image based on detection coordinates. + + !!! warning + + This annotator uses the `sv.Detections.tracker_id`. Read + [here](/latest/trackers/) to learn how to plug + tracking into your inference pipeline. + """ + + def __init__( + self, + color: Color | ColorPalette = ColorPalette.DEFAULT, + position: Position = Position.CENTER, + trace_length: int = 30, + thickness: int = 2, + color_lookup: ColorLookup = ColorLookup.CLASS, + ): + """ + Args: + color (Union[Color, ColorPalette]): The color to draw the trace, can be + a single color or a color palette. + position (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`. + color_lookup (ColorLookup): Strategy for mapping colors to annotations. + Options are `INDEX`, `CLASS`, `TRACK`. + """ + self.color: Color | ColorPalette = color + self.trace = Trace(max_size=trace_length, anchor=position) + self.thickness = thickness + self.color_lookup: ColorLookup = color_lookup + + @ensure_cv2_image_for_annotation + def annotate(self, scene: ImageType, detections: Detections, custom_color_lookup: np.ndarray | None = None) -> ImageType: + assert isinstance(scene, np.ndarray) + + if detections.tracker_id is None: + raise ValueError( + "The `tracker_id` field is missing in the provided detections." + " See more: https://supervision.roboflow.com/latest/how_to/track_objects" + ) + + detections = detections[detections.tracker_id != PENDING_TRACK_ID] + + self.trace.put(detections) + + for detection_idx in range(len(detections)): + tracker_id = int(detections.tracker_id[detection_idx]) + color = resolve_color( + color=self.color, + detections=detections, + detection_idx=detection_idx, + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, + ) + xy = self.trace.get(tracker_id=tracker_id) + spline_points = None; + + if len(xy) > 3: + x, y = xy[:, 0], xy[:, 1] + tck, u = splprep([x, y], s=100) + x_new, y_new = splev(np.linspace(0, 1, 100), tck) + spline_points = np.stack([x_new, y_new], axis=1).astype(np.int32) + else: + spline_points = xy.astype(np.int32) + + if len(xy) > 1: + scene = cv2.polylines( + scene, + [spline_points], + False, + color=color.as_bgr(), + thickness=self.thickness, + ) + return scene class HeatMapAnnotator(BaseAnnotator): """