38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from image_prediction.estimator.preprocessor.preprocessors.basic import BasicPreprocessor
|
|
from image_prediction.estimator.preprocessor.utils import image_to_normalized_tensor, images_to_batch_tensor
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [0, 1, 2, 16, 32, 64])
|
|
def image_conversion_is_correct(image):
|
|
tensor = image_to_normalized_tensor(image)
|
|
image_re = Image.fromarray(np.uint8(tensor * 255), mode="RGB")
|
|
return image == image_re and tensor.ndim == 3
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [0, 1, 2, 16, 32, 64])
|
|
def images_conversion_is_correct(images, tensor):
|
|
if not (images or tensor.size > 0):
|
|
return True
|
|
return all([isinstance(tensor, np.ndarray), tensor.ndim == 4, tensor.shape[0] == len(images)])
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [0, 1, 2, 16, 32, 64])
|
|
def test_image_to_tensor(images):
|
|
assert all(map(image_conversion_is_correct, images))
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [0, 1, 2, 16, 32, 64])
|
|
def test_images_to_batch_tensor(images):
|
|
tensor = images_to_batch_tensor(images)
|
|
assert images_conversion_is_correct(images, tensor)
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [0, 1, 2, 4, 6], scope="session")
|
|
def test_basic_preprocessor(images):
|
|
tensor = BasicPreprocessor()(images)
|
|
assert images_conversion_is_correct(images, tensor)
|