def op_script(data: torch.Tensor) -> torch.Tensor:
return kornia.rgb_to_bgr(data)
data = torch.tensor([[[1., 1.], [1., 1.]], [[2., 2.], [2., 2.]],
[[3., 3.], [3., 3.]]]) # 3x2x2
actual = op_script(data)
expected = kornia.rgb_to_bgr(data)
assert_allclose(actual, expected)
def test_bgr(self, device):
a_val: float = 1.
x_rgb = torch.rand(3, 4, 4).to(device)
x_bgr = kornia.rgb_to_bgr(x_rgb)
x_rgba = kornia.rgb_to_rgba(x_rgb, a_val)
x_rgba_new = kornia.bgr_to_rgba(x_bgr, a_val)
assert_allclose(x_rgba, x_rgba_new)
def test_back_and_forth(self, device):
data_bgr = torch.rand(1, 3, 3, 2).to(device)
data_rgb = kornia.bgr_to_rgb(data_bgr)
data_bgr_new = kornia.rgb_to_bgr(data_rgb)
assert_allclose(data_bgr, data_bgr_new)
def op_script(data: torch.Tensor) -> torch.Tensor:
return kornia.rgb_to_bgr(data)