def test_n_rectangles(self, device, batch, N, fill): points_list = [] h, w = 20, 20 for b in range(batch): points_list.append([]) for n in range(N): points_list[b].append([]) points_list[b][n].append(int(torch.randint(0, w - 1, (1,)))) points_list[b][n].append(int(torch.randint(0, h - 1, (1,)))) points_list[b][n].append(int(torch.randint(points_list[b][n][-2] + 1, w, (1,)))) points_list[b][n].append(int(torch.randint(points_list[b][n][-2] + 1, h, (1,)))) points = torch.tensor(points_list).to(device) random_background = torch.rand(batch, 3, h, w, device=device) random_w_rectangle = random_background.clone() draw_rectangle(random_w_rectangle, points, color=torch.tensor([1.0, 1.0, 1.0]), fill=fill) for b in range(batch): for n in range(N): if fill: assert ( 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, ].sum() == (points_list[b][n][3] - points_list[b][n][1] + 1) * (points_list[b][n][2] - points_list[b][n][0] + 1) * 3 ) else: assert ( random_w_rectangle[ 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] ].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 ].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 ].sum() == (points_list[b][n][2] - points_list[b][n][0] + 1) * 3 )
def test_smoke(self, device, batch, color): black_image = torch.zeros(batch, 1, 3, 3, device=device) # 1 channel 3x3 black_image points = torch.tensor([1.0, 1.0, 1.0, 1.0]).to(device).expand(batch, 1, 4) # single pixel rectangle draw_rectangle(black_image, points, color=color) target = torch.zeros(batch, 1, 3, 3, device=device) target[:, :, 1, 1] = color assert torch.all(black_image == target)
def test_color_foreground(self, device, color): image = torch.zeros(1, 3, 50, 40, device=device) image_w_rectangle = image.clone() p1 = (10, 4) p2 = (11, 40) points = torch.tensor([[[p1[1], p1[0], p2[1], p2[0]]]], device=device) draw_rectangle(image_w_rectangle, points, color=color) # corners are double counted, no plus 1 for y since p2[1] of 40 already lies outside of the image assert torch.abs((image_w_rectangle - image).sum() - (color[0]) * (2 * (p2[0] - p1[0] + 1) + 2 * (p2[1] - p1[1]) - 4) - (color[1]) * (2 * (p2[0] - p1[0] + 1) + 2 * (p2[1] - p1[1]) - 4) - (color[2]) * (2 * (p2[0] - p1[0] + 1) + 2 * (p2[1] - p1[1]) - 4)) \ <= 0.0001
def test_color_background(self, device, color): image = torch.zeros(1, 3, 40, 40, device=device) image[:, 0, :, :] = color[0] image[:, 1, :, :] = color[1] image[:, 2, :, :] = color[2] image_w_rectangle = image.clone() p1 = (1, 5) p2 = (30, 39) points = torch.tensor([[[p1[1], p1[0], p2[1], p2[0]]]], device=device) draw_rectangle(image_w_rectangle, points, color=torch.tensor([1.0])) assert torch.abs((image_w_rectangle - image).sum() - (1 - color[0]) * (2 * (p2[0] - p1[0] + 1) + 2 * (p2[1] - p1[1] + 1) - 4) - (1 - color[1]) * (2 * (p2[0] - p1[0] + 1) + 2 * (p2[1] - p1[1] + 1) - 4) - (1 - color[2]) * (2 * (p2[0] - p1[0] + 1) + 2 * (p2[1] - p1[1] + 1) - 4)) \ <= 0.0001
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_fill_and_edges(self, device, batch, fill, height, width): black_image = torch.zeros(batch, 3, height, width, device=device) # we should pass height - 1 and width - 1 but rectangle should clip correctly points = torch.tensor([0, 0, width, height]).to(device).expand(batch, 1, 4) image_w_rectangle = draw_rectangle(black_image, points, color=torch.tensor([1.0]), fill=fill) assert image_w_rectangle is black_image if fill: assert image_w_rectangle.sum() == batch * 3 * height * width else: # corners are double counted assert image_w_rectangle.sum() == batch * 3 * (2 * height + 2 * width - 4)
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)