Merge pull request roboflow#422 from hardikdava/box_mask_annotator

SkalskiP · web-flow · commit d30d80717c8b · 2023-10-07T00:47:09.000+02:00
Added box mask annotators
diff --git a/docs/annotators.md b/docs/annotators.md
@@ -19,87 +19,108 @@
 
     </div>
 
-=== "Mask"
+=== "BoxCorner"
 
     ```python
     >>> import supervision as sv
 
     >>> 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
     ... )
     ```
 
     <div class="result" markdown>
 
-    ![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" }
 
     </div>
 
-=== "Ellipse"
+=== "BoxMaskAnnotator"
 
     ```python
     >>> import supervision as sv
 
     >>> 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
     ... )
     ```
 
     <div class="result" markdown>
 
-    ![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" }
 
     </div>
 
-=== "BoxCorner"
+=== "Circle"
 
     ```python
     >>> import supervision as sv
 
     >>> 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
     ... )
     ```
 
     <div class="result" markdown>
 
-    ![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" }
 
     </div>
 
-=== "Circle"
+=== "Ellipse"
 
     ```python
     >>> import supervision as sv
 
     >>> 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
     ... )
     ```
 
     <div class="result" markdown>
 
-    ![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" }
+
+    </div>
+
+=== "Mask"
+
+    ```python
+    >>> import supervision as sv
+
+    >>> image = ...
+    >>> detections = sv.Detections(...)
+
+    >>> mask_annotator = sv.MaskAnnotator()
+    >>> annotated_frame = mask_annotator.annotate(
+    ...     scene=image.copy(),
+    ...     detections=detections
+    ... )
+    ```
+
+    <div class="result" markdown>
+
+    ![mask-annotator-example](https://media.roboflow.com/supervision-annotator-examples/mask-annotator-example.png){ align=center width="800" }
 
     </div>
 
@@ -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
@@ -10,6 +10,7 @@
     BlurAnnotator,
     BoundingBoxAnnotator,
     BoxCornerAnnotator,
+    BoxMaskAnnotator,
     CircleAnnotator,
     EllipseAnnotator,
     LabelAnnotator,
diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py
@@ -156,6 +156,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.
