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

import gzip
import io

import numpy as np
import pytest
from PIL import Image
from funcy import first

from cv_analysis.utils.structures import Rectangle
from incl.pyinfra.pyinfra.server.packing import bytes_to_string


@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": {
                "page_info": {"width": 1000, "height": 2000, "rotation": 90},
                "image_info": {"dpi": 200},
            },
        }
    ]


@pytest.fixture
def expected_analyse_metadata(operation, random_image_metadata_package, image_size):
    metadata = first(random_image_metadata_package)
    metadata = metadata["metadata"]

    if image_size == (200, 200):
        result_metadata = {"cells": [{"height": 72.0, "width": 71.99999999999999, "x": 0.0, "y": 1928.0}]}
    elif image_size == (500, 500):
        result_metadata = {"cells": [{"height": 180.0, "width": 179.99999999999997, "x": 0.0, "y": 1820.0}]}
    elif image_size == (800, 800):
        result_metadata = {"cells": [{"height": 288.0, "width": 287.99999999999994, "x": 0.0, "y": 1712.0}]}
    else:
        result_metadata = {}

    if operation == "mock":
        return {**metadata, **result_metadata}


@pytest.fixture
def analysis_fn_mock(operation):
    def analyse_mock(image: np.ndarray):
        return [Rectangle.from_xywh((0, 0, image.shape[1], image.shape[0]))]

    return analyse_mock