Repository: roboflow/supervision
File: supervision/keypoint/annotators.py (Python)

Code before:
from abc import ABC, abstractmethod
from logging import warn
from typing import List, Optional, Tuple, Union

import cv2
import numpy as np

from supervision import Rect, pad_boxes
from supervision.annotators.base import ImageType
from supervision.draw.color import Color
from supervision.draw.utils import draw_rounded_rectangle
from supervision.keypoint.core import KeyPoints
from supervision.keypoint.skeletons import SKELETONS_BY_VERTEX_COUNT
from supervision.utils.conversion import convert_for_annotation_method


class BaseKeyPointAnnotator(ABC):
    @abstractmethod
    def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
        pass


class VertexAnnotator(BaseKeyPointAnnotator):
    """
    A class that specializes in drawing skeleton vertices on images. It uses
    specified key points to determine the locations where the vertices should be
    drawn.
    """

    def __init__(
        self,
        color: Color = Color.ROBOFLOW,
        radius: int = 4,
    ) -> None:
        """
        Args:
            color (Color, optional): The color to use for annotating key points.
            radius (int, optional): The radius of the circles used to represent the key
                points.
        """
        self.color = color
        self.radius = radius

    @convert_for_annotation_method
    def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
        """
        Annotates the given scene with skeleton vertices based on the provided key
        points. It draws circles at each key point location.

        Args:
            scene (ImageType): The image where bounding boxes will be drawn. `ImageType`
                is a flexible type, accepting either `numpy.ndarray` or
                `PIL.Image.Image`.
            key_points (KeyPoints): A collection of key points where each key point
                consists of x and y coordinates.

        Returns:
            The annotated image, matching the type of `scene` (`numpy.ndarray`
                or `PIL.Image.Image`)

        Example:
            ```python
            import supervision as sv

            image = ...
            key_points = sv.KeyPoints(...)

            vertex_annotator = sv.VertexAnnotator(color=sv.Color.GREEN, radius=10)
            annotated_frame = vertex_annotator.annotate(
                scene=image.copy(),
                key_points=key_points
            )
            ```

        ![vertex-annotator-example](https://media.roboflow.com/
        supervision-annotator-examples/vertex-annotator-example.png)
        """
        if len(key_points) == 0:
            return scene

        for xy in key_points.xy:
            for x, y in xy:
                cv2.circle(
                    img=scene,
                    center=(int(x), int(y)),
                    radius=self.radius,
                    color=self.color.as_bgr(),
                    thickness=-1,
                )

        return scene


class EdgeAnnotator(BaseKeyPointAnnotator):
    """
    A class that specializes in drawing skeleton edges on images using specified key
    points. It connects key points with lines to form the skeleton structure.
    """

    def __init__(
        self,
        color: Color = Color.ROBOFLOW,
        thickness: int = 2,
        edges: Optional[List[Tuple[int, int]]] = None,
    ) -> None:
        """
        Args:
            color (Color, optional): The color to use for the edges.
            thickness (int, optional): The thickness of the edges.
            edges (Optional[List[Tuple[int, int]]]): The edges to draw.
                If set to `None`, will attempt to select automatically.
        """
        self.color = color
        self.thickness = thickness
        self.edges = edges

    @convert_for_annotation_method
    def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
        """
        Annotates the given scene by drawing lines between specified key points to form
        edges.

        Args:
            scene (ImageType): The image where bounding boxes will be drawn. `ImageType`
                is a flexible type, accepting either `numpy.ndarray` or
                `PIL.Image.Image`.
            key_points (KeyPoints): A collection of key points where each key point
                consists of x and y coordinates.

        Returns:
            Returns:
                The annotated image, matching the type of `scene` (`numpy.ndarray`
                    or `PIL.Image.Image`)

        Example:
            ```python
            import supervision as sv

            image = ...
            key_points = sv.KeyPoints(...)

            edge_annotator = sv.EdgeAnnotator(color=sv.Color.GREEN, thickness=5)
            annotated_frame = edge_annotator.annotate(
                scene=image.copy(),
                key_points=key_points
            )
            ```

        ![edge-annotator-example](https://media.roboflow.com/
        supervision-annotator-examples/edge-annotator-example.png)
        """
        if len(key_points) == 0:
            return scene

        for xy in key_points.xy:
            edges = self.edges
            if not edges:
                edges = SKELETONS_BY_VERTEX_COUNT.get(len(xy))
            if not edges:
                warn(f"No skeleton found with {len(xy)} vertices")
                return scene

            for class_a, class_b in edges:
                xy_a = xy[class_a - 1]
                xy_b = xy[class_b - 1]
                missing_a = np.allclose(xy_a, 0)
                missing_b = np.allclose(xy_b, 0)
                if missing_a or missing_b:
                    continue

                cv2.line(
                    img=scene,
                    pt1=(int(xy_a[0]), int(xy_a[1])),
                    pt2=(int(xy_b[0]), int(xy_b[1])),
                    color=self.color.as_bgr(),
                    thickness=self.thickness,
                )

        return scene


class VertexLabelAnnotator:
    """
    A class for annotating vertex labels on an image using provided detections.
    """

    def __init__(
        self,
        color: Union[Color, List[Color]] = Color.ROBOFLOW,
        text_color: Color = Color.WHITE,
        text_scale: float = 0.5,
        text_thickness: int = 1,
        text_padding: int = 10,
        border_radius: int = 0,
    ):
        self.border_radius: int = border_radius
        self.color: Union[Color, List[Color]] = color
        self.text_color: Color = text_color
        self.text_scale: float = text_scale
        self.text_thickness: int = text_thickness
        self.text_padding: int = text_padding

    @staticmethod
    def get_text_bounding_box(
        text: str,
        font: int,
        text_scale: float,
        text_thickness: int,
        center_coordinates: Tuple[int, int],
    ) -> Tuple[int, int, int, int]:
        text_w, text_h = cv2.getTextSize(
            text=text,
            fontFace=font,
            fontScale=text_scale,
            thickness=text_thickness,
        )[0]
        center_x, center_y = center_coordinates
        return (
            center_x - text_w // 2,
            center_y - text_h // 2,
            center_x + text_w // 2,
            center_y + text_h // 2,
        )

    def annotate(
        self, scene: ImageType, key_points: KeyPoints, labels: List[str] = None
    ) -> ImageType:
        font = cv2.FONT_HERSHEY_SIMPLEX

        N, K, _ = key_points.xy.shape

        if N == 0:
            return scene

        anchors = key_points.xy.reshape(K * N, 2).astype(int)
        colors = (
            np.array(self.color * N)
            if isinstance(self.color, list)
            else np.array([self.color] * K * N)
        )
        labels = np.array(labels * N)

        mask = np.all(anchors != 0, axis=1)

        if np.all(mask == False):
            return scene

        anchors = anchors[mask]
        colors = colors[mask]
        labels = labels[mask]

        xyxy = np.array(
            [
                self.get_text_bounding_box(
                    text=label,
                    font=font,
                    text_scale=self.text_scale,
                    text_thickness=self.text_thickness,
                    center_coordinates=tuple(anchor),
                )
                for anchor, label in zip(anchors, labels)
            ]
        )

        xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding)

        for text, color, box, box_padded in zip(labels, colors, xyxy, xyxy_padded):
            draw_rounded_rectangle(
                scene=scene,
                rect=Rect.from_xyxy(box_padded),
                color=color,
                border_radius=self.border_radius,
            )
            cv2.putText(
                img=scene,
                text=text,
                org=(box[0], box[1] + self.text_padding),
                fontFace=font,
                fontScale=self.text_scale,
                color=self.text_color.as_rgb(),
                thickness=self.text_thickness,
                lineType=cv2.LINE_AA,
            )

        return scene


Review comment (bug):
When an incorrect number of labels is passed in, this produces:
```
IndexError: boolean index did not match indexed array along dimension 0; dimension is 64 but corresponding boolean dimension is 68
```

It would be awesome if we could tell users exactly how the problem can be fixed.
```python
if labels and len(labels) != K:
    raise ValueError(
        f"Number of labels ({len(labels)}) must match number of key points ({K})"
    )
```


Code after:
        """
        font = cv2.FONT_HERSHEY_SIMPLEX

        skeletons_count, points_count, _ = key_points.xy.shape
        if skeletons_count == 0:
            return scene

        anchors = key_points.xy.reshape(points_count * skeletons_count, 2).astype(int)
        mask = np.all(anchors != 0, axis=1)

        if not np.any(mask):
            return scene

        colors = self.preprocess_and_validate_colors(
            colors=self.color,
            points_count=points_count,
            skeletons_count=skeletons_count,
        )

        labels = self.preprocess_and_validate_labels(
            labels=labels, points_count=points_count, skeletons_count=skeletons_count
        )

        anchors = anchors[mask]
        colors = colors[mask]
        labels = labels[mask]

        xyxy = np.array(
            [
                self.get_text_bounding_box(
                    text=label,
                    font=font,
                    text_scale=self.text_scale,
                    text_thickness=self.text_thickness,
                    center_coordinates=tuple(anchor),
                )
                for anchor, label in zip(anchors, labels)
            ]
        )

        xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding)

        for text, color, box, box_padded in zip(labels, colors, xyxy, xyxy_padded):
            draw_rounded_rectangle(
                scene=scene,
                rect=Rect.from_xyxy(box_padded),
                color=color,
                border_radius=self.border_radius,
            )
            cv2.putText(
                img=scene,