From ef43640c88bbab0d1d47fbe65fa43842c8ebdbcf Mon Sep 17 00:00:00 2001 From: Hardik Dava Date: Thu, 5 Oct 2023 17:56:32 +0200 Subject: [PATCH 1/5] Added box mask annotators --- supervision/__init__.py | 1 + supervision/annotators/core.py | 72 ++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/supervision/__init__.py b/supervision/__init__.py index d6492af1..2db639c5 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -10,6 +10,7 @@ from supervision.annotators.core import ( BlurAnnotator, BoundingBoxAnnotator, BoxCornerAnnotator, + BoxMaskAnnotator, CircleAnnotator, EllipseAnnotator, LabelAnnotator, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 8edca92b..7de9d5d7 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -158,6 +158,78 @@ class MaskAnnotator(BaseAnnotator): return scene +class BoxMaskAnnotator(BaseAnnotator): + """ + A class for drawing box masks on an image using provided detections. + """ + + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + opacity: float = 0.5, + color_map: str = "class", + ): + """ + Args: + color (Union[Color, ColorPalette]): The color or color palette to use for + annotating detections. + color_map (str): Strategy for mapping colors to annotations. + Options are `index`, `class`, or `track`. + """ + self.color: Union[Color, ColorPalette] = color + self.color_map: ColorMap = ColorMap(color_map) + self.opacity = opacity + + def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: + """ + Annotates the given scene with box masks based on the provided detections. + + Args: + scene (np.ndarray): The image where bounding boxes will be drawn. + detections (Detections): Object detections to annotate. + + Returns: + np.ndarray: The annotated image. + + Example: + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> box_mask_annotator = sv.BoxMaskAnnotator() + >>> annotated_frame = box_mask_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + + ![bounding-box-annotator-example](https://media.roboflow.com/ + supervision-annotator-examples/bounding-box-annotator-example.png) + """ + mask_image = scene.copy() + for detection_idx in range(len(detections)): + x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int) + idx = resolve_color_idx( + detections=detections, + detection_idx=detection_idx, + color_map=self.color_map, + ) + color = resolve_color(color=self.color, idx=idx) + cv2.rectangle( + img=scene, + pt1=(x1, y1), + pt2=(x2, y2), + color=color.as_bgr(), + thickness=-1, + ) + scene = cv2.addWeighted( + scene, self.opacity, mask_image, 1 - self.opacity, gamma=0 + ) + return scene + + class EllipseAnnotator(BaseAnnotator): """ A class for drawing ellipses on an image using provided detections. From e7c242bdd47f2ca12fd6a91c5f3cde49acc99cd2 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 6 Oct 2023 20:51:31 +0200 Subject: [PATCH 2/5] Update annotators in docs and code --- docs/annotators.md | 133 ++++++++++++++++++++------------- supervision/annotators/core.py | 4 +- 2 files changed, 81 insertions(+), 56 deletions(-) diff --git a/docs/annotators.md b/docs/annotators.md index 2fb1ccee..dc4d956f 100644 --- a/docs/annotators.md +++ b/docs/annotators.md @@ -19,48 +19,6 @@ -=== "Mask" - - ```python - >>> import supervision as sv - - >>> image = ... - >>> detections = sv.Detections(...) - - >>> mask_annotator = sv.MaskAnnotator() - >>> annotated_frame = mask_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) - ``` - -
- - ![mask-annotator-example](https://media.roboflow.com/supervision-annotator-examples/mask-annotator-example.png){ align=center width="800" } - -
- -=== "Ellipse" - - ```python - >>> import supervision as sv - - >>> image = ... - >>> detections = sv.Detections(...) - - >>> ellipse_annotator = sv.EllipseAnnotator() - >>> annotated_frame = ellipse_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) - ``` - -
- - ![ellipse-annotator-example](https://media.roboflow.com/supervision-annotator-examples/ellipse-annotator-example.png){ align=center width="800" } - -
- === "BoxCorner" ```python @@ -82,6 +40,27 @@ +=== "BoxMaskAnnotator" + + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> box_mask_annotator = sv.BoxMaskAnnotator() + >>> annotated_frame = box_mask_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + +
+ + ![box-mask-annotator-example](https://media.roboflow.com/supervision-annotator-examples/box-mask-annotator-example.png){ align=center width="800" } + +
+ === "Circle" ```python @@ -103,6 +82,48 @@ +=== "Ellipse" + + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> ellipse_annotator = sv.EllipseAnnotator() + >>> annotated_frame = ellipse_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + +
+ + ![ellipse-annotator-example](https://media.roboflow.com/supervision-annotator-examples/ellipse-annotator-example.png){ align=center width="800" } + +
+ +=== "Mask" + + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> mask_annotator = sv.MaskAnnotator() + >>> annotated_frame = mask_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + +
+ + ![mask-annotator-example](https://media.roboflow.com/supervision-annotator-examples/mask-annotator-example.png){ align=center width="800" } + +
+ === "Label" ```python @@ -170,30 +191,34 @@ :::supervision.annotators.core.BoundingBoxAnnotator -## MaskAnnotator - -:::supervision.annotators.core.MaskAnnotator - -## EllipseAnnotator - -:::supervision.annotators.core.EllipseAnnotator - ## BoxCornerAnnotator :::supervision.annotators.core.BoxCornerAnnotator +## BoxMaskAnnotator + +:::supervision.annotators.core.BoxMaskAnnotator + ## CircleAnnotator :::supervision.annotators.core.CircleAnnotator +## EllipseAnnotator + +:::supervision.annotators.core.EllipseAnnotator + +## MaskAnnotator + +:::supervision.annotators.core.MaskAnnotator + ## LabelAnnotator :::supervision.annotators.core.LabelAnnotator -## TraceAnnotator - -:::supervision.annotators.core.TraceAnnotator - ## BlurAnnotator :::supervision.annotators.core.BlurAnnotator + +## TraceAnnotator + +:::supervision.annotators.core.TraceAnnotator \ No newline at end of file diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 7de9d5d7..3fa26bf8 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -205,8 +205,8 @@ class BoxMaskAnnotator(BaseAnnotator): ... ) ``` - ![bounding-box-annotator-example](https://media.roboflow.com/ - supervision-annotator-examples/bounding-box-annotator-example.png) + ![box-mask-annotator-example](https://media.roboflow.com/ + supervision-annotator-examples/box-mask-annotator-example.png) """ mask_image = scene.copy() for detection_idx in range(len(detections)): From cc93bba741fcda543d2e020370c25550b6969527 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 6 Oct 2023 18:52:00 +0000 Subject: [PATCH 3/5] =?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 --- docs/annotators.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/annotators.md b/docs/annotators.md index dc4d956f..0626faad 100644 --- a/docs/annotators.md +++ b/docs/annotators.md @@ -221,4 +221,4 @@ ## TraceAnnotator -:::supervision.annotators.core.TraceAnnotator \ No newline at end of file +:::supervision.annotators.core.TraceAnnotator From bed28e3db4f3efd134bfb6ec8f60a53ca21f9b15 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 6 Oct 2023 20:54:09 +0200 Subject: [PATCH 4/5] Improve the documentation for annotators class --- supervision/annotators/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 3fa26bf8..ff3ec602 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -162,7 +162,6 @@ class BoxMaskAnnotator(BaseAnnotator): """ A class for drawing box masks on an image using provided detections. """ - def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), @@ -173,6 +172,7 @@ class BoxMaskAnnotator(BaseAnnotator): Args: color (Union[Color, ColorPalette]): The color or color palette to use for annotating detections. + opacity (float): Opacity of the overlay mask. Must be between `0` and `1`. color_map (str): Strategy for mapping colors to annotations. Options are `index`, `class`, or `track`. """ From 669c0658fee276fc454925aa4f9c6f99984e25f5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 6 Oct 2023 18:54:40 +0000 Subject: [PATCH 5/5] =?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 | 1 + 1 file changed, 1 insertion(+) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index ff3ec602..6dd63c7a 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -162,6 +162,7 @@ class BoxMaskAnnotator(BaseAnnotator): """ A class for drawing box masks on an image using provided detections. """ + def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(),