made input size adjustable via test fixture

This commit is contained in:
Matthias Bisping 2022-03-28 19:22:31 +02:00
parent f18e183ab0
commit 3c6dfed508
2 changed files with 7 additions and 6 deletions

View File

@ -100,9 +100,9 @@ def array_to_image(array):
return Image.fromarray(np.uint8(array * 255), mode="RGB")
@pytest.fixture
def input_size(depth=3, width=10, height=15):
return width, height, depth
@pytest.fixture(params=[{'depth': 3}, {'width': 10}, {'height': 15}])
def input_size(request):
return request.param.get("width", 10), request.param.get("height", 15), request.param.get("depth", 3)
@pytest.fixture
@ -135,7 +135,7 @@ def metadata(images):
def current_page_idx():
nonlocal page_idx
page_idx += random.randint(0, 2)
page_idx += random.randint(0, 3)
return min(page_idx, len(images) - 1)
def build_image_metadata(image):

View File

@ -13,8 +13,9 @@ def test_image_extractor_mock(image_extractor, images):
@pytest.mark.parametrize("extractor_type", ["parsable_pdf"])
@pytest.mark.parametrize("batch_size", [0, 1, 2, 4, 8])
def test_parsable_pdf_image_extractor(image_extractor, pdf, images, metadata):
@pytest.mark.parametrize("batch_size", [128])
@pytest.mark.parametrize("input_size", [{"depth": 3, "width": 100, "height": 90}], indirect=["input_size"])
def test_parsable_pdf_image_extractor(image_extractor, pdf, images, metadata, input_size):
images_extracted, metadata_extracted = map(list, extract_images_from_pdf(pdf, image_extractor))
assert np.allclose(images_to_batch_tensor(images_extracted), images_to_batch_tensor(images))
assert list(metadata_extracted) == metadata