Skip to content

Commit

Permalink
feat: 🚀 Triangle and Dot annotators outline thickness level and trian…
Browse files Browse the repository at this point in the history
…gle and circle shapes added (#1294)

* feat: 🚀 Triangle_annotator and Dot_annotator outline thickness value added
* feat: 🚀 Triangle_annotator and Dot_annotator outline triangle and circle added
* docs: 📝 outline_thickness py docs added
* refactor: remove else statement for thickness in circle not needed 
* fix(pre_commit): 🎨 auto format pre-commit hooks

---------

Signed-off-by: Onuralp SEZER <thunderbirdtr@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Onuralp SEZER <thunderbirdtr@gmail.com>
  • Loading branch information
onuralpszr and pre-commit-ci[bot] committed Jun 20, 2024
2 parents fbb770a + eec3615 commit b2b251a
Showing 1 changed file with 15 additions and 1 deletion.
16 changes: 15 additions & 1 deletion supervision/annotators/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -833,6 +833,7 @@ def __init__(
radius: int = 4,
position: Position = Position.CENTER,
color_lookup: ColorLookup = ColorLookup.CLASS,
outline_thickness: int = 0,
):
"""
Args:
Expand All @@ -842,11 +843,13 @@ def __init__(
position (Position): The anchor position for placing the dot.
color_lookup (ColorLookup): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACK`.
outline_thickness (int): Thickness of the outline of the dot.
"""
self.color: Union[Color, ColorPalette] = color
self.radius: int = radius
self.position: Position = position
self.color_lookup: ColorLookup = color_lookup
self.outline_thickness = outline_thickness

@convert_for_annotation_method
def annotate(
Expand Down Expand Up @@ -898,7 +901,12 @@ def annotate(
else custom_color_lookup,
)
center = (int(xy[detection_idx, 0]), int(xy[detection_idx, 1]))

cv2.circle(scene, center, self.radius, color.as_bgr(), -1)
if self.outline_thickness:
cv2.circle(
scene, center, self.radius, (0, 0, 0), self.outline_thickness
)
return scene


Expand Down Expand Up @@ -1613,6 +1621,7 @@ def __init__(
height: int = 10,
position: Position = Position.TOP_CENTER,
color_lookup: ColorLookup = ColorLookup.CLASS,
outline_thickness: int = 0,
):
"""
Args:
Expand All @@ -1623,12 +1632,14 @@ def __init__(
position (Position): The anchor position for placing the triangle.
color_lookup (ColorLookup): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACK`.
outline_thickness (int): Thickness of the outline of the triangle.
"""
self.color: Union[Color, ColorPalette] = color
self.base: int = base
self.height: int = height
self.position: Position = position
self.color_lookup: ColorLookup = color_lookup
self.outline_thickness: int = outline_thickness

@convert_for_annotation_method
def annotate(
Expand Down Expand Up @@ -1690,7 +1701,10 @@ def annotate(
)

cv2.fillPoly(scene, [vertices], color.as_bgr())

if self.outline_thickness:
cv2.polylines(
scene, [vertices], True, (0, 0, 0), thickness=self.outline_thickness
)
return scene


Expand Down

0 comments on commit b2b251a

Please sign in to comment.