import numpy as np
from PIL import Image

from image_prediction.estimator.preprocessor.utils import image_to_normalized_tensor, images_to_batch_tensor


def test_image_to_tensor(images):

    def inner(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

    assert all(map(inner, images))


def test_images_to_batch_tensor(images):
    tensor = images_to_batch_tensor(images)
    assert isinstance(tensor, np.ndarray)
    assert tensor.ndim == 4