Spaces:
Sleeping
Sleeping
File size: 1,790 Bytes
93277a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import cv2
import numpy as np
import supervision as sv
from supervision.annotators.base import ImageType
from supervision.keypoint.core import KeyPoints
from supervision.utils.conversion import ensure_cv2_image_for_annotation
class VertexAnnotatorHeart(sv.VertexAnnotator):
"""
Extend the original class to add vertices with the shape of heart or star.
"""
def __init__(self, color=sv.Color.RED, radius=5):
super().__init__(color=color, radius=radius)
@ensure_cv2_image_for_annotation
def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
assert isinstance(scene, np.ndarray)
if len(key_points) == 0:
return scene
for xy in key_points.xy:
for x, y in xy:
# Two circles (top lobes of heart)
cv2.circle(
img=scene,
center=(int(x) - self.radius // 2, int(y)),
radius=self.radius // 2,
color=self.color.as_bgr(),
thickness=-1,
)
cv2.circle(
img=scene,
center=(int(x) + self.radius // 2, int(y)),
radius=self.radius // 2,
color=self.color.as_bgr(),
thickness=-1,
)
# Triangle (bottom of heart)
pts = np.array(
[
[int(x) - self.radius, int(y)],
[int(x) + self.radius, int(y)],
[int(x), y + self.radius * 2],
],
np.int32,
).reshape((-1, 1, 2))
cv2.fillPoly(scene, [pts], self.color.as_bgr())
return scene
|