eye-cv 1.0.0__py3-none-any.whl

eye/geometry/core.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from math import sqrt
+from typing import Tuple
+
+
+class Position(Enum):
+    """
+    Enum representing the position of an anchor point.
+    """
+
+    CENTER = "CENTER"
+    CENTER_LEFT = "CENTER_LEFT"
+    CENTER_RIGHT = "CENTER_RIGHT"
+    TOP_CENTER = "TOP_CENTER"
+    TOP_LEFT = "TOP_LEFT"
+    TOP_RIGHT = "TOP_RIGHT"
+    BOTTOM_LEFT = "BOTTOM_LEFT"
+    BOTTOM_CENTER = "BOTTOM_CENTER"
+    BOTTOM_RIGHT = "BOTTOM_RIGHT"
+    CENTER_OF_MASS = "CENTER_OF_MASS"
+
+    @classmethod
+    def list(cls):
+        return list(map(lambda c: c.value, cls))
+
+
+@dataclass
+class Point:
+    x: float
+    y: float
+
+    def as_xy_int_tuple(self) -> Tuple[int, int]:
+        return int(self.x), int(self.y)
+
+    def as_xy_float_tuple(self) -> Tuple[float, float]:
+        return self.x, self.y
+
+
+@dataclass
+class Vector:
+    start: Point
+    end: Point
+
+    @property
+    def magnitude(self) -> float:
+        """
+        Calculate the magnitude (length) of the vector.
+
+        Returns:
+            float: The magnitude of the vector.
+        """
+        dx = self.end.x - self.start.x
+        dy = self.end.y - self.start.y
+        return sqrt(dx**2 + dy**2)
+
+    @property
+    def center(self) -> Point:
+        """
+        Calculate the center point of the vector.
+
+        Returns:
+            Point: The center point of the vector.
+        """
+        return Point(
+            x=(self.start.x + self.end.x) / 2,
+            y=(self.start.y + self.end.y) / 2,
+        )
+
+    def cross_product(self, point: Point) -> float:
+        """
+        Calculate the 2D cross product (also known as the vector product or outer
+        product) of the vector and a point, treated as vectors in 2D space.
+
+        Args:
+            point (Point): The point to be evaluated, treated as the endpoint of a
+                vector originating from the 'start' of the main vector.
+
+        Returns:
+            float: The scalar value of the cross product. It is positive if 'point'
+                lies to the left of the vector (when moving from 'start' to 'end'),
+                negative if it lies to the right, and 0 if it is collinear with the
+                vector.
+        """
+        dx_vector = self.end.x - self.start.x
+        dy_vector = self.end.y - self.start.y
+        dx_point = point.x - self.start.x
+        dy_point = point.y - self.start.y
+        return (dx_vector * dy_point) - (dy_vector * dx_point)
+
+
+@dataclass
+class Rect:
+    x: float
+    y: float
+    width: float
+    height: float
+
+    @classmethod
+    def from_xyxy(cls, xyxy: Tuple[float, float, float, float]) -> Rect:
+        x1, y1, x2, y2 = xyxy
+        return cls(x=x1, y=y1, width=x2 - x1, height=y2 - y1)
+
+    @property
+    def top_left(self) -> Point:
+        return Point(x=self.x, y=self.y)
+
+    @property
+    def bottom_right(self) -> Point:
+        return Point(x=self.x + self.width, y=self.y + self.height)
+
+    def pad(self, padding) -> Rect:
+        return Rect(
+            x=self.x - padding,
+            y=self.y - padding,
+            width=self.width + 2 * padding,
+            height=self.height + 2 * padding,
+        )
+
+    def as_xyxy_int_tuple(self) -> Tuple[int, int, int, int]:
+        return (
+            int(self.x),
+            int(self.y),
+            int(self.x + self.width),
+            int(self.y + self.height),
+        )

eye/geometry/utils.py

@@ -0,0 +1,47 @@
+import numpy as np
+
+from eye.geometry.core import Point
+
+
+def get_polygon_center(polygon: np.ndarray) -> Point:
+    """
+    Calculate the center of a polygon. The center is calculated as the center
+    of the solid figure formed by the points of the polygon
+
+    Parameters:
+        polygon (np.ndarray): A 2-dimensional numpy ndarray representing the
+            vertices of the polygon.
+
+    Returns:
+        Point: The center of the polygon, represented as a
+            Point object with x and y attributes.
+
+    Raises:
+        ValueError: If the polygon has no vertices.
+
+    Examples:
+        ```python
+        import numpy as np
+        import eye as sv
+
+        polygon = np.array([[0, 0], [0, 2], [2, 2], [2, 0]])
+        sv.get_polygon_center(polygon=polygon)
+        # Point(x=1, y=1)
+        ```
+    """
+
+    # This is one of the 3 candidate algorithms considered for centroid calculation.
+    # For a more detailed discussion, see PR #1084 and commit eb33176
+
+    if len(polygon) == 0:
+        raise ValueError("Polygon must have at least one vertex.")
+
+    shift_polygon = np.roll(polygon, -1, axis=0)
+    signed_areas = np.cross(polygon, shift_polygon) / 2
+    if signed_areas.sum() == 0:
+        center = np.mean(polygon, axis=0).round()
+        return Point(x=center[0], y=center[1])
+    centroids = (polygon + shift_polygon) / 3.0
+    center = np.average(centroids, axis=0, weights=signed_areas).round()
+
+    return Point(x=center[0], y=center[1])

eye/keypoint/annotators.py

@@ -0,0 +1,442 @@
+from abc import ABC, abstractmethod
+from logging import warn
+from typing import List, Optional, Tuple, Union
+
+import cv2
+import numpy as np
+
+from eye.annotators.base import ImageType
+from eye.detection.utils import pad_boxes, spread_out_boxes
+from eye.draw.color import Color
+from eye.draw.utils import draw_rounded_rectangle
+from eye.geometry.core import Rect
+from eye.keypoint.core import KeyPoints
+from eye.keypoint.skeletons import SKELETONS_BY_VERTEX_COUNT
+from eye.utils.conversion import ensure_cv2_image_for_annotation
+
+
+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.EYE_ORANGE,
+        radius: int = 4,
+    ) -> None:
+        """
+        Args:
+            color (Color): The color to use for annotating key points.
+            radius (int): The radius of the circles used to represent the key
+                points.
+        """
+        self.color = color
+        self.radius = radius
+
+    @ensure_cv2_image_for_annotation
+    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 skeleton vertices 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 eye 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/
+        eye-annotator-examples/vertex-annotator-example.png)
+        """
+        assert isinstance(scene, np.ndarray)
+        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): The color to use for the edges.
+            thickness (int): 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
+
+    @ensure_cv2_image_for_annotation
+    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 skeleton edges 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 eye 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/
+        eye-annotator-examples/edge-annotator-example.png)
+        """
+        assert isinstance(scene, np.ndarray)
+        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 that draws labels of skeleton vertices on images. It uses specified key
+    points to determine the locations where the vertices should be drawn.
+    """
+
+    def __init__(
+        self,
+        color: Union[Color, List[Color]] = Color.EYE_ORANGE,
+        text_color: Union[Color, List[Color]] = Color.WHITE,
+        text_scale: float = 0.5,
+        text_thickness: int = 1,
+        text_padding: int = 10,
+        border_radius: int = 0,
+        smart_position: bool = False,
+    ):
+        """
+        Args:
+            color (Union[Color, List[Color]]): The color to use for each
+                keypoint label. If a list is provided, the colors will be used in order
+                for each keypoint.
+            text_color (Union[Color, List[Color]]): The color to use
+                for the labels. If a list is provided, the colors will be used in order
+                for each keypoint.
+            text_scale (float): The scale of the text.
+            text_thickness (int): The thickness of the text.
+            text_padding (int): The padding around the text.
+            border_radius (int): The radius of the rounded corners of the
+                boxes. Set to a high value to produce circles.
+            smart_position (bool): Spread out the labels to avoid overlap.
+        """
+        self.border_radius: int = border_radius
+        self.color: Union[Color, List[Color]] = color
+        self.text_color: Union[Color, List[Color]] = text_color
+        self.text_scale: float = text_scale
+        self.text_thickness: int = text_thickness
+        self.text_padding: int = text_padding
+        self.smart_position = smart_position
+
+    def annotate(
+        self,
+        scene: ImageType,
+        key_points: KeyPoints,
+        labels: Optional[List[str]] = None,
+    ) -> ImageType:
+        """
+        A class that draws labels of skeleton vertices on images. It uses specified key
+            points to determine the locations where the vertices should be drawn.
+
+        Args:
+            scene (ImageType): The image where vertex labels 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.
+            labels (Optional[List[str]]): A list of labels to be displayed on the
+                annotated image. If not provided, keypoint indices will be used.
+
+        Returns:
+            The annotated image, matching the type of `scene` (`numpy.ndarray`
+                or `PIL.Image.Image`)
+
+        Example:
+            ```python
+            import eye as sv
+
+            image = ...
+            key_points = sv.KeyPoints(...)
+
+            vertex_label_annotator = sv.VertexLabelAnnotator(
+                color=sv.Color.GREEN,
+                text_color=sv.Color.BLACK,
+                border_radius=5
+            )
+            annotated_frame = vertex_label_annotator.annotate(
+                scene=image.copy(),
+                key_points=key_points
+            )
+            ```
+
+        ![vertex-label-annotator-example](https://media.roboflow.com/
+        eye-annotator-examples/vertex-label-annotator-example.png)
+
+        !!! tip
+
+            `VertexLabelAnnotator` allows to customize the color of each keypoint label
+            values.
+
+        Example:
+            ```python
+            import eye as sv
+
+            image = ...
+            key_points = sv.KeyPoints(...)
+
+            LABELS = [
+                "nose", "left eye", "right eye", "left ear",
+                "right ear", "left shoulder", "right shoulder", "left elbow",
+                "right elbow", "left wrist", "right wrist", "left hip",
+                "right hip", "left knee", "right knee", "left ankle",
+                "right ankle"
+            ]
+
+            COLORS = [
+                "#FF6347", "#FF6347", "#FF6347", "#FF6347",
+                "#FF6347", "#FF1493", "#00FF00", "#FF1493",
+                "#00FF00", "#FF1493", "#00FF00", "#FFD700",
+                "#00BFFF", "#FFD700", "#00BFFF", "#FFD700",
+                "#00BFFF"
+            ]
+            COLORS = [sv.Color.from_hex(color_hex=c) for c in COLORS]
+
+            vertex_label_annotator = sv.VertexLabelAnnotator(
+                color=COLORS,
+                text_color=sv.Color.BLACK,
+                border_radius=5
+            )
+            annotated_frame = vertex_label_annotator.annotate(
+                scene=image.copy(),
+                key_points=key_points,
+                labels=labels
+            )
+            ```
+        ![vertex-label-annotator-custom-example](https://media.roboflow.com/
+        eye-annotator-examples/vertex-label-annotator-custom-example.png)
+        """
+        assert isinstance(scene, np.ndarray)
+        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,
+        )
+
+        text_colors = self.preprocess_and_validate_colors(
+            colors=self.text_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]
+        text_colors = text_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)
+
+        if self.smart_position:
+            xyxy_padded = spread_out_boxes(xyxy_padded)
+            xyxy = pad_boxes(xyxy=xyxy_padded, px=-self.text_padding)
+
+        for text, color, text_color, box, box_padded in zip(
+            labels, colors, text_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[3]),
+                fontFace=font,
+                fontScale=self.text_scale,
+                color=text_color.as_bgr(),
+                thickness=self.text_thickness,
+                lineType=cv2.LINE_AA,
+            )
+
+        return scene
+
+    @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,
+        )
+
+    @staticmethod
+    def preprocess_and_validate_labels(
+        labels: Optional[List[str]], points_count: int, skeletons_count: int
+    ) -> np.ndarray:
+        if labels and len(labels) != points_count:
+            raise ValueError(
+                f"Number of labels ({len(labels)}) must match number of key points "
+                f"({points_count})."
+            )
+        if labels is None:
+            labels = [str(i) for i in range(points_count)]
+
+        return np.array(labels * skeletons_count)
+
+    @staticmethod
+    def preprocess_and_validate_colors(
+        colors: Optional[Union[Color, List[Color]]],
+        points_count: int,
+        skeletons_count: int,
+    ) -> np.ndarray:
+        if isinstance(colors, list) and len(colors) != points_count:
+            raise ValueError(
+                f"Number of colors ({len(colors)}) must match number of key points "
+                f"({points_count})."
+            )
+        return (
+            np.array(colors * skeletons_count)
+            if isinstance(colors, list)
+            else np.array([colors] * points_count * skeletons_count)
+        )