Skip to content

Commit 208bfe5

Browse files
Add docstrings
Signed-off-by: Ashwin Vaidya <[email protected]>
1 parent 7a98e5f commit 208bfe5

File tree

1 file changed

+42
-1
lines changed
  • model_api/python/model_api/visualizer/primitive

1 file changed

+42
-1
lines changed

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

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,21 @@
1717

1818

1919
class Polygon(Primitive):
20-
"""Polygon primitive."""
20+
"""Polygon primitive.
21+
22+
Args:
23+
points: List of points.
24+
mask: Mask to draw the polygon.
25+
color: Color of the polygon.
26+
27+
Examples:
28+
>>> polygon = Polygon(points=[(10, 10), (100, 10), (100, 100), (10, 100)], color="red")
29+
>>> polygon = Polygon(mask=mask, color="red")
30+
>>> polygon.compute(image).save("polygon.jpg")
31+
32+
>>> polygon = Polygon(mask=mask, color="red")
33+
>>> polygon.compute(image).save("polygon.jpg")
34+
"""
2135

2236
def __init__(
2337
self,
@@ -29,6 +43,17 @@ def __init__(
2943
self.color = color
3044

3145
def _get_points(self, points: list[tuple[int, int]] | None, mask: np.ndarray | None) -> list[tuple[int, int]]:
46+
"""Get points from either points or mask.
47+
Note:
48+
Either points or mask should be provided.
49+
50+
Args:
51+
points: List of points.
52+
mask: Mask to draw the polygon.
53+
54+
Returns:
55+
List of points.
56+
"""
3257
if points is not None and mask is not None:
3358
msg = "Either points or mask should be provided, not both."
3459
raise ValueError(msg)
@@ -42,11 +67,27 @@ def _get_points(self, points: list[tuple[int, int]] | None, mask: np.ndarray | N
4267
return points_
4368

4469
def _get_points_from_mask(self, mask: np.ndarray) -> list[tuple[int, int]]:
70+
"""Get points from mask.
71+
72+
Args:
73+
mask: Mask to draw the polygon.
74+
75+
Returns:
76+
List of points.
77+
"""
4578
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
4679
points_ = contours[0].squeeze().tolist()
4780
return [tuple(point) for point in points_]
4881

4982
def compute(self, image: Image) -> Image:
83+
"""Compute the polygon.
84+
85+
Args:
86+
image: Image to draw the polygon on.
87+
88+
Returns:
89+
Image with the polygon drawn on it.
90+
"""
5091
draw = ImageDraw.Draw(image)
5192
draw.polygon(self.points, fill=self.color)
5293
return image

0 commit comments

Comments
 (0)