cv-analysis-service/test/fixtures/pyinfra_compat.py

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)