Merge pull request #930 from dfreilich/dfreilich/add-icon-annotator

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

docs/detection/annotators.md

@@ -316,6 +316,33 @@ status: new
 
     </div>
 
+=== "Icon"
+
+    ```python
+    import supervision as sv
+
+    image = ...
+    detections = sv.Detections(...)
+
+    icon_paths = [
+        "<ICON_PATH>"
+        for _ in detections
+    ]
+
+    icon_annotator = sv.IconAnnotator()
+    annotated_frame = icon_annotator.annotate(
+        scene=image.copy(),
+        detections=detections,
+        icon_path=icon_paths
+    )
+    ```
+
+    <div class="result" markdown>
+
+    ![icon-annotator-example](https://media.roboflow.com/supervision-annotator-examples/icon-annotator-example.png){ align=center width="800" }
+
+    </div>
+
 === "Crop"
 
     ```python
@@ -550,6 +577,12 @@ status: new
 
 :::supervision.annotators.core.RichLabelAnnotator
 
+<div class="md-typeset">
+    <h2><a href="#supervision.annotators.core.IconAnnotator">IconAnnotator</a></h2>
+</div>
+
+:::supervision.annotators.core.IconAnnotator
+
 <div class="md-typeset">
     <h2><a href="#supervision.annotators.core.BlurAnnotator">BlurAnnotator</a></h2>
 </div>

supervision/__init__.py

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

supervision/annotators/core.py

@@ -1,3 +1,4 @@
+from functools import lru_cache
 from math import sqrt
 from typing import List, Optional, Tuple, Union
 
@@ -22,7 +23,12 @@
     ensure_cv2_image_for_annotation,
     ensure_pil_image_for_annotation,
 )
-from supervision.utils.image import crop_image, overlay_image, scale_image
+from supervision.utils.image import (
+    crop_image,
+    letterbox_image,
+    overlay_image,
+    scale_image,
+)
 from supervision.utils.internal import deprecated
 
 
@@ -1409,6 +1415,106 @@ def _load_default_font(size):
         return font
 
 
+class IconAnnotator(BaseAnnotator):
+    """
+    A class for drawing an icon on an image, using provided detections.
+    """
+
+    def __init__(
+        self,
+        icon_resolution_wh: Tuple[int, int] = (64, 64),
+        icon_position: Position = Position.TOP_CENTER,
+        offset_xy: Tuple[int, int] = (0, 0),
+    ):
+        """
+        Args:
+            icon_resolution_wh (Tuple[int, int]): The size of drawn icons.
+                All icons will be resized to this resolution, keeping the aspect ratio.
+            icon_position (Position): The position of the icon.
+            offset_xy (Tuple[int, int]): The offset to apply to the icon position,
+                in pixels. Can be both positive and negative.
+        """
+        self.icon_resolution_wh = icon_resolution_wh
+        self.position = icon_position
+        self.offset_xy = offset_xy
+
+    @ensure_cv2_image_for_annotation
+    def annotate(
+        self, scene: ImageType, detections: Detections, icon_path: Union[str, List[str]]
+    ) -> ImageType:
+        """
+        Annotates the given scene with given icons.
+
+        Args:
+            scene (ImageType): The image where labels will be drawn.
+                `ImageType` is a flexible type, accepting either `numpy.ndarray`
+                or `PIL.Image.Image`.
+            detections (Detections): Object detections to annotate.
+            icon_path (Union[str, List[str]]): The path to the PNG image to use as an
+                icon. Must be a single path or a list of paths, one for each detection.
+                Pass an empty string `""` to draw nothing.
+
+        Returns:
+            The annotated image, matching the type of `scene` (`numpy.ndarray`
+                or `PIL.Image.Image`)
+
+        Example:
+            ```python
+            import supervision as sv
+
+            image = ...
+            detections = sv.Detections(...)
+
+            available_icons = ["roboflow.png", "lenny.png"]
+            icon_paths = [np.random.choice(available_icons) for _ in detections]
+
+            icon_annotator = sv.IconAnnotator()
+            annotated_frame = icon_annotator.annotate(
+                scene=image.copy(),
+                detections=detections,
+                icon_path=icon_paths
+            )
+            ```
+
+        ![icon-annotator-example](https://media.roboflow.com/
+        supervision-annotator-examples/icon-annotator-example.png)
+        """
+        assert isinstance(scene, np.ndarray)
+        if isinstance(icon_path, list) and len(icon_path) != len(detections):
+            raise ValueError(
+                f"The number of icon paths provided ({len(icon_path)}) does not match "
+                f"the number of detections ({len(detections)}). Either provide a single"
+                f" icon path or one for each detection."
+            )
+
+        xy = detections.get_anchors_coordinates(anchor=self.position).astype(int)
+
+        for detection_idx in range(len(detections)):
+            current_path = (
+                icon_path if isinstance(icon_path, str) else icon_path[detection_idx]
+            )
+            if current_path == "":
+                continue
+            icon = self._load_icon(current_path)
+            icon_h, icon_w = icon.shape[:2]
+
+            x = int(xy[detection_idx, 0] - icon_w / 2 + self.offset_xy[0])
+            y = int(xy[detection_idx, 1] - icon_h / 2 + self.offset_xy[1])
+
+            scene[:] = overlay_image(scene, icon, (x, y))
+        return scene
+
+    @lru_cache
+    def _load_icon(self, icon_path: str) -> np.ndarray:
+        icon = cv2.imread(icon_path, cv2.IMREAD_UNCHANGED)
+        if icon is None:
+            raise FileNotFoundError(
+                f"Error: Couldn't load the icon image from {icon_path}"
+            )
+        icon = letterbox_image(image=icon, resolution_wh=self.icon_resolution_wh)
+        return icon
+
+
 class BlurAnnotator(BaseAnnotator):
     """
     A class for blurring regions in an image using provided detections.

supervision/utils/image.py

@@ -270,6 +270,7 @@ def letterbox_image(
 
     ![letterbox_image](https://media.roboflow.com/supervision-docs/letterbox-image.png){ align=center width="800" }
     """  # noqa E501 // docs
+    assert isinstance(image, np.ndarray)
     color = unify_to_bgr(color=color)
     resized_image = resize_image(
         image=image, resolution_wh=resolution_wh, keep_aspect_ratio=True
@@ -279,7 +280,7 @@ def letterbox_image(
     padding_bottom = resolution_wh[1] - height_new - padding_top
     padding_left = (resolution_wh[0] - width_new) // 2
     padding_right = resolution_wh[0] - width_new - padding_left
-    return cv2.copyMakeBorder(
+    image_with_borders = cv2.copyMakeBorder(
         resized_image,
         padding_top,
         padding_bottom,
@@ -289,6 +290,14 @@ def letterbox_image(
         value=color,
     )
 
+    if image.shape[2] == 4:
+        image[:padding_top, :, 3] = 0
+        image[height_new - padding_bottom :, :, 3] = 0
+        image[:, :padding_left, 3] = 0
+        image[:, width_new - padding_right :, 3] = 0
+
+    return image_with_borders
+
 
 def overlay_image(
     image: npt.NDArray[np.uint8],
@@ -341,9 +350,20 @@ def overlay_image(
     crop_x_max = image_width - max((anchor_x + image_width) - scene_width, 0)
     crop_y_max = image_height - max((anchor_y + image_height) - scene_height, 0)
 
-    image[y_min:y_max, x_min:x_max] = overlay[
-        crop_y_min:crop_y_max, crop_x_min:crop_x_max
-    ]
+    if overlay.shape[2] == 4:
+        b, g, r, alpha = cv2.split(
+            overlay[crop_y_min:crop_y_max, crop_x_min:crop_x_max]
+        )
+        alpha = alpha[:, :, None] / 255.0
+        overlay_color = cv2.merge((b, g, r))
+
+        roi = image[y_min:y_max, x_min:x_max]
+        roi[:] = roi * (1 - alpha) + overlay_color * alpha
+        image[y_min:y_max, x_min:x_max] = roi
+    else:
+        image[y_min:y_max, x_min:x_max] = overlay[
+            crop_y_min:crop_y_max, crop_x_min:crop_x_max
+        ]
 
     return image