Skip to content

Commit

Permalink
Add docstrings
Browse files Browse the repository at this point in the history
Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>
  • Loading branch information
ashwinvaidya17 committed Jan 16, 2025
1 parent 7a98e5f commit 208bfe5
Showing 1 changed file with 42 additions and 1 deletion.
43 changes: 42 additions & 1 deletion model_api/python/model_api/visualizer/primitive/polygon.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,21 @@


class Polygon(Primitive):
"""Polygon primitive."""
"""Polygon primitive.
Args:
points: List of points.
mask: Mask to draw the polygon.
color: Color of the polygon.
Examples:
>>> polygon = Polygon(points=[(10, 10), (100, 10), (100, 100), (10, 100)], color="red")
>>> polygon = Polygon(mask=mask, color="red")
>>> polygon.compute(image).save("polygon.jpg")
>>> polygon = Polygon(mask=mask, color="red")
>>> polygon.compute(image).save("polygon.jpg")
"""

def __init__(
self,
Expand All @@ -29,6 +43,17 @@ def __init__(
self.color = color

def _get_points(self, points: list[tuple[int, int]] | None, mask: np.ndarray | None) -> list[tuple[int, int]]:
"""Get points from either points or mask.
Note:
Either points or mask should be provided.
Args:
points: List of points.
mask: Mask to draw the polygon.
Returns:
List of points.
"""
if points is not None and mask is not None:
msg = "Either points or mask should be provided, not both."
raise ValueError(msg)
Expand All @@ -42,11 +67,27 @@ def _get_points(self, points: list[tuple[int, int]] | None, mask: np.ndarray | N
return points_

def _get_points_from_mask(self, mask: np.ndarray) -> list[tuple[int, int]]:
"""Get points from mask.
Args:
mask: Mask to draw the polygon.
Returns:
List of points.
"""
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:
"""Compute the polygon.
Args:
image: Image to draw the polygon on.
Returns:
Image with the polygon drawn on it.
"""
draw = ImageDraw.Draw(image)
draw.polygon(self.points, fill=self.color)
return image

0 comments on commit 208bfe5

Please sign in to comment.