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, device=device,
dtype=dtype)[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_polygon(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_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_polygon(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_close(rect_im, poly_im)
pts = -150 + torch.tensor(
[[[50, 50], [200, 50], [200, 250], [50, 250]]],
device=device,
dtype=dtype)
poly_im = draw_convex_polygon(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_close(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)
color = torch.tensor([[0.5, 0.5, 0.5], [0.5, 0.5, 0.75]],
dtype=dtype,
device=device)
poly_im = draw_convex_polygon(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_close(rect_im, poly_im)