import gzip import io import numpy as np import pytest from PIL import Image from funcy import first from cv_analysis.pyinfra_compat import get_analysis_fn from cv_analysis.utils.preprocessing import open_img_from_bytes from cv_analysis.utils.structures import Rectangle from incl.pyinfra.pyinfra.server.packing import bytes_to_string, string_to_bytes @pytest.fixture def random_image(): return np.random.rand(100, 100, 3) * 255 @pytest.fixture def random_image_as_bytes_and_compressed(random_image): image = Image.fromarray(random_image.astype("uint8")).convert("RGBA") img_byte_arr = io.BytesIO() image.save(img_byte_arr, format="PNG") return gzip.compress(img_byte_arr.getvalue()) @pytest.fixture def random_image_metadata_package(random_image_as_bytes_and_compressed): data = bytes_to_string(random_image_as_bytes_and_compressed) return [{"data": data, "metadata": {"key": "value", "key2": "value2"}}] @pytest.fixture def expected_analyse_metadata(operation, random_image_metadata_package): metadata = first(random_image_metadata_package) image = open_img_from_bytes(gzip.decompress(string_to_bytes(metadata["data"]))) wrapped_metadata = { **metadata["metadata"], "pageWidth": image.shape[1], "pageHeight": image.shape[0], } if operation == "mock": return { **wrapped_metadata, "cells": [{"x": 0, "y": 0, "width": image.shape[1], "height": image.shape[0]}], } if operation == "table_parsing": return {} @pytest.fixture def analyse_fn(operation): if operation == "mock": def analyse_mock(image: np.ndarray): return [Rectangle.from_xywh((0, 0, image.shape[1], image.shape[0]))] return analyse_mock return get_analysis_fn(operation)