From f94a8cb9a1bc3c0862c3875b26e8648cfc9617d2 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 6 Dec 2023 13:32:48 +0100 Subject: [PATCH 1/9] Add TriangleAnnotator to annotators Added the ability to draw triangle markers on an image based on the provided detections with the TriangleAnnotator class. This is useful when needing to specifically mark or highlight certain areas of an image. The user can now specify both the color of the annotation and the metrics of the triangle such as base and height, as well as its position. --- supervision/__init__.py | 1 + supervision/annotators/core.py | 70 ++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/supervision/__init__.py b/supervision/__init__.py index a40c58555..4d924acf1 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -21,6 +21,7 @@ PixelateAnnotator, PolygonAnnotator, TraceAnnotator, + TriangleAnnotator ) from supervision.annotators.utils import ColorLookup from supervision.classification.core import Classifications diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index c2ac4d873..23ba95546 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1262,3 +1262,73 @@ def annotate( scene[y1:y2, x1:x2] = scaled_down_roi return scene + + +class TriangleAnnotator(BaseAnnotator): + """ + A class for drawing triangle markers on an image at specific coordinates based on + provided detections. + """ + + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + base: int = 10, + height: int = 10, + position: Position = Position.CENTER, + color_lookup: ColorLookup = ColorLookup.CLASS, + ): + """ + Args: + color (Union[Color, ColorPalette]): The color or color palette to use for + annotating detections. + base (int): The base width of the triangle. + height (int): The height of the triangle. + position (Position): The anchor position for placing the triangle. + color_lookup (ColorLookup): Strategy for mapping colors to annotations. + Options are `INDEX`, `CLASS`, `TRACK`. + """ + self.color: Union[Color, ColorPalette] = color + self.base: int = base + self.height: int = height + self.position: Position = position + self.color_lookup: ColorLookup = color_lookup + + def annotate( + self, + scene: np.ndarray, + detections: Detections, + custom_color_lookup: Optional[np.ndarray] = None, + ) -> np.ndarray: + """ + Annotates the given scene with triangles based on the provided detections. + + Args: + scene (np.ndarray): The image where triangles will be drawn. + detections (Detections): Object detections to annotate. + custom_color_lookup (Optional[np.ndarray]): Custom color lookup array. + Allows to override the default color mapping strategy. + + Returns: + np.ndarray: The annotated image. + """ + xy = detections.get_anchors_coordinates(anchor=self.position) + for detection_idx in range(len(detections)): + color = resolve_color( + color=self.color, + detections=detections, + detection_idx=detection_idx, + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, + ) + center_x, center_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1]) + vertices = np.array([ + [center_x, center_y - self.height // 2], + [center_x - self.base // 2, center_y + self.height // 2], + [center_x + self.base // 2, center_y + self.height // 2] + ], np.int32) + + cv2.fillPoly(scene, [vertices], color.as_bgr()) + + return scene From a618b7c36935950eef055d807d3d05970dbed159 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 6 Dec 2023 12:34:54 +0000 Subject: [PATCH 2/9] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/__init__.py | 2 +- supervision/annotators/core.py | 33 ++++++++++++++++++--------------- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 4d924acf1..85b07850a 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -21,7 +21,7 @@ PixelateAnnotator, PolygonAnnotator, TraceAnnotator, - TriangleAnnotator + TriangleAnnotator, ) from supervision.annotators.utils import ColorLookup from supervision.classification.core import Classifications diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 23ba95546..45241b970 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1271,12 +1271,12 @@ class TriangleAnnotator(BaseAnnotator): """ def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - base: int = 10, - height: int = 10, - position: Position = Position.CENTER, - color_lookup: ColorLookup = ColorLookup.CLASS, + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + base: int = 10, + height: int = 10, + position: Position = Position.CENTER, + color_lookup: ColorLookup = ColorLookup.CLASS, ): """ Args: @@ -1295,10 +1295,10 @@ def __init__( self.color_lookup: ColorLookup = color_lookup def annotate( - self, - scene: np.ndarray, - detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None, + self, + scene: np.ndarray, + detections: Detections, + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with triangles based on the provided detections. @@ -1323,11 +1323,14 @@ def annotate( else custom_color_lookup, ) center_x, center_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1]) - vertices = np.array([ - [center_x, center_y - self.height // 2], - [center_x - self.base // 2, center_y + self.height // 2], - [center_x + self.base // 2, center_y + self.height // 2] - ], np.int32) + vertices = np.array( + [ + [center_x, center_y - self.height // 2], + [center_x - self.base // 2, center_y + self.height // 2], + [center_x + self.base // 2, center_y + self.height // 2], + ], + np.int32, + ) cv2.fillPoly(scene, [vertices], color.as_bgr()) From 366d39db48535805d45e6ab7356b15f9d8651764 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 6 Dec 2023 13:42:52 +0100 Subject: [PATCH 3/9] Refactor vertices calculation in TriangleAnnotator --- supervision/annotators/core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 23ba95546..947d4949b 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1324,9 +1324,9 @@ def annotate( ) center_x, center_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1]) vertices = np.array([ - [center_x, center_y - self.height // 2], - [center_x - self.base // 2, center_y + self.height // 2], - [center_x + self.base // 2, center_y + self.height // 2] + [center_x - self.base // 2, center_y - self.height // 2], + [center_x + self.base // 2, center_y - self.height // 2], + [center_x, center_y + self.height // 2] ], np.int32) cv2.fillPoly(scene, [vertices], color.as_bgr()) From 8d9228f7fdca5f940e91894d002290ef2e61f024 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 6 Dec 2023 12:43:56 +0000 Subject: [PATCH 4/9] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/annotators/core.py | 33 ++++++++++++++++++--------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 947d4949b..0450dcc55 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1271,12 +1271,12 @@ class TriangleAnnotator(BaseAnnotator): """ def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - base: int = 10, - height: int = 10, - position: Position = Position.CENTER, - color_lookup: ColorLookup = ColorLookup.CLASS, + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + base: int = 10, + height: int = 10, + position: Position = Position.CENTER, + color_lookup: ColorLookup = ColorLookup.CLASS, ): """ Args: @@ -1295,10 +1295,10 @@ def __init__( self.color_lookup: ColorLookup = color_lookup def annotate( - self, - scene: np.ndarray, - detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None, + self, + scene: np.ndarray, + detections: Detections, + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with triangles based on the provided detections. @@ -1323,11 +1323,14 @@ def annotate( else custom_color_lookup, ) center_x, center_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1]) - vertices = np.array([ - [center_x - self.base // 2, center_y - self.height // 2], - [center_x + self.base // 2, center_y - self.height // 2], - [center_x, center_y + self.height // 2] - ], np.int32) + vertices = np.array( + [ + [center_x - self.base // 2, center_y - self.height // 2], + [center_x + self.base // 2, center_y - self.height // 2], + [center_x, center_y + self.height // 2], + ], + np.int32, + ) cv2.fillPoly(scene, [vertices], color.as_bgr()) From 84713239eeae14bdd57706fb5b196b1bcca5d31f Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 6 Dec 2023 13:45:56 +0100 Subject: [PATCH 5/9] Adjusted the default position value in the TriangleAnnotator class. --- supervision/annotators/core.py | 33 +++++++++++++++------------------ 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 0450dcc55..d39f29298 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1271,12 +1271,12 @@ class TriangleAnnotator(BaseAnnotator): """ def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - base: int = 10, - height: int = 10, - position: Position = Position.CENTER, - color_lookup: ColorLookup = ColorLookup.CLASS, + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + base: int = 10, + height: int = 10, + position: Position = Position.TOP_CENTER, + color_lookup: ColorLookup = ColorLookup.CLASS, ): """ Args: @@ -1295,10 +1295,10 @@ def __init__( self.color_lookup: ColorLookup = color_lookup def annotate( - self, - scene: np.ndarray, - detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None, + self, + scene: np.ndarray, + detections: Detections, + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with triangles based on the provided detections. @@ -1323,14 +1323,11 @@ def annotate( else custom_color_lookup, ) center_x, center_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1]) - vertices = np.array( - [ - [center_x - self.base // 2, center_y - self.height // 2], - [center_x + self.base // 2, center_y - self.height // 2], - [center_x, center_y + self.height // 2], - ], - np.int32, - ) + vertices = np.array([ + [center_x - self.base // 2, center_y - self.height // 2], + [center_x + self.base // 2, center_y - self.height // 2], + [center_x, center_y + self.height // 2] + ], np.int32) cv2.fillPoly(scene, [vertices], color.as_bgr()) From bf5d8ac29ce4c4b4403b8a29a5c15c3dc9e3d48b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 6 Dec 2023 12:46:11 +0000 Subject: [PATCH 6/9] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/annotators/core.py | 33 ++++++++++++++++++--------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index d39f29298..87d63e46d 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1271,12 +1271,12 @@ class TriangleAnnotator(BaseAnnotator): """ def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - base: int = 10, - height: int = 10, - position: Position = Position.TOP_CENTER, - color_lookup: ColorLookup = ColorLookup.CLASS, + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + base: int = 10, + height: int = 10, + position: Position = Position.TOP_CENTER, + color_lookup: ColorLookup = ColorLookup.CLASS, ): """ Args: @@ -1295,10 +1295,10 @@ def __init__( self.color_lookup: ColorLookup = color_lookup def annotate( - self, - scene: np.ndarray, - detections: Detections, - custom_color_lookup: Optional[np.ndarray] = None, + self, + scene: np.ndarray, + detections: Detections, + custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: """ Annotates the given scene with triangles based on the provided detections. @@ -1323,11 +1323,14 @@ def annotate( else custom_color_lookup, ) center_x, center_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1]) - vertices = np.array([ - [center_x - self.base // 2, center_y - self.height // 2], - [center_x + self.base // 2, center_y - self.height // 2], - [center_x, center_y + self.height // 2] - ], np.int32) + vertices = np.array( + [ + [center_x - self.base // 2, center_y - self.height // 2], + [center_x + self.base // 2, center_y - self.height // 2], + [center_x, center_y + self.height // 2], + ], + np.int32, + ) cv2.fillPoly(scene, [vertices], color.as_bgr()) From f73829e941a53c2349c167d99fb08df9097f2ba1 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 6 Dec 2023 13:51:01 +0100 Subject: [PATCH 7/9] The function was initially determining the center point of the triangle, which did not align with the desired annotation style. The changes now ensure that 'xy' refers to the coordinate of the triangle's tip instead of its center. --- supervision/annotators/core.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 87d63e46d..be495c886 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1322,15 +1322,12 @@ def annotate( if custom_color_lookup is None else custom_color_lookup, ) - center_x, center_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1]) - vertices = np.array( - [ - [center_x - self.base // 2, center_y - self.height // 2], - [center_x + self.base // 2, center_y - self.height // 2], - [center_x, center_y + self.height // 2], - ], - np.int32, - ) + tip_x, tip_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1]) + vertices = np.array([ + [tip_x - self.base // 2, tip_y - self.height], + [tip_x + self.base // 2, tip_y - self.height], + [tip_x, tip_y] + ], np.int32) cv2.fillPoly(scene, [vertices], color.as_bgr()) From c46ff15143e86fb80bd97d6aef9348e41fd28acf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 6 Dec 2023 12:52:06 +0000 Subject: [PATCH 8/9] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/annotators/core.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index be495c886..0defdf74c 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1323,11 +1323,14 @@ def annotate( else custom_color_lookup, ) tip_x, tip_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1]) - vertices = np.array([ - [tip_x - self.base // 2, tip_y - self.height], - [tip_x + self.base // 2, tip_y - self.height], - [tip_x, tip_y] - ], np.int32) + vertices = np.array( + [ + [tip_x - self.base // 2, tip_y - self.height], + [tip_x + self.base // 2, tip_y - self.height], + [tip_x, tip_y], + ], + np.int32, + ) cv2.fillPoly(scene, [vertices], color.as_bgr()) From 26ad007b7c9bc891501bfef60c52ace47a0fd10f Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 6 Dec 2023 13:59:05 +0100 Subject: [PATCH 9/9] Add TriangleAnnotator example to documentation --- docs/annotators.md | 25 +++++++++++++++++++++++++ supervision/annotators/core.py | 17 +++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/docs/annotators.md b/docs/annotators.md index b7bbcd1f8..f2dcd2765 100644 --- a/docs/annotators.md +++ b/docs/annotators.md @@ -103,6 +103,27 @@ +=== "Triangle" + + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> triangle_annotator = sv.TriangleAnnotator() + >>> annotated_frame = triangle_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + +
+ + ![triangle-annotator-example](https://media.roboflow.com/supervision-annotator-examples/triangle-annotator-example.png){ align=center width="800" } + +
+ === "Ellipse" ```python @@ -330,6 +351,10 @@ :::supervision.annotators.core.DotAnnotator +## TriangleAnnotator + +:::supervision.annotators.core.TriangleAnnotator + ## EllipseAnnotator :::supervision.annotators.core.EllipseAnnotator diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index be495c886..f61f818e7 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1311,6 +1311,23 @@ def annotate( Returns: np.ndarray: The annotated image. + + Example: + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> triangle_annotator = sv.TriangleAnnotator() + >>> annotated_frame = triangle_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + + ![triangle-annotator-example](https://media.roboflow.com/ + supervision-annotator-examples/triangle-annotator-example.png) """ xy = detections.get_anchors_coordinates(anchor=self.position) for detection_idx in range(len(detections)):