315 lines
7.3 KiB
Python
315 lines
7.3 KiB
Python
import itertools
|
|
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 loguru import logger
|
|
|
|
from cv_analysis.utils import zipmap, every_nth
|
|
from cv_analysis.utils.conversion import normalize_image_format_to_array, normalize_image_format_to_pil
|
|
from cv_analysis.utils.geometric import is_square_like
|
|
from cv_analysis.utils.image_operations import blur, sharpen, overlay, superimpose
|
|
from cv_analysis.utils.merging import merge_related_rectangles
|
|
from cv_analysis.utils.postprocessing import remove_overlapping, remove_included
|
|
from synthesis.partitioner.two_column import TwoColumnPagePartitioner
|
|
from synthesis.random import rnd
|
|
from synthesis.segment.segments import (
|
|
generate_random_plot_with_caption,
|
|
generate_recursive_random_table_with_caption,
|
|
generate_random_text_block,
|
|
)
|
|
from synthesis.segment.table.table import paste_contents
|
|
|
|
logger.remove()
|
|
logger.add(sys.stderr, level="INFO")
|
|
|
|
|
|
from funcy import (
|
|
juxt,
|
|
compose,
|
|
identity,
|
|
lsplit,
|
|
lfilter,
|
|
)
|
|
|
|
from cv_analysis.locations import TEST_PAGE_TEXTURES_DIR
|
|
|
|
from cv_analysis.utils.display import show_image
|
|
from cv_analysis.utils.rectangle import Rectangle
|
|
|
|
|
|
@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)
|
|
# texture.putalpha(255) # ISSUE 1!!!
|
|
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
|
|
|
|
|
|
class ContentGenerator:
|
|
def __init__(self):
|
|
self.constrain_layouts = True
|
|
|
|
def __call__(self, boxes: List[Rectangle]) -> Image:
|
|
rnd.shuffle(boxes)
|
|
|
|
figure_boxes, text_boxes = lsplit(is_square_like, boxes)
|
|
|
|
if self.constrain_layouts:
|
|
figure_boxes = merge_related_rectangles(figure_boxes)
|
|
figure_boxes = lfilter(is_square_like, figure_boxes)
|
|
text_boxes = merge_related_rectangles(text_boxes)
|
|
|
|
boxes = list(
|
|
itertools.chain(
|
|
map(generate_random_text_block, every_nth(2, text_boxes)),
|
|
*zipmap(generate_recursive_random_table_with_caption, every_nth(2, text_boxes[1:])),
|
|
*zipmap(generate_recursive_random_table_with_caption, every_nth(2, figure_boxes)),
|
|
*zipmap(generate_random_plot_with_caption, every_nth(2, figure_boxes[1:])),
|
|
)
|
|
)
|
|
|
|
if self.constrain_layouts:
|
|
boxes = remove_included(boxes)
|
|
boxes = remove_overlapping(boxes)
|
|
|
|
return boxes
|
|
|
|
|
|
# TODO: deduplicate with generate_random_table_with_caption
|
|
|
|
|
|
@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
|
|
def page_with_content(
|
|
page_with_translucent_content,
|
|
# page_with_opaque_content,
|
|
) -> np.ndarray:
|
|
|
|
page, boxes = page_with_translucent_content
|
|
# page, boxes = page_with_opaque_content
|
|
|
|
draw_boxes(page, boxes)
|
|
|
|
return page
|
|
|
|
|
|
def draw_boxes(page: Image, boxes: Iterable[Rectangle]):
|
|
from cv_analysis.utils.drawing import draw_rectangles
|
|
|
|
page = draw_rectangles(page, boxes, filled=False, annotate=True)
|
|
show_image(page, backend="pil")
|