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(