diff --git a/docs/annotators.md b/docs/annotators.md index 2fb1ccee3..0626faadf 100644 --- a/docs/annotators.md +++ b/docs/annotators.md @@ -19,7 +19,7 @@ -=== "Mask" +=== "BoxCorner" ```python >>> import supervision as sv @@ -27,8 +27,8 @@ >>> image = ... >>> detections = sv.Detections(...) - >>> mask_annotator = sv.MaskAnnotator() - >>> annotated_frame = mask_annotator.annotate( + >>> corner_annotator = sv.BoxCornerAnnotator() + >>> annotated_frame = corner_annotator.annotate( ... scene=image.copy(), ... detections=detections ... ) @@ -36,11 +36,11 @@
- ![mask-annotator-example](https://media.roboflow.com/supervision-annotator-examples/mask-annotator-example.png){ align=center width="800" } + ![box-corner-annotator-example](https://media.roboflow.com/supervision-annotator-examples/box-corner-annotator-example.png){ align=center width="800" }
-=== "Ellipse" +=== "BoxMaskAnnotator" ```python >>> import supervision as sv @@ -48,8 +48,8 @@ >>> image = ... >>> detections = sv.Detections(...) - >>> ellipse_annotator = sv.EllipseAnnotator() - >>> annotated_frame = ellipse_annotator.annotate( + >>> box_mask_annotator = sv.BoxMaskAnnotator() + >>> annotated_frame = box_mask_annotator.annotate( ... scene=image.copy(), ... detections=detections ... ) @@ -57,11 +57,11 @@
- ![ellipse-annotator-example](https://media.roboflow.com/supervision-annotator-examples/ellipse-annotator-example.png){ align=center width="800" } + ![box-mask-annotator-example](https://media.roboflow.com/supervision-annotator-examples/box-mask-annotator-example.png){ align=center width="800" }
-=== "BoxCorner" +=== "Circle" ```python >>> import supervision as sv @@ -69,8 +69,8 @@ >>> image = ... >>> detections = sv.Detections(...) - >>> corner_annotator = sv.BoxCornerAnnotator() - >>> annotated_frame = corner_annotator.annotate( + >>> circle_annotator = sv.CircleAnnotator() + >>> annotated_frame = circle_annotator.annotate( ... scene=image.copy(), ... detections=detections ... ) @@ -78,11 +78,11 @@
- ![box-corner-annotator-example](https://media.roboflow.com/supervision-annotator-examples/box-corner-annotator-example.png){ align=center width="800" } + ![circle-annotator-example](https://media.roboflow.com/supervision-annotator-examples/circle-annotator-example.png){ align=center width="800" }
-=== "Circle" +=== "Ellipse" ```python >>> import supervision as sv @@ -90,8 +90,8 @@ >>> image = ... >>> detections = sv.Detections(...) - >>> circle_annotator = sv.CircleAnnotator() - >>> annotated_frame = circle_annotator.annotate( + >>> ellipse_annotator = sv.EllipseAnnotator() + >>> annotated_frame = ellipse_annotator.annotate( ... scene=image.copy(), ... detections=detections ... ) @@ -99,7 +99,28 @@
- ![circle-annotator-example](https://media.roboflow.com/supervision-annotator-examples/circle-annotator-example.png){ align=center width="800" } + ![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" }
@@ -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 -## LabelAnnotator +## EllipseAnnotator -:::supervision.annotators.core.LabelAnnotator +:::supervision.annotators.core.EllipseAnnotator -## TraceAnnotator +## MaskAnnotator -:::supervision.annotators.core.TraceAnnotator +:::supervision.annotators.core.MaskAnnotator + +## LabelAnnotator + +:::supervision.annotators.core.LabelAnnotator ## BlurAnnotator :::supervision.annotators.core.BlurAnnotator + +## TraceAnnotator + +:::supervision.annotators.core.TraceAnnotator diff --git a/supervision/__init__.py b/supervision/__init__.py index d6492af1a..2db639c57 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -10,6 +10,7 @@ BlurAnnotator, BoundingBoxAnnotator, BoxCornerAnnotator, + BoxMaskAnnotator, CircleAnnotator, EllipseAnnotator, LabelAnnotator, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 8edca92b5..6dd63c7a8 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -158,6 +158,79 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: 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. + 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`. + """ + 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 + ... ) + ``` + + ![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)): + 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.