Merge pull request #1527 from roboflow/feature/add-opacity-to-polygon…

…-zone-annotator

Add opacity to PolygonZoneAnnotator

docs/utils/draw.md

@@ -28,6 +28,12 @@ comments: true
 
 :::supervision.draw.utils.draw_polygon
 
+<div class="md-typeset">
+    <h2><a href="#supervision.draw.utils.draw_filled_polygon">draw_filled_polygon</a></h2>
+</div>
+
+:::supervision.draw.utils.draw_filled_polygon
+
 <div class="md-typeset">
     <h2><a href="#supervision.draw.utils.draw_text">draw_text</a></h2>
 </div>

supervision/__init__.py

@@ -76,6 +76,7 @@
 from supervision.draw.utils import (
     calculate_optimal_line_thickness,
     calculate_optimal_text_scale,
+    draw_filled_polygon,
     draw_filled_rectangle,
     draw_image,
     draw_line,

supervision/detection/tools/polygon_zone.py

@@ -9,7 +9,7 @@
 from supervision import Detections
 from supervision.detection.utils import clip_boxes, polygon_to_mask
 from supervision.draw.color import Color
-from supervision.draw.utils import draw_polygon, draw_text
+from supervision.draw.utils import draw_filled_polygon, draw_polygon, draw_text
 from supervision.geometry.core import Position
 from supervision.geometry.utils import get_polygon_center
 from supervision.utils.internal import SupervisionWarnings
@@ -109,6 +109,7 @@ class PolygonZoneAnnotator:
             default is cv2.FONT_HERSHEY_SIMPLEX
         center (Tuple[int, int]): The center of the polygon for text placement
         display_in_zone_count (bool): Show the label of the zone or not. Default is True
+        opacity: The opacity of zone filling when drawn on the scene. Default is 0
     """
 
     def __init__(
@@ -121,6 +122,7 @@ def __init__(
         text_thickness: int = 1,
         text_padding: int = 10,
         display_in_zone_count: bool = True,
+        opacity: float = 0,
     ):
         self.zone = zone
         self.color = color
@@ -132,6 +134,7 @@ def __init__(
         self.font = cv2.FONT_HERSHEY_SIMPLEX
         self.center = get_polygon_center(polygon=zone.polygon)
         self.display_in_zone_count = display_in_zone_count
+        self.opacity = opacity
 
     def annotate(self, scene: np.ndarray, label: Optional[str] = None) -> np.ndarray:
         """
@@ -145,12 +148,26 @@ def annotate(self, scene: np.ndarray, label: Optional[str] = None) -> np.ndarray
         Returns:
             np.ndarray: The image with the polygon zone and count of detected objects
         """
-        annotated_frame = draw_polygon(
-            scene=scene,
-            polygon=self.zone.polygon,
-            color=self.color,
-            thickness=self.thickness,
-        )
+        if self.opacity == 0:
+            annotated_frame = draw_polygon(
+                scene=scene,
+                polygon=self.zone.polygon,
+                color=self.color,
+                thickness=self.thickness,
+            )
+        else:
+            annotated_frame = draw_filled_polygon(
+                scene=scene.copy(),
+                polygon=self.zone.polygon,
+                color=self.color,
+                opacity=self.opacity,
+            )
+            annotated_frame = draw_polygon(
+                scene=annotated_frame,
+                polygon=self.zone.polygon,
+                color=self.color,
+                thickness=self.thickness,
+            )
 
         if self.display_in_zone_count:
             annotated_frame = draw_text(

supervision/draw/utils.py

@@ -59,25 +59,42 @@ def draw_rectangle(
     return scene
 
 
-def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.ndarray:
+def draw_filled_rectangle(
+    scene: np.ndarray, rect: Rect, color: Color, opacity: float = 1
+) -> np.ndarray:
     """
     Draws a filled rectangle on an image.
 
     Parameters:
         scene (np.ndarray): The scene on which the rectangle will be drawn
         rect (Rect): The rectangle to be drawn
         color (Color): The color of the rectangle
+        opacity (float): The opacity of rectangle when drawn on the scene.
 
     Returns:
         np.ndarray: The scene with the rectangle drawn on it
     """
-    cv2.rectangle(
-        scene,
-        rect.top_left.as_xy_int_tuple(),
-        rect.bottom_right.as_xy_int_tuple(),
-        color.as_bgr(),
-        -1,
-    )
+    if opacity == 1:
+        cv2.rectangle(
+            scene,
+            rect.top_left.as_xy_int_tuple(),
+            rect.bottom_right.as_xy_int_tuple(),
+            color.as_bgr(),
+            -1,
+        )
+    else:
+        scene_with_annotations = scene.copy()
+        cv2.rectangle(
+            scene_with_annotations,
+            rect.top_left.as_xy_int_tuple(),
+            rect.bottom_right.as_xy_int_tuple(),
+            color.as_bgr(),
+            -1,
+        )
+        cv2.addWeighted(
+            scene_with_annotations, opacity, scene, 1 - opacity, gamma=0, dst=scene
+        )
+
     return scene
 
 
@@ -153,6 +170,32 @@ def draw_polygon(
     return scene
 
 
+def draw_filled_polygon(
+    scene: np.ndarray, polygon: np.ndarray, color: Color, opacity: float = 1
+) -> np.ndarray:
+    """Draw a filled polygon on a scene.
+
+    Parameters:
+        scene (np.ndarray): The scene to draw the polygon on.
+        polygon (np.ndarray): The polygon to be drawn, given as a list of vertices.
+        color (Color): The color of the polygon.
+        opacity (float): The opacity of polygon when drawn on the scene.
+
+    Returns:
+        np.ndarray: The scene with the polygon drawn on it.
+    """
+    if opacity == 1:
+        cv2.fillPoly(scene, [polygon], color=color.as_bgr())
+    else:
+        scene_with_annotations = scene.copy()
+        cv2.fillPoly(scene_with_annotations, [polygon], color=color.as_bgr())
+        cv2.addWeighted(
+            scene_with_annotations, opacity, scene, 1 - opacity, gamma=0, dst=scene
+        )
+
+    return scene
+
+
 def draw_text(
     scene: np.ndarray,
     text: str,

test/detection/test_polygon_zone_annotator.py

@@ -0,0 +1,60 @@
+import numpy as np
+import pytest
+
+import supervision as sv
+
+COLOR = sv.Color(r=255, g=0, b=0)
+THICKNESS = 2
+POLYGON = np.array([[100, 100], [200, 100], [200, 200], [100, 200]])
+SCENE = np.random.randint(0, 255, (1000, 1000, 3), dtype=np.uint8)
+ANNOTATED_SCENE_NO_OPACITY = sv.draw_polygon(
+    scene=SCENE.copy(),
+    polygon=POLYGON,
+    color=COLOR,
+    thickness=THICKNESS,
+)
+ANNOTATED_SCENE_0_5_OPACITY = sv.draw_filled_polygon(
+    scene=ANNOTATED_SCENE_NO_OPACITY.copy(),
+    polygon=POLYGON,
+    color=COLOR,
+    opacity=0.5,
+)
+
+
+@pytest.mark.parametrize(
+    "scene, polygon_zone_annotator, expected_results",
+    [
+        (
+            SCENE,
+            sv.PolygonZoneAnnotator(
+                zone=sv.PolygonZone(
+                    POLYGON,
+                ),
+                color=COLOR,
+                thickness=THICKNESS,
+                display_in_zone_count=False,
+            ),
+            ANNOTATED_SCENE_NO_OPACITY,
+        ),  # Test no opacity (default)
+        (
+            SCENE,
+            sv.PolygonZoneAnnotator(
+                zone=sv.PolygonZone(
+                    POLYGON,
+                ),
+                color=COLOR,
+                thickness=THICKNESS,
+                display_in_zone_count=False,
+                opacity=0.5,
+            ),
+            ANNOTATED_SCENE_0_5_OPACITY,
+        ),  # Test 10% opacity
+    ],
+)
+def test_polygon_zone_annotator(
+    scene: np.ndarray,
+    polygon_zone_annotator: sv.PolygonZoneAnnotator,
+    expected_results: np.ndarray,
+) -> None:
+    annotated_scene = polygon_zone_annotator.annotate(scene=scene)
+    assert np.all(annotated_scene == expected_results)