22 lines
1.0 KiB
Python
22 lines
1.0 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from image_prediction.estimator.preprocessor.utils import images_to_batch_tensor
|
|
from image_prediction.extraction import extract_images_from_pdf
|
|
|
|
|
|
@pytest.mark.parametrize("extractor_type", ["mock"])
|
|
@pytest.mark.parametrize("batch_size", [1, 2, 4])
|
|
def test_image_extractor_mock(image_extractor, images):
|
|
images_extracted, metadata = map(list, zip(*image_extractor(images)))
|
|
assert images_extracted == images
|
|
|
|
|
|
@pytest.mark.parametrize("extractor_type", ["parsable_pdf", "default"])
|
|
@pytest.mark.parametrize("batch_size", [0, 1, 2, 64])
|
|
@pytest.mark.parametrize("input_size", [{"depth": 3, "width": 170, "height": 220}], 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
|