cv-analysis-service/test/fixtures/page_generation/page.py

import sys
from typing import Tuple, Iterable, List

import blend_modes
import numpy as np
import pytest
from PIL import Image, ImageEnhance
from PIL.Image import Transpose
from funcy import (
    juxt,
    compose,
    identity,
)
from loguru import logger

from cv_analysis.locations import TEST_PAGE_TEXTURES_DIR
from cv_analysis.utils.conversion import normalize_image_format_to_array, normalize_image_format_to_pil
from cv_analysis.utils.image_operations import blur, sharpen, overlay, superimpose
from cv_analysis.utils.rectangle import Rectangle
from synthesis.content_generator import ContentGenerator
from synthesis.partitioner.two_column import TwoColumnPagePartitioner
from synthesis.random import rnd
from synthesis.segment.table.table import paste_contents

logger.remove()
logger.add(sys.stderr, level="INFO")


@pytest.fixture(
    params=[
        # "rough_grain",
        # "plain",
        # "digital",
        "crumpled",
    ]
)
def base_texture(request, size):
    texture = Image.open(TEST_PAGE_TEXTURES_DIR / (request.param + ".jpg"))
    texture = texture.resize(size)
    return texture


@pytest.fixture(
    params=[
        # "portrait",
        "landscape",
    ]
)
def orientation(request):
    return request.param


@pytest.fixture(
    params=[
        # 30,
        100,
    ]
)
def dpi(request):
    return request.param


@pytest.fixture(
    params=[
        # "brown",
        "sepia",
        # "gray",
        # "white",
        # "light_red",
        # "light_blue",
    ]
)
def color_name(request):
    return request.param


@pytest.fixture(
    params=[
        # "smooth",
        # "coarse",
        "neutral",
    ]
)
def texture_name(request):
    return request.param


@pytest.fixture(
    params=[
        # 30,
        70,
        # 150,
    ]
)
def color_intensity(request):
    return request.param


def random_flip(image):
    if rnd.choice([True, False]):
        image = image.transpose(Transpose.FLIP_LEFT_RIGHT)
    if rnd.choice([True, False]):
        image = image.transpose(Transpose.FLIP_TOP_BOTTOM)
    return image


@pytest.fixture
def color(color_name):
    return {
        "brown": "#7d6c5b",
        "sepia": "#b8af88",
        "gray": "#9c9c9c",
        "white": "#ffffff",
        "light_red": "#d68c8b",
        "light_blue": "#8bd6d6",
    }[color_name]


@pytest.fixture
def texture_fn(texture_name, size):
    if texture_name == "smooth":
        fn = blur
    elif texture_name == "coarse":
        fn = compose(overlay, juxt(blur, sharpen))
    else:
        fn = identity

    return normalize_image_function(fn)


def normalize_image_function(func):
    def inner(image):
        image = normalize_image_format_to_array(image)
        image = func(image)
        image = normalize_image_format_to_pil(image)
        return image

    return inner


@pytest.fixture
def texture(tinted_blank_page, base_texture):
    texture = superimpose(base_texture, tinted_blank_page)
    return texture


@pytest.fixture
def tinted_blank_page(size, color, color_intensity):
    tinted_page = Image.new("RGBA", size, color)
    tinted_page.putalpha(color_intensity)
    return tinted_page


@pytest.fixture
def blank_page(size, color, color_intensity):
    page = Image.new("RGBA", size, color=(255, 255, 255, 0))
    return page


@pytest.fixture
def size(dpi, orientation):
    if orientation == "portrait":
        size = (8.5 * dpi, 11 * dpi)
    elif orientation == "landscape":
        size = (11 * dpi, 8.5 * dpi)
    else:
        raise ValueError(f"Unknown orientation: {orientation}")
    size = tuple(map(int, size))
    return size


@pytest.fixture(
    params=[
        TwoColumnPagePartitioner,
        # RandomPagePartitioner
    ]
)
def page_partitioner(request):
    return request.param()


@pytest.fixture
def boxes(page_partitioner, blank_page):
    boxes = page_partitioner(blank_page)
    return boxes


@pytest.fixture
def prepared_texture(texture, texture_fn):
    texture = random_flip(texture)
    texture = texture_fn(texture)
    return texture


@pytest.fixture
def content_boxes(boxes):
    content_generator = ContentGenerator()
    content_boxes = content_generator(boxes)
    return content_boxes


@pytest.fixture
def page_with_opaque_content(
    blank_page, tinted_blank_page, prepared_texture, content_boxes
) -> Tuple[np.ndarray, Iterable[Rectangle]]:
    """Creates a page with content"""
    page = paste_contents(prepared_texture, content_boxes)

    return page, content_boxes


@pytest.fixture
def page_with_translucent_content(
    blank_page, tinted_blank_page, prepared_texture, content_boxes
) -> Tuple[np.ndarray, List[Rectangle]]:
    """Creates a page with content"""
    page_content = paste_contents(blank_page, content_boxes)
    page = blend_by_multiply(page_content, prepared_texture)

    return page, content_boxes


def blend_by_multiply(page_content, texture):
    def to_array(image: Image) -> np.ndarray:
        return np.array(image).astype(np.float32)

    texture.putalpha(255)
    page_content.putalpha(255)
    factor = 1.2
    enhancer = ImageEnhance.Contrast(texture)
    texture = enhancer.enhance(factor)

    page = blend_modes.multiply(
        *map(
            to_array,
            (
                page_content,
                texture,
            ),
        ),
        opacity=1,
    ).astype(np.uint8)
    return page


@pytest.fixture(scope="function")
def random_seeding():
    from synthesis.segment.plot import pick_colormap

    seed = str(rnd.randint(0, 2**32 - 1))
    logger.info(f"Random seed: {seed}")
    rnd.seed(seed)
    pick_colormap.cache_clear()


@pytest.fixture
def page_with_content(
    random_seeding,
    page_with_translucent_content,
    # page_with_opaque_content,
) -> np.ndarray:

    page, boxes = page_with_translucent_content
    # page, boxes = page_with_opaque_content

    return page, boxes