📘 docs updates for LabelAnnotator

SkalskiP · SkalskiP · commit 3e8906d83079 · 2023-09-28T23:24:22.000+02:00
diff --git a/docs/detection/annotate.md b/docs/detection/annotate.md
@@ -13,3 +13,7 @@
 ## BoxCornerAnnotator
 
 :::supervision.annotators.core.BoxCornerAnnotator
+
+## LabelAnnotator
+
+:::supervision.annotators.core.LabelAnnotator
diff --git a/supervision/__init__.py b/supervision/__init__.py
@@ -10,10 +10,8 @@
     BoundingBoxAnnotator,
     BoxCornerAnnotator,
     EllipseAnnotator,
-    LabelAdvancedAnnotator,
     LabelAnnotator,
     MaskAnnotator,
-    TraceAnnotator,
 )
 from supervision.classification.core import Classifications
 from supervision.dataset.core import (
diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py
@@ -307,6 +307,9 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
 
 
 class LabelAnnotator:
+    """
+   A class for annotating labels on an image using provided detections.
+   """
     def __init__(
         self,
         color: Union[Color, ColorPalette] = ColorPalette.default(),
@@ -317,6 +320,19 @@ def __init__(
         text_position: Position = Position.TOP_LEFT,
         color_map: str = "class",
     ):
+        """
+        Args:
+            color (Union[Color, ColorPalette]): The color or color palette to use for
+                annotating the text background.
+            text_color (Color): The color to use for the text.
+            text_scale (float): Font scale for the text.
+            text_thickness (int): Thickness of the text characters.
+            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`.
+        """
         self.color: Union[Color, ColorPalette] = color
         self.text_color: Color = text_color
         self.text_scale: float = text_scale
@@ -367,6 +383,31 @@ def annotate(
         detections: Detections,
         labels: List[str] = None,
     ) -> np.ndarray:
+        """
+        Annotates the given scene with labels based on the provided detections.
+
+        Args:
+            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.
+
+        Returns:
+            np.ndarray: The annotated image.
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> image = ...
+            >>> detections = sv.Detections(...)
+
+            >>> label_annotator = sv.LabelAnnotator()
+            >>> annotated_frame = label_annotator.annotate(
+            ...     scene=image.copy(),
+            ...     detections=detections
+            ... )
+            ```
+        """
         font = cv2.FONT_HERSHEY_SIMPLEX
         for detection_idx in range(len(detections)):
             detection_xyxy = detections.xyxy[detection_idx].astype(int)
@@ -416,213 +457,3 @@ def annotate(
                 lineType=cv2.LINE_AA,
             )
         return scene
-
-
-class LabelAdvancedAnnotator(BaseAnnotator):
-    def __init__(
-        self,
-        color: Union[Color, ColorPalette] = ColorPalette.default(),
-        text_color: Color = Color.black(),
-        text_padding: int = 20,
-        color_by_track: bool = False,
-        font: Optional[str] = None,
-        font_size: Optional[int] = 15,
-    ):
-        if font and os.path.exists(font):
-            self.font = ImageFont.truetype(font, font_size)
-        else:
-            self.font = ImageFont.load_default()
-        self.color: Union[Color, ColorPalette] = color
-        self.text_color: Color = text_color
-        self.text_padding: int = text_padding
-        self.color_by_track = color_by_track
-
-    def annotate(
-        self,
-        scene: np.ndarray,
-        detections: Detections,
-        labels: Optional[List[str]] = None,
-    ) -> np.ndarray:
-        """
-        Draws text on the frame using the detections provided and label.
-
-        Args:
-            scene (np.ndarray): The image on which the bounding boxes will be drawn
-            detections (Detections): The detections for which
-                the bounding boxes will be drawn
-            labels (Optional[List[str]]): An optional list of labels corresponding
-                to each detection. If `labels` are not provided,
-                corresponding `class_id` will be used as label.
-        Returns:
-            np.ndarray: The image with the bounding boxes drawn on it
-
-        Example:
-            ```python
-            >>> import supervision as sv
-
-            >>> classes = ['person', ...]
-            >>> image = ...
-            >>> detections = sv.Detections(...)
-
-            >>> pil_label_annotator = sv.LabelAdvancedAnnotator()
-            >>> labels = [
-            ...     f"{classes[class_id]} {confidence:0.2f}"
-            ...     for _, _, confidence, class_id, _
-            ...     in detections
-            ... ]
-            >>> annotated_frame = pil_label_annotator.annotate(
-            ...     scene=image.copy(),
-            ...     detections=detections,
-            ...     labels=labels,
-            ... )
-            ```
-        """
-        pil_image = Image.fromarray(scene)
-        draw = ImageDraw.Draw(pil_image)
-        text_color = "#fff"
-
-        for i in range(len(detections)):
-            x1, y1, x2, y2 = detections.xyxy[i].astype(int)
-            if self.color_by_track:
-                tracker_id = (
-                    detections.tracker_id[i]
-                    if detections.tracker_id is not None
-                    else None
-                )
-                idx = tracker_id if tracker_id is not None else i
-            else:
-                class_id = (
-                    detections.class_id[i] if detections.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
-            )
-
-            text = (
-                f"{idx}"
-                if (labels is None or len(detections) != len(labels))
-                else labels[i]
-            )
-
-            text_bbox = draw.textbbox((x1, y1), text, font=self.font)
-
-            text_height = text_bbox[3] - text_bbox[1]
-            text_width = text_bbox[2] - text_bbox[0]
-
-            text_x = x1 + self.text_padding / 2
-            text_y = y1 - self.text_padding / 2 - text_height
-
-            text_background_x1 = x1
-            text_background_y1 = y1 - self.text_padding / 2 - text_height
-
-            text_background_x2 = x1 + 2 * self.text_padding / 2 + text_width
-            text_background_y2 = y1  # correct
-
-            draw.rectangle(
-                (
-                    text_background_x1,
-                    text_background_y1,
-                    text_background_x2,
-                    text_background_y2,
-                ),
-                fill=color.as_bgr(),
-            )
-            draw.text((text_x, text_y), text, font=self.font, fill=text_color)
-
-        scene = np.asarray(pil_image)
-        return scene
-
-
-class TraceAnnotator(BaseAnnotator):
-    """
-    A class for drawing trajectory of a tracker on an image using detections provided.
-
-    Attributes:
-        color (Union[Color, ColorPalette]): The color to draw the trajectory,
-            can be a single color or a color palette
-        color_by_track (bool): Whther to use tracker id to pick the color
-        position (Optional[Position]): Choose position of trajectory such as
-            center position, top left corner, etc
-        trace_length (int): Length of the previous points
-        thickness (int): thickness of the line
-    """
-
-    def __init__(
-        self,
-        color: Union[Color, ColorPalette] = ColorPalette.default(),
-        color_by_track: bool = False,
-        position: Optional[Position] = Position.CENTER,
-        trace_length: int = 30,
-        thickness: int = 2,
-    ):
-        self.color: Union[Color, ColorPalette] = color
-        self.color_by_track = color_by_track
-        self.position = position
-        self.tracker_storage = defaultdict(lambda: [])
-        self.trace_length = trace_length
-        self.thickness = thickness
-
-    def annotate(
-        self, scene: np.ndarray, detections: Detections, **kwargs
-    ) -> np.ndarray:
-        """
-        Draw the object trajectory based on history of tracked objects
-
-        Args:
-            scene (np.ndarray): The image on which the trace will be drawn
-            detections (Detections): The detections for trajectory and points
-
-        Returns:
-            np.ndarray: The image with the masks overlaid
-        Example:
-            ```python
-            >>> import supervision as sv
-
-            >>> classes = ['person', ...]
-            >>> image = ...
-            >>> detections = sv.Detections(...)
-
-            >>> trace_annotator = sv.TraceAnnotator()
-            >>> annotated_frame = trace_annotator.annotate(
-            ...     scene=image.copy(),
-            ...     detections=detections
-            ... )
-            ```
-        """
-        if detections.tracker_id is None:
-            return scene
-
-        anchor_points = detections.get_anchor_coordinates(anchor=self.position)
-
-        for i, tracker_id in enumerate(detections.tracker_id):
-            track = self.tracker_storage[tracker_id]
-            track.append((anchor_points[i][0], anchor_points[i][1]))
-            if len(track) > self.trace_length:
-                track.pop(0)
-            points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
-
-            if self.color_by_track:
-                idx = tracker_id if tracker_id is not None else i
-            else:
-                class_id = (
-                    detections.class_id[i] if detections.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
-            )
-            cv2.polylines(
-                scene,
-                [points],
-                isClosed=False,
-                color=color.as_bgr(),
-                thickness=self.thickness,
-            )
-
-        return scene
