import numpy as np import pytest from PIL import Image from funcy import lmap, compose from test.utils.image import image_to_bytes @pytest.fixture def data(data_type, pdf): if data_type == "pdf": return pdf elif data_type == "bytestring": return "content".encode("latin1") @pytest.fixture def data_items(item_type, n_items): if item_type == "string": return [bytes(f"content{i}", encoding="utf8") for i in range(n_items)] elif item_type == "image": return images(n_items) else: raise ValueError(f"Unknown item type {item_type}") @pytest.fixture(params=[0, 1, 2, 5]) def n_items(request): return request.param def array_to_image(array) -> Image.Image: return Image.fromarray(np.uint8(array * 255), mode="RGB") def input_batch(n_items): return np.random.random_sample(size=(n_items, 3, 30, 30)) def images(n_items): return lmap(compose(image_to_bytes, array_to_image), input_batch(n_items)) @pytest.fixture def metadata(): return {"idx": [1, 100, 101], "path": "asd/asd"}