From 22b6d3d464fd23033a94f925f56eea852e2dce6b Mon Sep 17 00:00:00 2001 From: Ashp116 Date: Wed, 30 Jul 2025 02:04:24 -0400 Subject: [PATCH 1/7] ADD: Added a spline annotator --- supervision/__init__.py | 1 + supervision/annotators/core.py | 82 ++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+) 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): """ From 8c6ac2dd2970fe12565f6ad83beb767e4641ba57 Mon Sep 17 00:00:00 2001 From: Ashp116 Date: Wed, 30 Jul 2025 02:53:58 -0400 Subject: [PATCH 2/7] UPDATE: Made smoothing and spline_order as hyperparameters --- supervision/annotators/core.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 09073728..61632a04 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1937,25 +1937,34 @@ class SplineAnnotator(BaseAnnotator): position: Position = Position.CENTER, trace_length: int = 30, thickness: int = 2, + smoothing_factor: int = 20, + spline_order: int = 3, color_lookup: ColorLookup = ColorLookup.CLASS, ): """ Args: - color (Union[Color, ColorPalette]): The color to draw the trace, can be + color (Union[Color, ColorPalette]): The color to draw the spline, can be a single color or a color palette. - position (Position): The position of the trace. + position (Position): The position of the spline. Defaults to `CENTER`. - trace_length (int): The maximum length of the trace in terms of historical + trace_length (int): The maximum length of the spline in terms of historical points. Defaults to `30`. - thickness (int): The thickness of the trace lines. Defaults to `2`. + thickness (int): The thickness of the spline lines. Defaults to `2`. + smoothing_factor (int): The smoothing factor of the spline. + spline_order (int): The order of the spline. Use odd numbers that are 1 <= x <= 5. 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.smoothing_factor = smoothing_factor self.color_lookup: ColorLookup = color_lookup + if spline_order % 2 == 0 or spline_order < 1 or spline_order > 5: + raise ValueError("Spline order must be an odd number between 1 and 5.") + self.spline_order = spline_order + @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) @@ -1985,7 +1994,7 @@ class SplineAnnotator(BaseAnnotator): if len(xy) > 3: x, y = xy[:, 0], xy[:, 1] - tck, u = splprep([x, y], s=100) + tck, u = splprep([x, y], s=self.smoothing_factor, k=self.spline_order) 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: From f447b6fa95690f58e4c95dded17754605eb9dd39 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 30 Jul 2025 07:10:13 +0000 Subject: [PATCH 3/7] =?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/__init__.py | 2 +- supervision/annotators/core.py | 17 ++++++++++++----- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 2a3a642d..1b2266cd 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -28,8 +28,8 @@ from supervision.annotators.core import ( PolygonAnnotator, RichLabelAnnotator, RoundBoxAnnotator, - TraceAnnotator, SplineAnnotator, + TraceAnnotator, TriangleAnnotator, ) from supervision.annotators.utils import ColorLookup diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 61632a04..b1ecc674 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -7,7 +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 scipy.interpolate import splev, splprep from supervision.annotators.base import BaseAnnotator, ImageType from supervision.annotators.utils import ( @@ -1920,6 +1920,7 @@ class TraceAnnotator(BaseAnnotator): ) return scene + class SplineAnnotator(BaseAnnotator): """ A class for drawing trace paths on an image based on detection coordinates. @@ -1966,7 +1967,12 @@ class SplineAnnotator(BaseAnnotator): self.spline_order = spline_order @ensure_cv2_image_for_annotation - def annotate(self, scene: ImageType, detections: Detections, custom_color_lookup: np.ndarray | None = None) -> ImageType: + 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: @@ -1974,11 +1980,11 @@ class SplineAnnotator(BaseAnnotator): "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( @@ -1990,7 +1996,7 @@ class SplineAnnotator(BaseAnnotator): else custom_color_lookup, ) xy = self.trace.get(tracker_id=tracker_id) - spline_points = None; + spline_points = None if len(xy) > 3: x, y = xy[:, 0], xy[:, 1] @@ -2010,6 +2016,7 @@ class SplineAnnotator(BaseAnnotator): ) return scene + class HeatMapAnnotator(BaseAnnotator): """ A class for drawing heatmaps on an image based on provided detections. From eaf6cdd47701bebd2a6e9157d080c40a08a1790b Mon Sep 17 00:00:00 2001 From: Ashp116 Date: Wed, 30 Jul 2025 03:14:00 -0400 Subject: [PATCH 4/7] Precommit --- supervision/annotators/core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index b1ecc674..303b2ded 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1952,7 +1952,8 @@ class SplineAnnotator(BaseAnnotator): points. Defaults to `30`. thickness (int): The thickness of the spline lines. Defaults to `2`. smoothing_factor (int): The smoothing factor of the spline. - spline_order (int): The order of the spline. Use odd numbers that are 1 <= x <= 5. + spline_order (int): The order of the spline. Use odd numbers that are + between 1 <= x <= 5. color_lookup (ColorLookup): Strategy for mapping colors to annotations. Options are `INDEX`, `CLASS`, `TRACK`. """ From 8407d5a6923bc7ba556cd720721df69892cdb7b4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 30 Jul 2025 07:14:56 +0000 Subject: [PATCH 5/7] =?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 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 303b2ded..26c08e80 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1952,7 +1952,7 @@ class SplineAnnotator(BaseAnnotator): points. Defaults to `30`. thickness (int): The thickness of the spline lines. Defaults to `2`. smoothing_factor (int): The smoothing factor of the spline. - spline_order (int): The order of the spline. Use odd numbers that are + spline_order (int): The order of the spline. Use odd numbers that are between 1 <= x <= 5. color_lookup (ColorLookup): Strategy for mapping colors to annotations. Options are `INDEX`, `CLASS`, `TRACK`. From cdca2e6b51eee680fd1858392290c91999244dbe Mon Sep 17 00:00:00 2001 From: Ashp116 Date: Wed, 30 Jul 2025 13:48:38 -0400 Subject: [PATCH 6/7] UPDATE: Removed SplineAnnotator and added smooth flag for TraceAnnotator --- supervision/__init__.py | 1 - supervision/annotators/core.py | 97 ++-------------------------------- 2 files changed, 5 insertions(+), 93 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 1b2266cd..ab45651a 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -28,7 +28,6 @@ from supervision.annotators.core import ( PolygonAnnotator, RichLabelAnnotator, RoundBoxAnnotator, - SplineAnnotator, TraceAnnotator, TriangleAnnotator, ) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 26c08e80..41189e07 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1821,6 +1821,7 @@ class TraceAnnotator(BaseAnnotator): position: Position = Position.CENTER, trace_length: int = 30, thickness: int = 2, + smooth: int = 0, color_lookup: ColorLookup = ColorLookup.CLASS, ): """ @@ -1832,12 +1833,14 @@ class TraceAnnotator(BaseAnnotator): 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`. + smooth (int): The smoothing factor of the trace lines. Defaults to `0` 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.smooth = smooth self.color_lookup: ColorLookup = color_lookup @ensure_cv2_image_for_annotation @@ -1909,103 +1912,13 @@ class TraceAnnotator(BaseAnnotator): else custom_color_lookup, ) 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 - - -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, - smoothing_factor: int = 20, - spline_order: int = 3, - color_lookup: ColorLookup = ColorLookup.CLASS, - ): - """ - Args: - color (Union[Color, ColorPalette]): The color to draw the spline, can be - a single color or a color palette. - position (Position): The position of the spline. - Defaults to `CENTER`. - trace_length (int): The maximum length of the spline in terms of historical - points. Defaults to `30`. - thickness (int): The thickness of the spline lines. Defaults to `2`. - smoothing_factor (int): The smoothing factor of the spline. - spline_order (int): The order of the spline. Use odd numbers that are - between 1 <= x <= 5. - 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.smoothing_factor = smoothing_factor - self.color_lookup: ColorLookup = color_lookup - - if spline_order % 2 == 0 or spline_order < 1 or spline_order > 5: - raise ValueError("Spline order must be an odd number between 1 and 5.") - self.spline_order = spline_order - - @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 + spline_points = xy.astype(np.int32) if len(xy) > 3: x, y = xy[:, 0], xy[:, 1] - tck, u = splprep([x, y], s=self.smoothing_factor, k=self.spline_order) + tck, u = splprep([x, y], s=self.smooth) 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( From 0c9fea4d9fd80366998be2f400dac91f58ee7998 Mon Sep 17 00:00:00 2001 From: Ashp116 Date: Wed, 30 Jul 2025 18:24:14 -0400 Subject: [PATCH 7/7] UPDATE: Change Smooth to bool --- supervision/annotators/core.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 41189e07..f951e681 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1821,7 +1821,7 @@ class TraceAnnotator(BaseAnnotator): position: Position = Position.CENTER, trace_length: int = 30, thickness: int = 2, - smooth: int = 0, + smooth: bool = False, color_lookup: ColorLookup = ColorLookup.CLASS, ): """ @@ -1833,7 +1833,7 @@ class TraceAnnotator(BaseAnnotator): 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`. - smooth (int): The smoothing factor of the trace lines. Defaults to `0` + smooth (bool): Smooth the trace lines. color_lookup (ColorLookup): Strategy for mapping colors to annotations. Options are `INDEX`, `CLASS`, `TRACK`. """ @@ -1914,9 +1914,9 @@ class TraceAnnotator(BaseAnnotator): xy = self.trace.get(tracker_id=tracker_id) spline_points = xy.astype(np.int32) - if len(xy) > 3: + if len(xy) > 3 and self.smooth: x, y = xy[:, 0], xy[:, 1] - tck, u = splprep([x, y], s=self.smooth) + tck, u = splprep([x, y], s=20) x_new, y_new = splev(np.linspace(0, 1, 100), tck) spline_points = np.stack([x_new, y_new], axis=1).astype(np.int32)