[IconAnnotator] - Add IconAnnotator to Mark Objects with Custom Icons/Images

supervision/__init__.py

@@ -16,6 +16,7 @@
     EllipseAnnotator,
     HaloAnnotator,
     HeatMapAnnotator,
+    IconAnnotator,
     LabelAnnotator,
     MaskAnnotator,
     OrientedBoxAnnotator,

supervision/annotators/core.py

@@ -1763,3 +1763,113 @@ def validate_custom_values(
 
             if not all(0 <= value <= 1 for value in custom_values):
                 raise ValueError("All values in custom_values must be between 0 and 1.")
+
+
+class IconAnnotator(BaseAnnotator):
+    """
+    A class for drawing an icon on an image, using provided detections.
+    """
+
+    def __init__(
+        self,
+        icon_path: str,
+        position: Position = Position.TOP_CENTER,
+        icon_size: float = 0.2,
+        color: Union[Color, ColorPalette] = ColorPalette.default(),
+        color_lookup: ColorLookup = ColorLookup.CLASS,
+    ):
+        """
+        Args:
+            icon_path (str): path of the icon in png format.
+            position (Position): The position of the icon. Defaults to
+                `TOP_CENTER`.
+            icon_size (float): Represents the fraction of the original icon size to
+              be displayed, with a default value of 0.2
+              (equivalent to 20% of the original size).
+            color (Union[Color, ColorPalette]): The color to draw the trace, can be
+                a single color or a color palette.
+            color_lookup (str): Strategy for mapping colors to annotations.
+                Options are `INDEX`, `CLASS`, `TRACE`.
+        """
+        self.color: Union[Color, ColorPalette] = color
+        self.color_lookup: ColorLookup = color_lookup
+        self.position = position
+        icon = cv2.imread(icon_path, cv2.IMREAD_UNCHANGED)
+        if icon is None:
+            print(f"Error: Couldn't load the icon image from {icon_path}")
+            return
+
+        resized_icon_h, resized_icon_w = (
+            int(icon.shape[0] * icon_size),
+            int(icon.shape[1] * icon_size),
+        )
+        self.icon = cv2.resize(
+            icon, (resized_icon_h, resized_icon_w), interpolation=cv2.INTER_AREA
+        )
+
+    def annotate(
+        self,
+        scene: np.ndarray,
+        detections: Detections,
+        custom_color_lookup: Optional[np.ndarray] = None,
+    ) -> np.ndarray:
+        """
+        Annotates the given scene with icons based on the provided detections.
+        Args:
+            scene (np.ndarray): The image where bounding boxes will be drawn.
+            detections (Detections): Object detections to annotate.
+            custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
+                Allows to override the default color mapping strategy.
+        Returns:
+            np.ndarray: The annotated image.
+        Example:
+            ```python
+            >>> import supervision as sv
+            >>> image = ...
+            >>> detections = sv.Detections(...)
+            >>> icon_annotator = sv.IconAnnotator(icon_path='path_of_icon')
+            >>> annotated_frame = icon_annotator.annotate(
+            ...     scene=image.copy(),
+            ...     detections=detections
+            ... )
+            ```
+        """
+        icon_h, icon_w = self.icon.shape[:2]
+
+        xy = detections.get_anchors_coordinates(anchor=self.position)
+        for detection_idx in range(len(detections)):
+            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
+                ),
+            )
+            # Convert Color type to numpy list in BGRA format, and color icon with it
+            colored_icon = np.ones_like(self.icon) * (list(color.as_bgr()) + [1])
+
+            # The current positions of the anchors don't account for annotations that
+            # have mass. This visually centers the anchor, given the size of the icon
+            x_offset = 0
+            if self.position in [
+                Position.TOP_CENTER,
+                Position.CENTER,
+                Position.BOTTOM_CENTER,
+            ]:
+                x_offset = icon_w / 2
+
+            x, y = (
+                int(xy[detection_idx, 0] - x_offset),
+                int(xy[detection_idx, 1]),
+            )
+
+            # Takes the alpha channel from the original icon
+            alpha_channel = self.icon[:, :, 3]
+            # Apply alpha blending to paste the icon onto the larger image
+            scene[y : y + icon_h, x : x + icon_w][alpha_channel != 0] = colored_icon[
+                :, :, :3
+            ][alpha_channel != 0]
+        return scene

test/annotators/test_core.py

@@ -0,0 +1,65 @@
+from contextlib import ExitStack as DoesNotRaise
+from test.test_utils import mock_detections
+from typing import Optional
+
+import cv2
+import numpy as np
+import pytest
+
+from supervision.annotators.core import IconAnnotator
+from supervision.annotators.utils import ColorLookup
+from supervision.detection.core import Detections
+
+
+@pytest.mark.parametrize(
+    "detections, detection_idx, color_lookup, expected_result, "
+    "input_image_path, expected_image_path, exception",
+    [
+        (
+            mock_detections(
+                xyxy=[
+                    [123.45, 197.18, 1110.6, 710.51],
+                    [746.57, 40.801, 1142.1, 712.37],
+                ],
+                class_id=[0, 0],
+                tracker_id=None,
+            ),
+            0,
+            ColorLookup.INDEX,
+            0,
+            "../data/zidane.jpg",
+            "../data/zidane-icon-annotated.jpg",
+            DoesNotRaise(),
+        ),  # multiple detections; index lookup
+    ],
+)
+def test_icon_annotator(
+    detections: Detections,
+    detection_idx: int,
+    color_lookup: ColorLookup,
+    expected_result: Optional[int],
+    input_image_path: str,
+    expected_image_path: str,
+    exception: Exception,
+) -> None:
+    with exception:
+        icon_path = "../data/icons8-diamond-50.png"
+
+        icon_annotator = IconAnnotator(icon_path=icon_path, icon_size=1)
+        image = cv2.imread(input_image_path)
+        result_image = icon_annotator.annotate(
+            scene=image.copy(), detections=detections
+        )
+        expected_image = cv2.imread(expected_image_path)
+        assert images_are_equal(expected_image, result_image, 10)
+
+
+def images_are_equal(image1: np.ndarray, image2: np.ndarray, threshold: int) -> bool:
+    h, w = image1.shape[:2]
+    diff = cv2.subtract(image1, image2)
+    err = np.sum(diff**2)
+    mse = err / (float(h * w))
+
+    print(mse)
+
+    return mse < threshold  # Adjust the threshold as needed