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 cv_analysis.locations import TEST_PAGE_TEXTURES_DIR from cv_analysis.logging import logger 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 @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