ludwig-ai--ludwig
wehub-resource-sync
593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
390 行
13 KiB
Python
390 行
13 KiB
Python
# Copyright (c) 2023 Predibase, Inc., 2019 Uber Technologies, Inc.
|
|
#
|
|
# 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.
|
|
# ==============================================================================
|
|
from collections.abc import Callable
|
|
|
|
import pytest
|
|
import torch
|
|
import torchvision.transforms.functional as F
|
|
|
|
from ludwig.utils.image_utils import (
|
|
crop,
|
|
crop_or_pad,
|
|
get_class_mask_from_image,
|
|
get_image_from_class_mask,
|
|
get_unique_channels,
|
|
grayscale,
|
|
is_image_score,
|
|
num_channels_in_image,
|
|
pad,
|
|
read_image_as_tif,
|
|
resize_image,
|
|
ResizeChannels,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("pad_fn", [pad, torch.jit.script(pad)])
|
|
@pytest.mark.parametrize(
|
|
"img,size,padded_img",
|
|
[
|
|
(
|
|
torch.arange(12, dtype=torch.int).reshape(3, 2, 2),
|
|
4,
|
|
torch.Tensor(
|
|
[
|
|
0,
|
|
0,
|
|
1,
|
|
1,
|
|
0,
|
|
0,
|
|
1,
|
|
1,
|
|
2,
|
|
2,
|
|
3,
|
|
3,
|
|
2,
|
|
2,
|
|
3,
|
|
3,
|
|
4,
|
|
4,
|
|
5,
|
|
5,
|
|
4,
|
|
4,
|
|
5,
|
|
5,
|
|
6,
|
|
6,
|
|
7,
|
|
7,
|
|
6,
|
|
6,
|
|
7,
|
|
7,
|
|
8,
|
|
8,
|
|
9,
|
|
9,
|
|
8,
|
|
8,
|
|
9,
|
|
9,
|
|
10,
|
|
10,
|
|
11,
|
|
11,
|
|
10,
|
|
10,
|
|
11,
|
|
11,
|
|
]
|
|
)
|
|
.type(torch.int)
|
|
.reshape(3, 4, 4),
|
|
)
|
|
],
|
|
)
|
|
def test_pad(pad_fn: Callable, img: torch.Tensor, size: int, padded_img: torch.Tensor):
|
|
output_img = pad_fn(img, size)
|
|
assert torch.equal(output_img, padded_img)
|
|
|
|
|
|
@pytest.mark.parametrize("crop_fn", [crop, torch.jit.script(crop)])
|
|
@pytest.mark.parametrize(
|
|
"img,size,cropped_img",
|
|
[
|
|
(
|
|
torch.arange(27, dtype=torch.int).reshape(3, 3, 3),
|
|
2,
|
|
torch.Tensor([0, 1, 3, 4, 9, 10, 12, 13, 18, 19, 21, 22]).type(torch.int).reshape(3, 2, 2),
|
|
)
|
|
],
|
|
)
|
|
def test_crop(crop_fn: Callable, img: torch.Tensor, size: int, cropped_img: torch.Tensor):
|
|
output_img = crop_fn(img, size)
|
|
assert torch.equal(output_img, cropped_img)
|
|
|
|
|
|
@pytest.mark.parametrize("crop_or_pad_fn", [crop_or_pad, torch.jit.script(crop_or_pad)])
|
|
@pytest.mark.parametrize(
|
|
"img,new_size,expected_img",
|
|
[
|
|
(
|
|
torch.arange(12, dtype=torch.int).reshape(3, 2, 2),
|
|
4,
|
|
torch.Tensor(
|
|
[
|
|
0,
|
|
0,
|
|
1,
|
|
1,
|
|
0,
|
|
0,
|
|
1,
|
|
1,
|
|
2,
|
|
2,
|
|
3,
|
|
3,
|
|
2,
|
|
2,
|
|
3,
|
|
3,
|
|
4,
|
|
4,
|
|
5,
|
|
5,
|
|
4,
|
|
4,
|
|
5,
|
|
5,
|
|
6,
|
|
6,
|
|
7,
|
|
7,
|
|
6,
|
|
6,
|
|
7,
|
|
7,
|
|
8,
|
|
8,
|
|
9,
|
|
9,
|
|
8,
|
|
8,
|
|
9,
|
|
9,
|
|
10,
|
|
10,
|
|
11,
|
|
11,
|
|
10,
|
|
10,
|
|
11,
|
|
11,
|
|
]
|
|
)
|
|
.type(torch.int)
|
|
.reshape(3, 4, 4),
|
|
),
|
|
(
|
|
torch.arange(27, dtype=torch.int).reshape(3, 3, 3),
|
|
2,
|
|
torch.Tensor([0, 1, 3, 4, 9, 10, 12, 13, 18, 19, 21, 22]).type(torch.int).reshape(3, 2, 2),
|
|
),
|
|
],
|
|
)
|
|
def test_crop_or_pad(crop_or_pad_fn: Callable, img: torch.Tensor, new_size: int, expected_img: torch.Tensor):
|
|
output_image = crop_or_pad_fn(img, new_size)
|
|
assert torch.equal(output_image, expected_img)
|
|
|
|
|
|
@pytest.mark.parametrize("resize_image_fn", [resize_image, torch.jit.script(resize_image)])
|
|
@pytest.mark.parametrize(
|
|
"img,new_size,resize_method",
|
|
[
|
|
(
|
|
torch.arange(27, dtype=torch.int).reshape(3, 3, 3),
|
|
2,
|
|
"crop_or_pad",
|
|
),
|
|
(
|
|
torch.arange(27, dtype=torch.int).reshape(3, 3, 3),
|
|
2,
|
|
"interpolate",
|
|
),
|
|
],
|
|
)
|
|
def test_resize_image(resize_image_fn: Callable, img: torch.Tensor, new_size: int, resize_method: str):
|
|
# Get the expected output from the underlying function
|
|
if resize_method == "crop_or_pad":
|
|
expected_img = crop_or_pad(img, new_size)
|
|
else:
|
|
expected_img = F.resize(img, new_size)
|
|
|
|
output_img = resize_image_fn(img, new_size, resize_method)
|
|
|
|
# Test that resize_image is equivalent to the underlying function output
|
|
assert torch.equal(output_img, expected_img)
|
|
|
|
|
|
@pytest.mark.parametrize("grayscale_fn", [grayscale, torch.jit.script(grayscale)])
|
|
@pytest.mark.parametrize(
|
|
"input_img,grayscale_img",
|
|
[(torch.arange(12).reshape(3, 2, 2).type(torch.int), torch.Tensor([[[3, 4], [5, 6]]]).type(torch.int))],
|
|
)
|
|
def test_grayscale(grayscale_fn: Callable, input_img: torch.Tensor, grayscale_img: torch.Tensor):
|
|
output_img = grayscale_fn(input_img)
|
|
assert torch.equal(output_img, grayscale_img)
|
|
|
|
|
|
def test_num_channels_in_image():
|
|
image_2d = torch.randint(0, 1, (10, 10))
|
|
image_3d = torch.randint(0, 1, (3, 10, 10))
|
|
assert num_channels_in_image(image_2d) == 1
|
|
assert num_channels_in_image(image_3d) == 3
|
|
|
|
with pytest.raises(ValueError):
|
|
num_channels_in_image(torch.rand(5))
|
|
num_channels_in_image(None)
|
|
|
|
|
|
@pytest.mark.parametrize("image_shape", [(1, 10, 10), (3, 10, 10), (5, 10, 10)])
|
|
@pytest.mark.parametrize("num_channels_expected", [1, 2, 3, 4])
|
|
def test_ResizeChannels_module(image_shape, num_channels_expected):
|
|
image = torch.randint(0, 1, image_shape)
|
|
fn = ResizeChannels(num_channels_expected)
|
|
assert fn(image).shape == tuple([num_channels_expected] + list(image_shape[1:]))
|
|
|
|
|
|
@pytest.mark.parametrize("image_shape", [(2, 1, 10, 10), (2, 3, 10, 10), (2, 5, 10, 10)])
|
|
@pytest.mark.parametrize("num_channels_expected", [1, 2, 3, 4])
|
|
def test_ResizeChannels_module_with_batch_dim(image_shape, num_channels_expected):
|
|
image = torch.randint(0, 1, image_shape)
|
|
fn = ResizeChannels(num_channels_expected)
|
|
assert fn(image).shape == tuple([image_shape[0], num_channels_expected] + list(image_shape[2:]))
|
|
|
|
|
|
def test_read_image_as_tif():
|
|
img_bytes = b"II*\x00\x0c\x00\x00\x00\x05 \x8c\xe5\x10\x00\x00\x01\x03\x00\x01\x00\x00\x00\x02\x00\x00\x00\x01\x01\x03\x00\x01\x00\x00\x00\x02\x00\x00\x00\x02\x01\x03\x00\x01\x00\x00\x00\x08\x00\x00\x00\x03\x01\x03\x00\x01\x00\x00\x00\x01\x00\x00\x00\x06\x01\x03\x00\x01\x00\x00\x00\x01\x00\x00\x00\x11\x01\x04\x00\x01\x00\x00\x00\x08\x00\x00\x00\x12\x01\x03\x00\x01\x00\x00\x00\x01\x00\x00\x00\x15\x01\x03\x00\x01\x00\x00\x00\x01\x00\x00\x00\x16\x01\x03\x00\x01\x00\x00\x00\x80\x00\x00\x00\x17\x01\x04\x00\x01\x00\x00\x00\x04\x00\x00\x00\x1a\x01\x05\x00\x01\x00\x00\x00\xd2\x00\x00\x00\x1b\x01\x05\x00\x01\x00\x00\x00\xda\x00\x00\x00\x1c\x01\x03\x00\x01\x00\x00\x00\x01\x00\x00\x00\x1d\x01\x02\x00\x07\x00\x00\x00\xe2\x00\x00\x00(\x01\x03\x00\x01\x00\x00\x00\x02\x00\x00\x00S\x01\x03\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00H\x00\x00\x00\x01\x00\x00\x00H\x00\x00\x00\x01\x00\x00\x004.tiff\x00"
|
|
tensor = read_image_as_tif(img_bytes)
|
|
assert tensor is not None
|
|
assert tensor.equal(torch.tensor([[[5, 32], [140, 229]]], dtype=torch.uint8))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"extension, score",
|
|
[
|
|
("data.png", 1),
|
|
("/home/peter/data.jpg", 1),
|
|
("./data/file.jpeg", 1),
|
|
("new.tiff", 1),
|
|
("b.tif", 1),
|
|
(".bmp", 1),
|
|
("a.gif", 1),
|
|
("b.tif", 1),
|
|
("audio.wav", 0),
|
|
(".png/video.mp4", 0),
|
|
],
|
|
)
|
|
def test_is_image_score(extension: str, score: int):
|
|
assert is_image_score(extension) == score
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"img_list,num_channels,num_classes,expected_class_map",
|
|
[
|
|
(
|
|
[
|
|
torch.Tensor([0, 0, 8, 8, 120, 120, 180, 180, 230, 230, 255, 255]).type(torch.uint8).reshape(3, 2, 2),
|
|
torch.Tensor([1, 2, 3, 4, 131, 132, 133, 134, 241, 242, 243, 244]).type(torch.uint8).reshape(3, 2, 2),
|
|
],
|
|
3,
|
|
None,
|
|
torch.Tensor(
|
|
[[0, 120, 230], [8, 180, 255], [1, 131, 241], [2, 132, 242], [3, 133, 243], [4, 134, 244]]
|
|
).type(torch.uint8),
|
|
),
|
|
(
|
|
[
|
|
torch.Tensor([0, 255, 255, 0, 255, 255, 255, 0, 0, 0, 255, 255, 0, 255, 255])
|
|
.type(torch.uint8)
|
|
.reshape(1, 3, 5),
|
|
],
|
|
1,
|
|
None,
|
|
torch.Tensor([[0], [255]]).type(torch.uint8),
|
|
),
|
|
(
|
|
[
|
|
torch.Tensor([0, 31, 17, 185, 192, 173, 55, 76, 24, 128, 255, 238]).type(torch.uint8).reshape(3, 4),
|
|
],
|
|
1,
|
|
2,
|
|
torch.Tensor([[0], [255]]).type(torch.uint8),
|
|
),
|
|
],
|
|
)
|
|
def test_unique_channels(
|
|
img_list: list[torch.Tensor], num_channels: int, num_classes: int, expected_class_map: torch.Tensor
|
|
):
|
|
channel_class_map = get_unique_channels(img_list, num_channels, num_classes)
|
|
|
|
channel_class_map, _ = channel_class_map.sort(dim=0)
|
|
expected_class_map, _ = expected_class_map.sort(dim=0)
|
|
assert torch.equal(channel_class_map, expected_class_map)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"img,channel_class_map,expected_mask",
|
|
[
|
|
(
|
|
torch.Tensor([1, 2, 3, 4, 131, 132, 133, 134, 241, 242, 243, 244]).type(torch.uint8).reshape(3, 2, 2),
|
|
torch.Tensor(
|
|
[[0, 120, 230], [8, 180, 255], [1, 131, 241], [2, 132, 242], [3, 133, 243], [4, 134, 244]]
|
|
).type(torch.uint8),
|
|
torch.Tensor([2, 3, 4, 5]).type(torch.uint8).reshape(2, 2),
|
|
),
|
|
(
|
|
torch.Tensor([0, 255, 255, 0, 255, 255, 255, 0, 0, 0, 255, 255, 0, 255, 255])
|
|
.type(torch.uint8)
|
|
.reshape(1, 3, 5),
|
|
torch.Tensor([[0], [255]]).type(torch.uint8),
|
|
torch.Tensor([0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 1]).type(torch.uint8).reshape(3, 5),
|
|
),
|
|
(
|
|
torch.Tensor([0, 31, 17, 185, 192, 173, 55, 76, 24, 128, 255, 238]).type(torch.uint8).reshape(3, 4),
|
|
torch.Tensor([[0], [255]]).type(torch.uint8),
|
|
torch.Tensor([0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1]).type(torch.uint8).reshape(3, 4),
|
|
),
|
|
],
|
|
)
|
|
def test_class_mask_from_image(img: torch.Tensor, channel_class_map: torch.Tensor, expected_mask: torch.Tensor):
|
|
mask = get_class_mask_from_image(channel_class_map, img)
|
|
assert torch.equal(mask, expected_mask)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"mask,channel_class_map,expected_img",
|
|
[
|
|
(
|
|
torch.Tensor([0, 0, 1, 1]).type(torch.uint8).reshape(2, 2),
|
|
torch.Tensor(
|
|
[[0, 120, 230], [8, 180, 255], [1, 131, 241], [2, 132, 242], [3, 133, 243], [4, 134, 244]]
|
|
).type(torch.uint8),
|
|
torch.Tensor([0, 0, 8, 8, 120, 120, 180, 180, 230, 230, 255, 255]).type(torch.uint8).reshape(3, 2, 2),
|
|
),
|
|
(
|
|
torch.Tensor([2, 3, 4, 5]).type(torch.uint8).reshape(2, 2),
|
|
torch.Tensor(
|
|
[[0, 120, 230], [8, 180, 255], [1, 131, 241], [2, 132, 242], [3, 133, 243], [4, 134, 244]]
|
|
).type(torch.uint8),
|
|
torch.Tensor([1, 2, 3, 4, 131, 132, 133, 134, 241, 242, 243, 244]).type(torch.uint8).reshape(3, 2, 2),
|
|
),
|
|
(
|
|
torch.Tensor([0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 1]).type(torch.uint8).reshape(3, 5),
|
|
torch.Tensor([[0], [255]]).type(torch.uint8),
|
|
torch.Tensor([0, 255, 255, 0, 255, 255, 255, 0, 0, 0, 255, 255, 0, 255, 255])
|
|
.type(torch.uint8)
|
|
.reshape(1, 3, 5),
|
|
),
|
|
],
|
|
)
|
|
def test_image_from_class_mask(mask: torch.Tensor, channel_class_map: torch.Tensor, expected_img: torch.Tensor):
|
|
img = get_image_from_class_mask(channel_class_map, mask.numpy())
|
|
assert torch.equal(torch.from_numpy(img), expected_img)
|