wehub-resource-sync
3a2c66702c
Tests on CPU (scheduled) / check-skip (push) Has been cancelled
Tests on CPU (scheduled) / pre-tests (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-ubuntu (float32) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-ubuntu (float64) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.11, float32, 2.5.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.11, float32, 2.9.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.11, float64, 2.5.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.11, float64, 2.9.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.12, float32, 2.5.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.12, float32, 2.9.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.12, float64, 2.5.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.12, float64, 2.9.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.13, float32, 2.9.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-windows (3.13, float64, 2.9.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-mac (3.11, float32, 2.5.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-mac (3.11, float32, 2.9.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-mac (3.12, float32, 2.5.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-mac (3.12, float32, 2.9.1) (push) Has been cancelled
Tests on CPU (scheduled) / tests-cpu-mac (3.13, float32, 2.9.1) (push) Has been cancelled
Tests on CPU (scheduled) / coverage (push) Has been cancelled
Tests on CPU (scheduled) / typing (push) Has been cancelled
Tests on CPU (scheduled) / tutorials (push) Has been cancelled
Tests on CPU (scheduled) / docs (push) Has been cancelled
Lint / TOML Format (push) Has been cancelled
551 lines
26 KiB
Python
551 lines
26 KiB
Python
# LICENSE HEADER MANAGED BY add-license-header
|
|
#
|
|
# Copyright 2018 Kornia Team
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import math
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from kornia.geometry import create_meshgrid
|
|
from kornia.image import draw_convex_polygon, draw_rectangle
|
|
from kornia.image.draw import draw_line, draw_point2d
|
|
|
|
from testing.base import BaseTester
|
|
|
|
|
|
class TestDrawPoint(BaseTester):
|
|
"""Test drawing individual pixels."""
|
|
|
|
def test_draw_point2d_rgb(self, dtype, device):
|
|
"""Test plotting multiple [x, y] points."""
|
|
points = torch.tensor([(1, 3), (2, 4)], device=device)
|
|
color = torch.tensor([5, 10, 15], dtype=dtype, device=device)
|
|
img = torch.zeros(3, 8, 8, dtype=dtype, device=device)
|
|
img = draw_point2d(img, points, color)
|
|
for x, y in points:
|
|
self.assert_close(img[:, y, x], color.to(img.dtype))
|
|
|
|
def test_draw_point2d_grayscale_third_order(self, dtype, device):
|
|
"""Test plotting multiple [x, y] points on a (1, m, n) image."""
|
|
points = torch.tensor([(1, 3), (2, 4)], device=device)
|
|
color = torch.tensor([100], dtype=dtype, device=device)
|
|
img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
|
|
img = draw_point2d(img, points, color)
|
|
for x, y in points:
|
|
self.assert_close(img[:, y, x], color.to(img.dtype))
|
|
|
|
def test_draw_point2d_grayscale_second_order(self, dtype, device):
|
|
"""Test plotting multiple [x, y] points on a (m, n) image."""
|
|
points = torch.tensor([(1, 3), (2, 4)], device=device)
|
|
color = torch.tensor([100], dtype=dtype, device=device)
|
|
img = torch.zeros(8, 8, dtype=dtype, device=device)
|
|
img = draw_point2d(img, points, color)
|
|
for x, y in points:
|
|
self.assert_close(torch.unsqueeze(img[y, x], dim=0), color.to(img.dtype))
|
|
|
|
def test_draw_point2d_with_mismatched_dims(self, dtype, device):
|
|
"""Test that we raise if the len of the color tensor != the # of image channels."""
|
|
points = torch.tensor([(1, 3), (2, 4)], device=device)
|
|
color = torch.tensor([100], dtype=dtype, device=device)
|
|
img = torch.zeros(3, 8, 8, dtype=dtype, device=device)
|
|
with pytest.raises(Exception):
|
|
draw_point2d(img, points, color)
|
|
|
|
def test_draw_point2d_with_mismatched_dtype(self, device):
|
|
"""Test that the color is correctly cast to the image dtype when drawing points."""
|
|
points = torch.tensor([(1, 3), (2, 4)], device=device)
|
|
color = torch.tensor([5, 10, 15], dtype=torch.float32, device=device)
|
|
img = torch.zeros(3, 8, 8, dtype=torch.uint8, device=device)
|
|
img = draw_point2d(img, points, color)
|
|
assert img.dtype is torch.uint8
|
|
for x, y in points:
|
|
self.assert_close(img[:, y, x], color.to(torch.uint8))
|
|
|
|
def test_draw_point2d_with_singleton_color_dims(self, dtype, device):
|
|
"""Ensure that plotting behavior is consistent if we have a singleton dim for the color."""
|
|
points = torch.tensor([(1, 3), (2, 4)], device=device)
|
|
# Plot given a color tensor of shape [3]
|
|
color_vec = torch.tensor([5, 10, 15], dtype=torch.float32, device=device)
|
|
vec_img = torch.zeros(3, 8, 8, dtype=torch.uint8, device=device)
|
|
drawn_vec_img = draw_point2d(vec_img, points, color_vec)
|
|
# Plot given a color tensor of shape [3, 1]
|
|
color_mat = torch.unsqueeze(color_vec, dim=1)
|
|
mat_img = vec_img.clone()
|
|
drawn_mat_img = draw_point2d(mat_img, points, color_mat)
|
|
# Ensure that we get the same underlying image back
|
|
self.assert_close(drawn_vec_img, drawn_mat_img)
|
|
|
|
|
|
class TestDrawLine(BaseTester):
|
|
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, 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,
|
|
)
|
|
self.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],
|
|
[0.0, 0.0, 0.0, 0.0, 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,
|
|
)
|
|
self.assert_close(img, img_mask)
|
|
|
|
def test_draw_line_with_big_coordinates(self, dtype, device):
|
|
"""Test drawing a line with big coordinates."""
|
|
img = torch.zeros(1, 500, 500, dtype=dtype, device=device)
|
|
img = draw_line(img, torch.tensor([200, 200]), torch.tensor([400, 200]), torch.tensor([255]))
|
|
img_mask = torch.zeros(1, 500, 500, dtype=dtype, device=device)
|
|
img_mask[:, 200, 200:401] = 255
|
|
self.assert_close(img, img_mask)
|
|
|
|
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, 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, 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,
|
|
)
|
|
self.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.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, 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, 0.0, 0.0, 0.0, 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,
|
|
)
|
|
self.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.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, 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,
|
|
)
|
|
self.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.0, 0.0, 0.0, 0.0, 0.0, 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, 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, 0.0, 255.0, 0.0, 0.0, 0.0, 0.0],
|
|
]
|
|
],
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
self.assert_close(img, img_mask)
|
|
|
|
def test_draw_lines_batched(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], [0, 7], [0, 2]]), torch.tensor([[6, 0], [0, 0], [7, 7]]), torch.tensor([255])
|
|
)
|
|
img_mask = torch.tensor(
|
|
[
|
|
[
|
|
[255.0, 0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0],
|
|
[255.0, 0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0],
|
|
[255.0, 0.0, 0.0, 0.0, 0.0, 255.0, 0.0, 0.0],
|
|
[255.0, 255.0, 0.0, 0.0, 255.0, 0.0, 0.0, 0.0],
|
|
[255.0, 0.0, 255.0, 255.0, 0.0, 0.0, 0.0, 0.0],
|
|
[255.0, 0.0, 255.0, 255.0, 255.0, 0.0, 0.0, 0.0],
|
|
[255.0, 255.0, 0.0, 0.0, 0.0, 255.0, 0.0, 0.0],
|
|
[255.0, 0.0, 0.0, 0.0, 0.0, 0.0, 255.0, 255.0],
|
|
]
|
|
],
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
self.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])]
|
|
)
|
|
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)
|
|
with pytest.raises(ValueError) as excinfo:
|
|
draw_line(img, p1, torch.tensor([0, 0]), torch.tensor([255]))
|
|
|
|
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])]
|
|
)
|
|
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)
|
|
with pytest.raises(ValueError) as excinfo:
|
|
draw_line(img, torch.tensor([0, 0]), p2, torch.tensor([255]))
|
|
|
|
assert "p2 is out of bounds." == str(excinfo.value)
|
|
|
|
@pytest.mark.parametrize(
|
|
"p1,p2", [(torch.tensor([0, 0]), torch.tensor([0, 0])), (torch.tensor([0, 0]), torch.tensor([7, 7]))]
|
|
)
|
|
def test_p1_p2_bounds(self, p1, p2, dtype, device):
|
|
"""Tests to verify that points on the bounds do not trigger errors."""
|
|
img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
|
|
img = draw_line(img, p1, p2, torch.tensor([255]))
|
|
assert img.shape == (1, 8, 8)
|
|
|
|
@pytest.mark.parametrize("img_size", [(200, 100), (32, 3, 20, 20)])
|
|
def test_image_size(self, img_size, dtype, device):
|
|
img = torch.zeros(*img_size, dtype=dtype, device=device)
|
|
with pytest.raises(ValueError) as excinfo:
|
|
draw_line(img, torch.tensor([0, 0]), torch.tensor([1, 1]), torch.tensor([255]))
|
|
|
|
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]))])
|
|
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:
|
|
draw_line(img, torch.tensor([0, 0]), torch.tensor([1, 1]), color)
|
|
|
|
assert "color must have the same number of channels as the image." == str(excinfo.value)
|
|
|
|
@pytest.mark.parametrize("p1,p2", [(torch.rand([10, 2]), torch.rand([20, 2])), (torch.rand([2]), torch.rand([3]))])
|
|
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:
|
|
draw_line(img, p1, p2, torch.tensor([255]))
|
|
|
|
assert "Input points must be 2D points with shape (2, ) or (B, 2) and must have the same batch sizes." == str(
|
|
excinfo.value
|
|
)
|
|
|
|
|
|
class TestDrawRectangle(BaseTester):
|
|
@pytest.mark.parametrize("batch", (4, 17))
|
|
@pytest.mark.parametrize("color", (torch.Tensor([1.0]), torch.Tensor([0.5])))
|
|
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)
|
|
|
|
@pytest.mark.parametrize("batch", (8, 11))
|
|
@pytest.mark.parametrize("fill", (True, False))
|
|
@pytest.mark.parametrize("height", (12, 106, 298))
|
|
@pytest.mark.parametrize("width", (7, 123, 537))
|
|
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)
|
|
|
|
@pytest.mark.parametrize("batch", (4, 6))
|
|
@pytest.mark.parametrize("N", (5, 12))
|
|
@pytest.mark.parametrize("fill", (True, False))
|
|
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, device=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
|
|
)
|
|
|
|
@pytest.mark.parametrize("color", (torch.tensor([0.5, 0.3, 0.15]), torch.tensor([0.23, 0.33, 0.8])))
|
|
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
|
|
)
|
|
|
|
@pytest.mark.parametrize("color", (torch.tensor([0.34, 0.63, 0.16]), torch.tensor([0.29, 0.13, 0.48])))
|
|
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
|
|
)
|
|
|
|
|
|
class TestFillConvexPolygon(BaseTester):
|
|
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, device=device, dtype=dtype)[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_polygon(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, 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_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_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)
|
|
self.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)
|
|
self.assert_close(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]], 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)
|
|
self.assert_close(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], 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, fill=True)
|
|
self.assert_close(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_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)
|
|
self.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)
|
|
self.assert_close(rect_im, poly_im)
|