File size: 3,324 Bytes
36c95ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
from typing import List
import numpy as np
import pytest
import torch
import kornia
from kornia.testing import assert_close
@pytest.mark.parametrize(
"input_dtype, expected_dtype", [(np.uint8, torch.uint8), (np.float32, torch.float32), (np.float64, torch.float64)]
)
def test_image_to_tensor_keep_dtype(input_dtype, expected_dtype):
image = np.ones((1, 3, 4, 5), dtype=input_dtype)
tensor = kornia.image_to_tensor(image)
assert tensor.dtype == expected_dtype
@pytest.mark.parametrize("num_of_images, image_shape", [(2, (4, 3, 1)), (0, (1, 2, 3)), (5, (2, 3, 2, 5))])
def test_list_of_images_to_tensor(num_of_images, image_shape):
images: List[np.array] = []
if num_of_images == 0:
with pytest.raises(ValueError):
kornia.utils.image_list_to_tensor([])
return
for _ in range(num_of_images):
images.append(np.ones(shape=image_shape))
if len(image_shape) != 3:
with pytest.raises(ValueError):
kornia.utils.image_list_to_tensor(images)
return
tensor = kornia.utils.image_list_to_tensor(images)
assert tensor.shape == (num_of_images, image_shape[-1], image_shape[-3], image_shape[-2])
@pytest.mark.parametrize(
"input_shape, expected",
[
((4, 4), (4, 4)),
((1, 4, 4), (4, 4)),
((1, 1, 4, 4), (4, 4)),
((3, 4, 4), (4, 4, 3)),
((2, 3, 4, 4), (2, 4, 4, 3)),
((1, 3, 4, 4), (4, 4, 3)),
],
)
def test_tensor_to_image(device, input_shape, expected):
tensor = torch.ones(input_shape).to(device)
image = kornia.utils.tensor_to_image(tensor)
assert image.shape == expected
assert isinstance(image, np.ndarray)
@pytest.mark.parametrize(
"input_shape, expected",
[
((4, 4), (4, 4)),
((1, 4, 4), (4, 4)),
((1, 1, 4, 4), (1, 4, 4)),
((3, 4, 4), (4, 4, 3)),
((2, 3, 4, 4), (2, 4, 4, 3)),
((1, 3, 4, 4), (1, 4, 4, 3)),
],
)
def test_tensor_to_image_keepdim(device, input_shape, expected):
tensor = torch.ones(input_shape).to(device)
image = kornia.utils.tensor_to_image(tensor, keepdim=True)
assert image.shape == expected
assert isinstance(image, np.ndarray)
@pytest.mark.parametrize(
"input_shape, expected",
[
((4, 4), (1, 1, 4, 4)),
((1, 4, 4), (1, 4, 1, 4)),
((2, 3, 4), (1, 4, 2, 3)),
((4, 4, 3), (1, 3, 4, 4)),
((2, 4, 4, 3), (2, 3, 4, 4)),
((1, 4, 4, 3), (1, 3, 4, 4)),
],
)
def test_image_to_tensor(input_shape, expected):
image = np.ones(input_shape)
tensor = kornia.utils.image_to_tensor(image, keepdim=False)
assert tensor.shape == expected
assert isinstance(tensor, torch.Tensor)
to_tensor = kornia.utils.ImageToTensor(keepdim=False)
assert_close(tensor, to_tensor(image))
@pytest.mark.parametrize(
"input_shape, expected",
[
((4, 4), (1, 4, 4)),
((1, 4, 4), (4, 1, 4)),
((2, 3, 4), (4, 2, 3)),
((4, 4, 3), (3, 4, 4)),
((2, 4, 4, 3), (2, 3, 4, 4)),
((1, 4, 4, 3), (1, 3, 4, 4)),
],
)
def test_image_to_tensor_keepdim(input_shape, expected):
image = np.ones(input_shape)
tensor = kornia.utils.image_to_tensor(image, keepdim=True)
assert tensor.shape == expected
assert isinstance(tensor, torch.Tensor)
|