Repository: kornia/kornia

Issue:
Equivalent function of cv2.fillConvexPoly in Kornia
### Discussed in https://github.com/kornia/kornia/discussions/1599

Include a tensor version of[ `cv2.fillConvexPoly`](https://docs.opencv.org/3.4/d6/d6e/group__imgproc__draw.html#ga9bb982be9d641dc51edd5e8ae3624e1f)

<div type='discussions-op-text'>

<sup>Originally posted by **vipulkumar-developer** March  4, 2022</sup>
Hi, I've started a new discussion to ask if there exists a method like cv2.fillConvexPoly in Kornia.</div>


Discussion:

Resolving patch:
diff --git a/kornia/utils/__init__.py b/kornia/utils/__init__.py
--- a/kornia/utils/__init__.py
+++ b/kornia/utils/__init__.py
@@ -1,4 +1,4 @@
-from .draw import draw_line, draw_rectangle
+from .draw import draw_convex_polygons, draw_line, draw_rectangle
 from .grid import create_meshgrid, create_meshgrid3d
 from .helpers import _extract_device_dtype, safe_inverse_with_mask, safe_solve_with_mask
 from .image import ImageToTensor, image_list_to_tensor, image_to_tensor, tensor_to_image
@@ -17,6 +17,7 @@
     "image_list_to_tensor",
     "save_pointcloud_ply",
     "load_pointcloud_ply",
+    "draw_convex_polygons",
     "draw_rectangle",
     "draw_line",
     "_extract_device_dtype",
diff --git a/kornia/utils/draw.py b/kornia/utils/draw.py
--- a/kornia/utils/draw.py
+++ b/kornia/utils/draw.py
@@ -1,6 +1,9 @@
-from typing import Optional
+from typing import List, Optional, Tuple, Union
 
 import torch
+from torch import Tensor
+
+from kornia.testing import KORNIA_CHECK, KORNIA_CHECK_SHAPE
 
 # TODO: implement width of the line
 
@@ -175,7 +178,6 @@ def draw_rectangle(
 
     if len(color.shape) == 1:
         color = color.expand(batch, num_rectangle, c)
-
     b, n, color_channels = color.shape
 
     if color_channels == 1 and c == 3:
@@ -205,3 +207,106 @@ def draw_rectangle(
                 ]
 
     return image
+
+
+def _get_convex_edges(polygon: Tensor, h: int, w: int) -> Tuple[Tensor, Tensor]:
+    r"""Gets the left and right edges of a polygon for each y-coordinate y \in [0, h)
+    Args:
+        polygons: represents polygons to draw in BxNx2
+            N is the number of points
+            2 is (x, y).
+        h: bottom most coordinate (top coordinate is assumed to be 0)
+        w: right most coordinate (left coordinate is assumed to be 0)
+    Returns:
+        The left and right edges of the polygon of shape (B,B).
+    """
+
+    dtype = polygon.dtype
+
+    # Check if polygons are in loop closed format, if not -> make it so
+    if not torch.allclose(polygon[..., -1, :], polygon[..., 0, :]):
+        polygon = torch.cat((polygon, polygon[..., :1, :]), dim=-2)  # (B, N+1, 2)
+
+    # Partition points into edges
+    x_start, y_start = polygon[..., :-1, 0], polygon[..., :-1, 1]
+    x_end, y_end = polygon[..., 1:, 0], polygon[..., 1:, 1]
+
+    # Create scanlines, edge dx/dy, and produce x values
+    ys = torch.arange(h, device=polygon.device, dtype=dtype)
+    dx = ((x_end - x_start) / (y_end - y_start + 1e-12)).clamp(-w, w)
+    xs = (ys[..., :, None] - y_start[..., None, :]) * dx[..., None, :] + x_start[..., None, :]
+
+    # Only count edge in their active regions (i.e between the vertices)
+    valid_edges = (y_start[..., None, :] <= ys[..., :, None]).logical_and(ys[..., :, None] <= y_end[..., None, :])
+    valid_edges |= (y_start[..., None, :] >= ys[..., :, None]).logical_and(ys[..., :, None] >= y_end[..., None, :])
+    x_left_edges = xs.clone()
+    x_left_edges[~valid_edges] = w
+    x_right_edges = xs.clone()
+    x_right_edges[~valid_edges] = -1
+
+    # Find smallest and largest x values for the valid edges
+    x_left = x_left_edges.min(dim=-1).values
+    x_right = x_right_edges.max(dim=-1).values
+    return x_left, x_right
+
+
+def _batch_polygons(polygons: List[Tensor]) -> Tensor:
+    r"""Converts a List of variable length polygons into a fixed size tensor.
+    Works by repeating the last element in the tensor.
+    Args:
+        polygon: List of variable length polygons of shape [N_1 x 2, N_2 x 2, ..., N_B x 2].
+                    B is the batch size,
+                    N_i is the number of points,
+                    2 is (x, y).
+    Returns:
+        A fixed size tensor of shape (B, N, 2) where N = max_i(N_i)
+    """
+    B, N = len(polygons), len(max(polygons, key=len))
+    batched_polygons = torch.zeros(B, N, 2, dtype=polygons[0].dtype, device=polygons[0].device)
+    for b, p in enumerate(polygons):
+        batched_polygons[b] = torch.cat((p, p[-1:].expand(N - len(p), 2))) if len(p) < N else p
+    return batched_polygons
+
+
+def draw_convex_polygons(images: Tensor, polygons: Union[Tensor, List[Tensor]], colors: Tensor) -> Tensor:
+    r"""Draws convex polygons on a batch of image tensors.
+
+    Args:
+        images: is tensor of BxCxHxW.
+        polygons: represents polygons as points, either BxNx2 or List of variable length polygons.
+            N is the number of points.
+            2 is (x, y).
+        color: a B x 3 tensor or 3 tensor with color to fill in.
+
+    Returns:
+        This operation modifies image inplace but also returns the drawn tensor for
+        convenience with same shape the of the input BxCxHxW.
+
+    Note:
+        This function assumes a coordinate system (0, h - 1), (0, w - 1) in the image, with (0, 0) being the center
+        of the top-left pixel and (w - 1, h - 1) being the center of the bottom-right coordinate.
+
+    Example:
+        >>> img = torch.rand(1, 3, 12, 16)
+        >>> poly = torch.tensor([[[4, 4], [12, 4], [12, 8], [4, 8]]])
+        >>> color = torch.tensor([[0.5, 0.5, 0.5]])
+        >>> out = draw_convex_polygons(img, poly, color)
+    """
+    # TODO: implement optional linetypes for smooth edges
+    KORNIA_CHECK_SHAPE(images, ["B", "C", "H", "W"])
+    b_i, c_i, h_i, w_i, device = *images.shape, images.device
+    if isinstance(polygons, List):
+        polygons = _batch_polygons(polygons)
+    b_p, _, xy, device_p, dtype_p = *polygons.shape, polygons.device, polygons.dtype
+    if len(colors.shape) == 1:
+        colors = colors.expand(b_i, c_i)
+    b_c, _, device_c = *colors.shape, colors.device
+    KORNIA_CHECK(xy == 2, "Polygon vertices must be xy, i.e. 2-dimensional")
+    KORNIA_CHECK(b_i == b_p == b_c, "Image, polygon, and color must have same batch dimension")
+    KORNIA_CHECK(device == device_p == device_c, "Image, polygon, and color must have same device")
+
+    x_left, x_right = _get_convex_edges(polygons, h_i, w_i)
+    ws = torch.arange(w_i, device=device, dtype=dtype_p)[None, None, :]
+    fill_region = (ws >= x_left[..., :, None]) & (ws <= x_right[..., :, None])
+    images = (~fill_region[:, None]) * images + fill_region[:, None] * colors[..., None, None]
+    return images


Test patch:
diff --git a/test/utils/test_draw.py b/test/utils/test_draw.py
--- a/test/utils/test_draw.py
+++ b/test/utils/test_draw.py
@@ -1,8 +1,10 @@
+import math
+
 import pytest
 import torch
 
 from kornia.testing import assert_close
-from kornia.utils import draw_rectangle
+from kornia.utils import create_meshgrid, draw_convex_polygons, draw_rectangle
 from kornia.utils.draw import draw_line
 
 
@@ -11,32 +13,44 @@ def test_draw_line_vertical(self, dtype, device):
         """Test drawing a vertical line."""
         img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
         img = draw_line(img, torch.tensor([6, 2]), torch.tensor([6, 0]), torch.tensor([255]))
-        img_mask = torch.tensor([[
-            [0., 0., 0., 0., 0., 0., 255., 0.],
-            [0., 0., 0., 0., 0., 0., 255., 0.],
-            [0., 0., 0., 0., 0., 0., 255., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-        ]], device=device, dtype=dtype)
+        img_mask = torch.tensor(
+            [
+                [
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                ]
+            ],
+            device=device,
+            dtype=dtype,
+        )
         assert_close(img, img_mask)
 
     def test_draw_line_horizontal(self, dtype, device):
         """Test drawing a horizontal line."""
         img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
         img = draw_line(img, torch.tensor([6, 4]), torch.tensor([0, 4]), torch.tensor([255]))
-        img_mask = torch.tensor([[
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [255., 255., 255., 255., 255., 255., 255., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-        ]], device=device, dtype=dtype)
+        img_mask = torch.tensor(
+            [
+                [
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [255.0, 255.0, 255.0, 255.0, 255.0, 255.0, 255.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                ]
+            ],
+            device=device,
+            dtype=dtype,
+        )
         assert_close(img, img_mask)
 
     def test_draw_line_with_big_coordinates(self, dtype, device):
@@ -51,73 +65,93 @@ def test_draw_line_m_lte_neg1(self, dtype, device):
         """Test drawing a line with m <= -1."""
         img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
         img = draw_line(img, torch.tensor([0, 7]), torch.tensor([6, 0]), torch.tensor([255]))
-        img_mask = torch.tensor([[
-            [0., 0., 0., 0., 0., 0., 255., 0.],
-            [0., 0., 0., 0., 0., 255., 0., 0.],
-            [0., 0., 0., 0., 255., 0., 0., 0.],
-            [0., 0., 0., 255., 0., 0., 0., 0.],
-            [0., 0., 0., 255., 0., 0., 0., 0.],
-            [0., 0., 255., 0., 0., 0., 0., 0.],
-            [0., 255., 0., 0., 0., 0., 0., 0.],
-            [255., 0., 0., 0., 0., 0., 0., 0.],
-        ]], device=device, dtype=dtype)
+        img_mask = torch.tensor(
+            [
+                [
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 255.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 255.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 255.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 255.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 255.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [255.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                ]
+            ],
+            device=device,
+            dtype=dtype,
+        )
         assert_close(img, img_mask)
 
     def test_draw_line_m_lt_0_gte_neg1(self, dtype, device):
         """Test drawing a line with -1 < m < 0."""
         img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
         img = draw_line(img, torch.tensor([1, 5]), torch.tensor([7, 0]), torch.tensor([255]))
-        img_mask = torch.tensor([[
-            [0., 0., 0., 0., 0., 0., 0., 255.],
-            [0., 0., 0., 0., 0., 0., 255., 0.],
-            [0., 0., 0., 0., 0., 255., 0., 0.],
-            [0., 0., 0., 255., 255., 0., 0., 0.],
-            [0., 0., 255., 0., 0., 0., 0., 0.],
-            [0., 255., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-        ]], device=device, dtype=dtype)
+        img_mask = torch.tensor(
+            [
+                [
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 255.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 255.0, 255.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 255.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 255.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                ]
+            ],
+            device=device,
+            dtype=dtype,
+        )
         assert_close(img, img_mask)
 
     def test_draw_line_m_gt_0_lt_1(self, dtype, device):
         """Test drawing a line with 0 < m < 1."""
         img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
         img = draw_line(img, torch.tensor([0, 0]), torch.tensor([6, 2]), torch.tensor([255]))
-        img_mask = torch.tensor([[
-            [255., 255., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 255., 255., 255., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 255., 255., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-        ]], device=device, dtype=dtype)
+        img_mask = torch.tensor(
+            [
+                [
+                    [255.0, 255.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 255.0, 255.0, 255.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 255.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                ]
+            ],
+            device=device,
+            dtype=dtype,
+        )
         assert_close(img, img_mask)
 
     def test_draw_line_m_gte_1(self, dtype, device):
         """Test drawing a line with m >= 1."""
         img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
         img = draw_line(img, torch.tensor([3, 7]), torch.tensor([1, 4]), torch.tensor([255]))
-        img_mask = torch.tensor([[
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 0., 0., 0., 0., 0.],
-            [0., 255., 0., 0., 0., 0., 0., 0.],
-            [0., 0., 255., 0., 0., 0., 0., 0.],
-            [0., 0., 255., 0., 0., 0., 0., 0.],
-            [0., 0., 0., 255., 0., 0., 0., 0.],
-        ]], device=device, dtype=dtype)
+        img_mask = torch.tensor(
+            [
+                [
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 255.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 255.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 255.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 255.0, 0.0, 0.0, 0.0, 0.0],
+                ]
+            ],
+            device=device,
+            dtype=dtype,
+        )
         assert_close(img, img_mask)
 
-    @pytest.mark.parametrize('p1',
-                             [
-                                 torch.tensor([-1, 0]),
-                                 torch.tensor([0, -1]),
-                                 torch.tensor([8, 0]),
-                                 torch.tensor([0, 8]),
-                             ])
+    @pytest.mark.parametrize(
+        'p1', [torch.tensor([-1, 0]), torch.tensor([0, -1]), torch.tensor([8, 0]), torch.tensor([0, 8])]
+    )
     def test_p1_out_of_bounds(self, p1, dtype, device):
         """Tests that an exception is raised if p1 is out of bounds."""
         img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
@@ -126,13 +160,9 @@ def test_p1_out_of_bounds(self, p1, dtype, device):
 
         assert 'p1 is out of bounds.' == str(excinfo.value)
 
-    @pytest.mark.parametrize('p2',
-                             [
-                                 torch.tensor([-1, 0]),
-                                 torch.tensor([0, -1]),
-                                 torch.tensor([8, 0]),
-                                 torch.tensor([0, 8]),
-                             ])
+    @pytest.mark.parametrize(
+        'p2', [torch.tensor([-1, 0]), torch.tensor([0, -1]), torch.tensor([8, 0]), torch.tensor([0, 8])]
+    )
     def test_p2_out_of_bounds(self, p2, dtype, device):
         """Tests that an exception is raised if p2 is out of bounds."""
         img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
@@ -149,11 +179,7 @@ def test_image_size(self, img_size, dtype, device):
 
         assert 'image must have 3 dimensions (C,H,W).' == str(excinfo.value)
 
-    @pytest.mark.parametrize('img_size,color',
-                             [
-                                 ((1, 8, 8), torch.tensor([23, 53])),
-                                 ((3, 8, 8), torch.tensor([255]))
-                             ])
+    @pytest.mark.parametrize('img_size,color', [((1, 8, 8), torch.tensor([23, 53])), ((3, 8, 8), torch.tensor([255]))])
     def test_color_image_channel_size(self, img_size, color, dtype, device):
         img = torch.zeros(*img_size, dtype=dtype, device=device)
         with pytest.raises(ValueError) as excinfo:
@@ -161,13 +187,15 @@ def test_color_image_channel_size(self, img_size, color, dtype, device):
 
         assert 'color must have the same number of channels as the image.' == str(excinfo.value)
 
-    @pytest.mark.parametrize('p1,p2',
-                             [
-                                 ((0, 1), (1, 5, 2)),
-                                 ((0, 1, 2), (1, 5)),
-                                 (torch.tensor([0, 1]), torch.tensor([0, 2, 3])),
-                                 (torch.tensor([0, 1, 5]), torch.tensor([0, 2])),
-                             ])
+    @pytest.mark.parametrize(
+        'p1,p2',
+        [
+            ((0, 1), (1, 5, 2)),
+            ((0, 1, 2), (1, 5)),
+            (torch.tensor([0, 1]), torch.tensor([0, 2, 3])),
+            (torch.tensor([0, 1, 5]), torch.tensor([0, 2])),
+        ],
+    )
     def test_point_size(self, p1, p2, dtype, device):
         img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
         with pytest.raises(ValueError) as excinfo:
@@ -237,8 +265,8 @@ def test_n_rectangles(self, device, batch, N, fill):
                         random_w_rectangle[
                             b,
                             :,
-                            points_list[b][n][1]: points_list[b][n][3] + 1,
-                            points_list[b][n][0]: points_list[b][n][2] + 1,
+                            points_list[b][n][1] : points_list[b][n][3] + 1,
+                            points_list[b][n][0] : points_list[b][n][2] + 1,
                         ].sum()
                         == (points_list[b][n][3] - points_list[b][n][1] + 1)
                         * (points_list[b][n][2] - points_list[b][n][0] + 1)
@@ -247,25 +275,25 @@ def test_n_rectangles(self, device, batch, N, fill):
                 else:
                     assert (
                         random_w_rectangle[
-                            b, :, points_list[b][n][1]: points_list[b][n][3] + 1, points_list[b][n][0]
+                            b, :, points_list[b][n][1] : points_list[b][n][3] + 1, points_list[b][n][0]
                         ].sum()
                         == (points_list[b][n][3] - points_list[b][n][1] + 1) * 3
                     )
                     assert (
                         random_w_rectangle[
-                            b, :, points_list[b][n][1]: points_list[b][n][3] + 1, points_list[b][n][2]
+                            b, :, points_list[b][n][1] : points_list[b][n][3] + 1, points_list[b][n][2]
                         ].sum()
                         == (points_list[b][n][3] - points_list[b][n][1] + 1) * 3
                     )
                     assert (
                         random_w_rectangle[
-                            b, :, points_list[b][n][1], points_list[b][n][0]: points_list[b][n][2] + 1
+                            b, :, points_list[b][n][1], points_list[b][n][0] : points_list[b][n][2] + 1
                         ].sum()
                         == (points_list[b][n][2] - points_list[b][n][0] + 1) * 3
                     )
                     assert (
                         random_w_rectangle[
-                            b, :, points_list[b][n][1], points_list[b][n][0]: points_list[b][n][2] + 1
+                            b, :, points_list[b][n][1], points_list[b][n][0] : points_list[b][n][2] + 1
                         ].sum()
                         == (points_list[b][n][2] - points_list[b][n][0] + 1) * 3
                     )
@@ -312,3 +340,103 @@ def test_color_foreground(self, device, color):
             )
             <= 0.0001
         )
+
+
+class TestFillConvexPolygon:
+    def test_circle(self, device, dtype):
+        b, c, h, w = 1, 3, 500, 500
+        n = 5000
+        im = torch.zeros(b, c, h, w, device=device, dtype=dtype)
+        t = torch.linspace(0, 1, steps=n)[None].expand(b, n)
+        color = torch.tensor([1, 1, 1], device=device, dtype=dtype)[None].expand(b, c)
+        x = (2 * math.pi * t).cos()
+        y = (2 * math.pi * t).sin()
+        ctr = 200
+        radius = 200
+        pts = ctr + radius * torch.stack((x, y), dim=-1)
+        poly_im = draw_convex_polygons(im, pts, color)
+        XY = create_meshgrid(h, w, normalized_coordinates=False, device=device, dtype=dtype)
+        inside = (((XY[..., 1] - ctr) ** 2 + (XY[..., 0] - ctr) ** 2).sqrt() <= radius)[:, None].expand(b, c, h, w)
+        circ_im = inside * color[..., None, None]
+        assert (circ_im - poly_im).abs().mean() <= 1e-4
+
+    def test_ellipse(self, device, dtype):
+        b, c, h, w = 1, 3, 500, 500
+        n = 5000
+        im = torch.zeros(b, c, h, w, device=device, dtype=dtype)
+        t = torch.linspace(0, 1, steps=n)[None].expand(b, n)
+        color = torch.tensor([1, 1, 1], device=device, dtype=dtype)[None].expand(b, c)
+        lam = 2
+        x = lam * (2 * math.pi * t).cos()
+        y = (2 * math.pi * t).sin()
+        ctr = 200
+        radius = 100
+        pts = ctr + radius * torch.stack((x, y), dim=-1)
+        poly_im = draw_convex_polygons(im, pts, color)
+        XY = create_meshgrid(h, w, normalized_coordinates=False, device=device, dtype=dtype)
+        inside = (((XY[..., 1] - ctr) ** 2 + ((XY[..., 0] - ctr) / lam) ** 2).sqrt() <= radius)[:, None].expand(
+            b, c, h, w
+        )
+        ellipse_im = inside * color[..., None, None]
+        assert (ellipse_im - poly_im).abs().mean() <= 1e-4
+
+    def test_rectangle(self, device, dtype):
+        b, c, h, w = 1, 3, 500, 500
+        im = torch.zeros(b, c, h, w, device=device, dtype=dtype)
+        color = torch.tensor([1, 1, 1], device=device, dtype=dtype)[None].expand(b, c)
+        pts = torch.tensor([[[50, 50], [200, 50], [200, 250], [50, 250]]], device=device, dtype=dtype)
+        poly_im = draw_convex_polygons(im.clone(), pts, color)
+        rect = torch.cat((pts[..., 0, :], pts[..., 2, :]), dim=-1)[:, None]
+        rect_im = draw_rectangle(im.clone(), rect, color[:, None], fill=True)
+        assert torch.allclose(rect_im, poly_im)
+
+    def test_batch(self, device, dtype):
+        im = torch.rand(2, 3, 12, 16, dtype=dtype, device=device)
+        pts = torch.tensor(
+            [[[4, 4], [12, 4], [12, 8], [4, 8]], [[0, 0], [4, 0], [4, 4], [0, 4]]], dtype=dtype, device=device
+        )
+        torch.tensor([[[0, 0, 4, 4]], [[4, 4, 10, 10]]])
+        color = torch.tensor([[0.5, 0.5, 0.5], [0.5, 0.5, 0.75]])
+        poly_im = draw_convex_polygons(im.clone(), pts, color)
+        rect = torch.tensor([[[4, 4, 12, 8]], [[0, 0, 4, 4]]], dtype=dtype, device=device)
+        rect_im = draw_rectangle(im.clone(), rect, color[:, None], fill=True)
+        assert torch.allclose(rect_im, poly_im)
+
+    def test_batch_variable_size(self, device, dtype):
+        im = torch.rand(2, 3, 12, 16, dtype=dtype, device=device)
+        pts = [
+            torch.tensor([[4, 4], [12, 4], [12, 8], [4, 8]], dtype=dtype, device=device),
+            torch.tensor([[0, 0], [2, 0], [4, 0], [4, 4], [0, 4]], dtype=dtype, device=device),
+        ]
+        color = torch.tensor([[0.5, 0.5, 0.5], [0.5, 0.5, 0.75]])
+        poly_im = draw_convex_polygons(im.clone(), pts, color)
+        rect = torch.tensor([[[4, 4, 12, 8]], [[0, 0, 4, 4]]], dtype=dtype, device=device)
+        rect_im = draw_rectangle(im.clone(), rect, color[:, None], fill=True)
+        assert torch.allclose(rect_im, poly_im)
+
+    def test_batch_color_no_batch(self, device, dtype):
+        im = torch.rand(2, 3, 12, 16, dtype=dtype, device=device)
+        pts = [
+            torch.tensor([[4, 4], [12, 4], [12, 8], [4, 8]], dtype=dtype, device=device),
+            torch.tensor([[0, 0], [2, 0], [4, 0], [4, 4], [0, 4]], dtype=dtype, device=device),
+        ]
+        color = torch.tensor([0.5, 0.5, 0.75])
+        poly_im = draw_convex_polygons(im.clone(), pts, color)
+        rect = torch.tensor([[[4, 4, 12, 8]], [[0, 0, 4, 4]]], dtype=dtype, device=device)
+        rect_im = draw_rectangle(im.clone(), rect, color, fill=True)
+        assert torch.allclose(rect_im, poly_im)
+
+    def test_out_of_bounds_rectangle(self, device, dtype):
+        b, c, h, w = 1, 3, 500, 500
+        im = torch.zeros(b, c, h, w, device=device, dtype=dtype)
+        color = torch.tensor([1, 1, 1], device=device, dtype=dtype)[None].expand(b, c)
+        pts = 350 + torch.tensor([[[50, 50], [200, 50], [200, 250], [50, 250]]], device=device, dtype=dtype)
+        poly_im = draw_convex_polygons(im.clone(), pts, color)
+        rect = torch.cat((pts[..., 0, :], pts[..., 2, :]), dim=-1)[:, None]
+        rect_im = draw_rectangle(im.clone(), rect, color[:, None], fill=True)
+        assert torch.allclose(rect_im, poly_im)
+        pts = -150 + torch.tensor([[[50, 50], [200, 50], [200, 250], [50, 250]]], device=device, dtype=dtype)
+        poly_im = draw_convex_polygons(im.clone(), pts, color)
+        rect = torch.cat((pts[..., 0, :], pts[..., 2, :]), dim=-1)[:, None]
+        rect_im = draw_rectangle(im.clone(), rect, color[:, None], fill=True)
+        assert torch.allclose(rect_im, poly_im)