Add polygon primitive

Signed-off-by: Ashwin Vaidya <[email protected]>

Author: ashwinvaidya17

model_api/python/model_api/visualizer/__init__.py

@@ -4,8 +4,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .layout import Flatten, HStack, Layout
-from .primitive import BoundingBox, Overlay
+from .primitive import BoundingBox, Overlay, Polygon
 from .scene import Scene
 from .visualizer import Visualizer
 
-__all__ = ["BoundingBox", "Overlay", "Scene", "Visualizer", "Layout", "Flatten", "HStack"]
+__all__ = ["BoundingBox", "Overlay", "Polygon", "Scene", "Visualizer", "Layout", "Flatten", "HStack"]

model_api/python/model_api/visualizer/primitive/__init__.py

@@ -5,6 +5,7 @@
 
 from .bounding_box import BoundingBox
 from .overlay import Overlay
+from .polygon import Polygon
 from .primitive import Primitive
 
-__all__ = ["Primitive", "BoundingBox", "Overlay"]
+__all__ = ["Primitive", "BoundingBox", "Overlay", "Polygon"]

model_api/python/model_api/visualizer/primitive/polygon.py

@@ -0,0 +1,52 @@
+"""Polygon primitive."""
+
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import cv2
+from PIL import Image, ImageDraw
+
+from .primitive import Primitive
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+class Polygon(Primitive):
+    """Polygon primitive."""
+
+    def __init__(
+        self,
+        points: list[tuple[int, int]] | None = None,
+        mask: np.ndarray | None = None,
+        color: str | tuple[int, int, int] = "blue",
+    ) -> None:
+        self.points = self._get_points(points, mask)
+        self.color = color
+
+    def _get_points(self, points: list[tuple[int, int]] | None, mask: np.ndarray | None) -> list[tuple[int, int]]:
+        if points is not None and mask is not None:
+            msg = "Either points or mask should be provided, not both."
+            raise ValueError(msg)
+        if points is not None:
+            points_ = points
+        elif mask is not None:
+            points_ = self._get_points_from_mask(mask)
+        else:
+            msg = "Either points or mask should be provided."
+            raise ValueError(msg)
+        return points_
+
+    def _get_points_from_mask(self, mask: np.ndarray) -> list[tuple[int, int]]:
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        points_ = contours[0].squeeze().tolist()
+        return [tuple(point) for point in points_]
+
+    def compute(self, image: Image) -> Image:
+        draw = ImageDraw.Draw(image)
+        draw.polygon(self.points, fill=self.color)
+        return image

tests/python/unit/visualizer/test_primitive.py

@@ -7,7 +7,7 @@
 import PIL
 from PIL import ImageDraw
 
-from model_api.visualizer import BoundingBox, Overlay
+from model_api.visualizer import BoundingBox, Overlay, Polygon
 
 
 def test_overlay(mock_image: PIL.Image):
@@ -32,3 +32,22 @@ def test_bounding_box(mock_image: PIL.Image):
     draw.rectangle((10, 10, 100, 100), outline="blue", width=2)
     bounding_box = BoundingBox(x1=10, y1=10, x2=100, y2=100)
     assert bounding_box.compute(mock_image) == expected_image
+
+
+def test_polygon(mock_image: PIL.Image):
+    """Test if the polygon is created correctly."""
+    # Test from points
+    expected_image = mock_image.copy()
+    draw = ImageDraw.Draw(expected_image)
+    draw.polygon([(10, 10), (100, 10), (100, 100), (10, 100)], fill="red")
+    polygon = Polygon(points=[(10, 10), (100, 10), (100, 100), (10, 100)], color="red")
+    assert polygon.compute(mock_image) == expected_image
+
+    # Test from mask
+    mask = np.zeros((100, 100), dtype=np.uint8)
+    mask[10:100, 10:100] = 255
+    expected_image = mock_image.copy()
+    draw = ImageDraw.Draw(expected_image)
+    draw.polygon([(10, 10), (100, 10), (100, 100), (10, 100)], fill="red")
+    polygon = Polygon(mask=mask, color="red")
+    assert polygon.compute(mock_image) == expected_image