841 lines
25 KiB
Python
841 lines
25 KiB
Python
import io
|
|
import itertools
|
|
import os
|
|
import random
|
|
import string
|
|
import textwrap
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
from typing import Tuple, Union, Iterable, List
|
|
|
|
import albumentations as A
|
|
import cv2 as cv
|
|
import loguru
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
from PIL import Image, ImageOps, ImageFont, ImageDraw
|
|
from PIL.Image import Transpose
|
|
from faker import Faker
|
|
from matplotlib import pyplot as plt
|
|
from tabulate import tabulate
|
|
|
|
from cv_analysis.utils import star, rconj
|
|
from cv_analysis.utils.merging import merge_related_rectangles
|
|
from cv_analysis.utils.postprocessing import remove_overlapping, remove_included
|
|
|
|
Image_t = Union[Image.Image, np.ndarray]
|
|
#
|
|
# transform = A.Compose(
|
|
# [
|
|
# # geometric transforms
|
|
# A.HorizontalFlip(p=0.2),
|
|
# A.RandomRotate90(p=0.2),
|
|
# A.VerticalFlip(p=0.2),
|
|
# # brightness and contrast transforms
|
|
# A.OneOf(
|
|
# [
|
|
# A.RandomGamma(p=0.5),
|
|
# A.RandomBrightnessContrast(p=0.5),
|
|
# ],
|
|
# p=0.5,
|
|
# ),
|
|
# # noise transforms
|
|
# A.SomeOf(
|
|
# [
|
|
# A.Emboss(p=0.05),
|
|
# A.ImageCompression(p=0.05),
|
|
# A.PixelDropout(p=0.05),
|
|
# ],
|
|
# p=0.5,
|
|
# n=2,
|
|
# ),
|
|
# # color transforms
|
|
# A.SomeOf(
|
|
# [
|
|
# A.ColorJitter(p=1),
|
|
# A.RGBShift(p=1, r_shift_limit=0.1, g_shift_limit=0.1, b_shift_limit=0.1),
|
|
# A.ChannelShuffle(p=1),
|
|
# ],
|
|
# p=0.5,
|
|
# n=3, # 3 => all
|
|
# ),
|
|
# # blurring and sharpening transforms
|
|
# A.OneOf(
|
|
# [
|
|
# A.GaussianBlur(p=0.05),
|
|
# A.MotionBlur(p=0.05, blur_limit=21),
|
|
# A.Sharpen(p=0.05),
|
|
# ],
|
|
# p=0.5,
|
|
# ),
|
|
# # environmental transforms
|
|
# A.OneOf(
|
|
# [
|
|
# A.RandomRain(p=0.2, rain_type="drizzle"),
|
|
# A.RandomFog(p=0.2, fog_coef_upper=0.4),
|
|
# A.RandomSnow(p=0.2),
|
|
# ],
|
|
# p=0.5,
|
|
# ),
|
|
# ],
|
|
# p=0.5,
|
|
# )
|
|
from funcy import (
|
|
juxt,
|
|
compose,
|
|
identity,
|
|
lflatten,
|
|
lmap,
|
|
first,
|
|
iterate,
|
|
take,
|
|
last,
|
|
rest,
|
|
rcompose,
|
|
lsplit,
|
|
lfilter,
|
|
lzip,
|
|
keep,
|
|
lkeep,
|
|
)
|
|
|
|
from cv_analysis.locations import TEST_PAGE_TEXTURES_DIR
|
|
|
|
# transform = A.Compose(
|
|
# [
|
|
# # brightness and contrast transforms
|
|
# A.OneOf(
|
|
# [
|
|
# A.RandomGamma(p=0.2),
|
|
# A.RandomBrightnessContrast(p=0.2, brightness_limit=0.05, contrast_limit=0.05),
|
|
# ],
|
|
# p=0.5,
|
|
# ),
|
|
# # color transforms
|
|
# A.SomeOf(
|
|
# [
|
|
# A.ColorJitter(p=1),
|
|
# A.RGBShift(p=1, r_shift_limit=0.3, g_shift_limit=0.3, b_shift_limit=0.3),
|
|
# A.ChannelShuffle(p=1),
|
|
# ],
|
|
# p=1.0,
|
|
# n=3, # 3 => all
|
|
# ),
|
|
# # # blurring and sharpening transforms
|
|
# # A.OneOf(
|
|
# # [
|
|
# # A.GaussianBlur(p=0.05),
|
|
# # A.MotionBlur(p=0.05, blur_limit=21),
|
|
# # A.Sharpen(p=0.05),
|
|
# # ],
|
|
# # p=0.0,
|
|
# # ),
|
|
# ]
|
|
# )
|
|
from cv_analysis.utils.display import show_image
|
|
from cv_analysis.utils.drawing import draw_rectangles
|
|
from cv_analysis.utils.rectangle import Rectangle
|
|
|
|
transform = A.Compose(
|
|
[
|
|
# A.ColorJitter(p=1),
|
|
]
|
|
)
|
|
|
|
|
|
Color = Tuple[int, int, int]
|
|
|
|
|
|
@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 random.choice([True, False]):
|
|
image = image.transpose(Transpose.FLIP_LEFT_RIGHT)
|
|
if random.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 blur(image: np.ndarray):
|
|
return cv.blur(image, (3, 3))
|
|
|
|
|
|
def normalize_image_format_to_array(image: Image_t):
|
|
if isinstance(image, Image.Image):
|
|
return np.array(image)
|
|
return image
|
|
|
|
|
|
def normalize_image_format_to_pil(image: Image_t):
|
|
if isinstance(image, np.ndarray):
|
|
return Image.fromarray(image)
|
|
return image
|
|
|
|
|
|
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
|
|
|
|
|
|
def sharpen(image: np.ndarray):
|
|
return cv.filter2D(image, -1, np.array([[-1, -1, -1], [-1, 6, -1], [-1, -1, -1]]))
|
|
|
|
|
|
def overlay(images, mode=np.sum):
|
|
assert mode in [np.sum, np.max]
|
|
images = np.stack(list(images))
|
|
image = mode(images, axis=0)
|
|
image = (image / image.max() * 255).astype(np.uint8)
|
|
return image
|
|
|
|
|
|
@pytest.fixture
|
|
def texture(base_texture, color, color_intensity):
|
|
color_image = Image.new("RGBA", base_texture.size, color)
|
|
color_image.putalpha(color_intensity)
|
|
texture = superimpose_texture_with_transparency(base_texture, color_image)
|
|
return texture
|
|
|
|
|
|
def tint_image(src, color="#FFFFFF"):
|
|
src.load()
|
|
r, g, b, alpha = src.split()
|
|
gray = ImageOps.grayscale(src)
|
|
result = ImageOps.colorize(gray, (0, 0, 0), color)
|
|
result.putalpha(alpha)
|
|
return result
|
|
|
|
|
|
def color_shift_array(image: np.ndarray, color: Color):
|
|
"""Creates a 3-tensor from a 2-tensor by stacking the 2-tensor three times weighted by the color tuple."""
|
|
assert image.ndim == 3
|
|
assert image.shape[-1] == 3
|
|
assert isinstance(color, tuple)
|
|
assert max(color) <= 255
|
|
assert image.max() <= 255
|
|
|
|
color = np.array(color)
|
|
weights = color / color.sum() / 10
|
|
assert max(weights) <= 1
|
|
|
|
colored = (image * weights).astype(np.uint8)
|
|
|
|
assert colored.shape == image.shape
|
|
|
|
return colored
|
|
|
|
|
|
@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
|
|
|
|
|
|
def superimpose_texture_with_transparency(page: Image, texture: Image) -> Image:
|
|
"""Superimposes a noise image with transparency onto a page image."""
|
|
assert page.mode == "RGB"
|
|
assert texture.mode == "RGBA"
|
|
assert page.size == texture.size
|
|
page.paste(texture, (0, 0), texture)
|
|
return page
|
|
|
|
|
|
@pytest.fixture
|
|
def blank_page(texture, texture_fn) -> np.ndarray:
|
|
"""Creates a blank page with a given orientation and dpi."""
|
|
page = random_flip(texture)
|
|
page = texture_fn(page)
|
|
page_partitioner = PagePartitioner()
|
|
boxes = page_partitioner(page)
|
|
content_generator = ContentGenerator()
|
|
boxes = content_generator(boxes)
|
|
page = paste_contents(page, boxes)
|
|
page_partitioner.draw_boxes(page, boxes)
|
|
|
|
page = np.array(page)
|
|
return page
|
|
|
|
|
|
class ContentRectangle(Rectangle):
|
|
def __init__(self, x1, y1, x2, y2, content=None):
|
|
super().__init__(x1, y1, x2, y2)
|
|
self.content = content
|
|
|
|
|
|
class ContentGenerator:
|
|
def __init__(self):
|
|
self.constrain_layouts = True
|
|
|
|
def __call__(self, boxes: List[Rectangle]) -> Image:
|
|
random.shuffle(boxes)
|
|
|
|
char_boxes, text_boxes = lsplit(is_square_like, boxes)
|
|
|
|
if self.constrain_layouts:
|
|
char_boxes = merge_related_rectangles(char_boxes)
|
|
char_boxes = lfilter(is_square_like, char_boxes)
|
|
text_boxes = merge_related_rectangles(text_boxes)
|
|
|
|
text_boxes = lmap(generate_random_text_block, text_boxes)
|
|
plots = lmap(generate_random_table, every_nth(2, char_boxes))
|
|
tables = lmap(generate_random_table, every_nth(2, char_boxes[1:]))
|
|
|
|
boxes = text_boxes + plots + tables
|
|
boxes = remove_included(boxes)
|
|
boxes = remove_overlapping(boxes)
|
|
|
|
return boxes
|
|
|
|
|
|
def is_square_like(box: Rectangle):
|
|
return box.width / box.height > 0.5 and box.height / box.width > 0.5
|
|
|
|
|
|
def every_nth(n, iterable):
|
|
return itertools.islice(iterable, 0, None, n)
|
|
|
|
|
|
def generate_random_plot(rectangle: Rectangle) -> ContentRectangle:
|
|
block = RandomPlot(*rectangle.coords)
|
|
block.generate_random_plot(rectangle)
|
|
return block
|
|
|
|
|
|
def generate_random_table(rectangle: Rectangle) -> ContentRectangle:
|
|
block = RandomTable(*rectangle.coords)
|
|
block.generate_random_table(rectangle)
|
|
return block
|
|
|
|
|
|
@lru_cache(maxsize=None)
|
|
def get_random_seed():
|
|
return random.randint(0, 2**32 - 1)
|
|
|
|
|
|
class RandomContentRectangle(ContentRectangle):
|
|
def __init__(self, x1, y1, x2, y2, content=None, seed=None):
|
|
super().__init__(x1, y1, x2, y2, content)
|
|
self.seed = seed or get_random_seed()
|
|
self.random = random.Random(self.seed)
|
|
|
|
|
|
class RandomTable(RandomContentRectangle):
|
|
def __init__(self, x1, y1, x2, y2, seed=None):
|
|
super().__init__(x1, y1, x2, y2, seed=seed)
|
|
|
|
def generate_random_table(self, rectangle: Rectangle):
|
|
"""Generates the image of a random table.
|
|
|
|
Reference: https://stackoverflow.com/questions/35634238/how-to-save-a-pandas-dataframe-table-as-a-png
|
|
"""
|
|
|
|
text_table = self.generate_random_ascii_table(rectangle)
|
|
table_lines = text_table.split("\n")
|
|
image = write_lines_to_image(table_lines, rectangle)
|
|
|
|
self.content = image
|
|
|
|
def generate_random_dataframe(self, rectangle: Rectangle):
|
|
"""Generates a random dataframe that has as many rows and columns as to fit the given rectangle."""
|
|
|
|
image = Image.new("RGB", (rectangle.width, rectangle.height), color="white")
|
|
draw = ImageDraw.Draw(image)
|
|
text_size = draw.textsize("dummy", font=ImageFont.load_default())[1]
|
|
|
|
rows = rectangle.height // text_size
|
|
|
|
col_names = list(string.ascii_uppercase)
|
|
|
|
cols = min(rectangle.width // text_size, len(col_names))
|
|
|
|
df = pd.DataFrame(
|
|
np.random.randint(0, 100, size=(rows, cols)),
|
|
columns=col_names[:cols],
|
|
)
|
|
|
|
return df
|
|
|
|
def generate_random_ascii_table(self, rectangle: Rectangle):
|
|
df = self.generate_random_dataframe(rectangle)
|
|
table_format = random.choice(
|
|
[
|
|
"simple",
|
|
"grid",
|
|
"presto",
|
|
"psql",
|
|
"rst",
|
|
]
|
|
)
|
|
text_table = tabulate(df, headers="keys", tablefmt=table_format)
|
|
return text_table
|
|
|
|
|
|
def dump_plt_to_image(rectangle):
|
|
buf = io.BytesIO()
|
|
plt.savefig(buf, format="png", transparent=True)
|
|
buf.seek(0)
|
|
image = Image.open(buf)
|
|
image = image.resize((rectangle.width, rectangle.height))
|
|
buf.close()
|
|
plt.close()
|
|
return image
|
|
|
|
|
|
class RandomFontPicker:
|
|
def __init__(self, font_dir=None):
|
|
self.fonts = get_fonts(font_dir)
|
|
self.fonts_lower = [font.lower() for font in self.fonts]
|
|
|
|
self.test_image = Image.new("RGB", (200, 200), (255, 255, 255))
|
|
self.draw = ImageDraw.Draw(self.test_image)
|
|
|
|
def pick_random_font_available_on_system(self, includes=None, excludes=None) -> ImageFont: # FIXME: Slow!
|
|
|
|
self.shuffle_fonts()
|
|
|
|
includes_pattern = (lambda f: includes.lower() in f) if includes else (lambda f: True)
|
|
excludes_pattern = (lambda f: excludes.lower() not in f) if excludes else (lambda f: True)
|
|
|
|
mask = lmap(lambda f: includes_pattern(f) and excludes_pattern(f), self.fonts_lower)
|
|
fonts = itertools.compress(self.fonts, mask)
|
|
fonts = keep(map(self.load_font, fonts))
|
|
fonts = filter(self.font_is_renderable, fonts) # FIXME: this does not work
|
|
|
|
font = first(fonts)
|
|
loguru.logger.debug(f"Using font: {font}")
|
|
return font
|
|
|
|
def shuffle_fonts(self):
|
|
l = lzip(self.fonts, self.fonts_lower)
|
|
random.shuffle(l)
|
|
self.fonts, self.fonts_lower = lzip(*l)
|
|
|
|
def pick_random_mono_space_font_available_on_system(self) -> ImageFont:
|
|
return self.pick_random_font_available_on_system(includes="mono", excludes="oblique")
|
|
|
|
@lru_cache(maxsize=None)
|
|
def load_font(self, font: str):
|
|
loguru.logger.trace(f"Loading font: {font}")
|
|
try:
|
|
return ImageFont.truetype(font, size=11)
|
|
except OSError:
|
|
return None
|
|
|
|
@lru_cache(maxsize=None)
|
|
def font_is_renderable(self, font):
|
|
text_size = self.draw.textsize("Test String", font=font)
|
|
return text_size[0] > 0 and text_size[1]
|
|
|
|
|
|
def get_fonts(path: Path = None) -> List[str]:
|
|
path = path or Path("/usr/share/fonts")
|
|
fonts = list(path.rglob("*.ttf"))
|
|
fonts = [font.name for font in fonts]
|
|
return fonts
|
|
|
|
|
|
@lru_cache(maxsize=None)
|
|
def get_font_picker(font_dir=None):
|
|
return RandomFontPicker(font_dir=font_dir)
|
|
|
|
|
|
@lru_cache(maxsize=None)
|
|
def pick_random_mono_space_font_available_on_system():
|
|
font_picker = get_font_picker()
|
|
return font_picker.pick_random_mono_space_font_available_on_system()
|
|
|
|
|
|
@lru_cache(maxsize=None)
|
|
def pick_random_font_available_on_system():
|
|
font_picker = get_font_picker()
|
|
return font_picker.pick_random_font_available_on_system(includes="mono")
|
|
|
|
|
|
class RandomPlot(RandomContentRectangle):
|
|
def __init__(self, x1, y1, x2, y2, seed=None):
|
|
super().__init__(x1, y1, x2, y2, seed=seed)
|
|
|
|
cmap_name = self.random.choice(
|
|
[
|
|
"viridis",
|
|
"plasma",
|
|
"inferno",
|
|
"magma",
|
|
"cividis",
|
|
],
|
|
)
|
|
self.cmap = plt.get_cmap(cmap_name)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
pass
|
|
|
|
def generate_random_plot(self, rectangle: Rectangle):
|
|
# noinspection PyArgumentList
|
|
random.choice(
|
|
[
|
|
self.generate_random_line_plot,
|
|
self.generate_random_bar_plot,
|
|
self.generate_random_scatter_plot,
|
|
self.generate_random_histogram,
|
|
self.generate_random_pie_chart,
|
|
]
|
|
)(rectangle)
|
|
|
|
def generate_random_bar_plot(self, rectangle: Rectangle):
|
|
x = sorted(np.random.randint(low=1, high=11, size=5))
|
|
y = np.random.randint(low=1, high=11, size=5)
|
|
image = self.__generate_random_plot(plt.bar, rectangle, x, y)
|
|
self.content = image
|
|
|
|
def generate_random_line_plot(self, rectangle: Rectangle):
|
|
f = random.choice([np.sin, np.cos, np.tan, np.exp, np.log, np.sqrt, np.square])
|
|
|
|
x = np.linspace(0, 10, 100)
|
|
y = f(x)
|
|
|
|
image = self.__generate_random_plot(plt.plot, rectangle, x, y)
|
|
self.content = image
|
|
|
|
def generate_random_scatter_plot(self, rectangle: Rectangle):
|
|
x = np.random.normal(size=100)
|
|
y = np.random.normal(size=100)
|
|
image = self.__generate_random_plot(plt.scatter, rectangle, x, y)
|
|
self.content = image
|
|
|
|
def generate_random_histogram(self, rectangle: Rectangle):
|
|
x = np.random.normal(size=100)
|
|
image = self.__generate_random_plot(plt.hist, rectangle, x, 10)
|
|
self.content = image
|
|
|
|
def generate_random_pie_chart(self, rectangle: Rectangle):
|
|
x = np.random.uniform(size=10)
|
|
image = self.__generate_random_plot(
|
|
plt.pie, rectangle, x, None, plot_kwargs=self.generate_plot_kwargs(keywords=["a"])
|
|
)
|
|
self.content = image
|
|
|
|
def generate_plot_kwargs(self, keywords=None):
|
|
|
|
kwargs = {
|
|
"color": random.choice(self.cmap.colors),
|
|
"linestyle": random.choice(["-", "--", "-.", ":"]),
|
|
"linewidth": random.uniform(0.5, 2),
|
|
}
|
|
|
|
return kwargs if not keywords else {k: v for k, v in kwargs.items() if k in keywords}
|
|
|
|
def __generate_random_plot(self, plot_fn, rectangle: Rectangle, x, y, plot_kwargs=None):
|
|
|
|
plot_kwargs = self.generate_plot_kwargs() if plot_kwargs is None else plot_kwargs
|
|
|
|
fig, ax = plt.subplots()
|
|
fig.set_size_inches(rectangle.width / 100, rectangle.height / 100)
|
|
fig.tight_layout(pad=0)
|
|
|
|
plot_fn(
|
|
x,
|
|
y,
|
|
**plot_kwargs,
|
|
)
|
|
ax.set_facecolor("none")
|
|
|
|
maybe() and ax.set_title("Figure Title")
|
|
|
|
# disable axes at random
|
|
maybe() and ax.set_xticks([])
|
|
maybe() and ax.set_yticks([])
|
|
maybe() and ax.set_xticklabels([])
|
|
maybe() and ax.set_yticklabels([])
|
|
maybe() and ax.set_xlabel("")
|
|
maybe() and ax.set_ylabel("")
|
|
maybe() and ax.set_title("")
|
|
maybe() and ax.set_frame_on(False)
|
|
|
|
# remove spines at random
|
|
maybe() and (ax.spines["top"].set_visible(False) or ax.spines["right"].set_visible(False))
|
|
|
|
image = dump_plt_to_image(rectangle)
|
|
return image
|
|
|
|
|
|
def maybe():
|
|
return random.random() > 0.9
|
|
|
|
|
|
def generate_random_text_block(rectangle: Rectangle) -> ContentRectangle:
|
|
block = RandomTextBlock(*rectangle.coords)
|
|
block.generate_random_text(rectangle)
|
|
return block
|
|
|
|
|
|
def write_lines_to_image(lines: List[str], rectangle: Rectangle, font=None) -> Image.Image:
|
|
def write_line(line, line_number):
|
|
draw.text((0, line_number * text_size), line, font=font, fill=(0, 0, 0, 200))
|
|
|
|
font = font or pick_random_mono_space_font_available_on_system()
|
|
|
|
image = Image.new("RGBA", (rectangle.width, rectangle.height), (0, 255, 255, 0))
|
|
draw = ImageDraw.Draw(image)
|
|
text_size = draw.textsize(first(lines), font=font)[1]
|
|
|
|
for line_number, line in enumerate(lines):
|
|
write_line(line, line_number)
|
|
|
|
return image
|
|
|
|
|
|
class RandomTextBlock(ContentRectangle):
|
|
def __init__(self, x1, y1, x2, y2):
|
|
super().__init__(x1, y1, x2, y2)
|
|
self.blank_line_percentage = random.uniform(0, 0.5)
|
|
self.font = pick_random_font_available_on_system()
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
pass
|
|
|
|
def generate_random_text(self, rectangle: Rectangle):
|
|
lines = generate_random_text_lines(rectangle, self.format_lines)
|
|
image = write_lines_to_image(lines, rectangle, self.font)
|
|
self.content = image
|
|
|
|
def format_lines(self, lines, last_full):
|
|
def truncate_current_line():
|
|
return random.random() < self.blank_line_percentage and last_full
|
|
|
|
# This is meant to be read from the bottom up.
|
|
current_line_shall_not_be_a_full_line = truncate_current_line()
|
|
line_formatter = self.truncate_line if current_line_shall_not_be_a_full_line else identity
|
|
format_current_line = compose(line_formatter, first)
|
|
move_current_line_to_back = star(rconj)
|
|
split_first_line_from_lines_and_format_the_former = juxt(rest, format_current_line)
|
|
split_off_current_line_then_format_it_then_move_it_to_the_back = rcompose(
|
|
split_first_line_from_lines_and_format_the_former,
|
|
move_current_line_to_back,
|
|
)
|
|
current_line_is_a_full_line = not current_line_shall_not_be_a_full_line
|
|
# Start reading here and move up.
|
|
return split_off_current_line_then_format_it_then_move_it_to_the_back(lines), current_line_is_a_full_line
|
|
|
|
def format_line(self, line, full=True):
|
|
line = self.truncate_line(line) if not full else line
|
|
return line, full
|
|
|
|
def truncate_line(self, line: str):
|
|
n_trailing_words = random.randint(0, 4)
|
|
line = " ".join(line.split()[-n_trailing_words - 1 : -1]).replace(".", "")
|
|
line = line + ".\n" if line else line
|
|
return line
|
|
|
|
|
|
def generate_random_text_lines(rectangle: Rectangle, line_formatter=identity) -> List[str]:
|
|
text = Faker().paragraph(nb_sentences=3000, variable_nb_sentences=False, ext_word_list=None)
|
|
unformatted_lines = textwrap.wrap(text, width=rectangle.width, break_long_words=False)
|
|
# each iteration of the line formatter function formats one more line and adds it to the back of the list
|
|
formatted_lines_generator = iterate(star(line_formatter), (unformatted_lines, True))
|
|
# hence do as many iterations as there are lines in the rectangle
|
|
lines_per_iteration = take(len(unformatted_lines), formatted_lines_generator)
|
|
# and then take the lines from the last iteration of the function
|
|
formatted_lines, _ = last(lines_per_iteration)
|
|
|
|
return formatted_lines
|
|
|
|
|
|
def paste_content(page, content_box: ContentRectangle):
|
|
assert page.mode == "RGB"
|
|
assert content_box.content.mode == "RGBA"
|
|
page.paste(content_box.content, (content_box.x1, content_box.y1), content_box.content)
|
|
return page
|
|
|
|
|
|
def paste_contents(page, contents: Iterable[ContentRectangle]):
|
|
for content in contents:
|
|
paste_content(page, content)
|
|
return page
|
|
|
|
|
|
class PagePartitioner:
|
|
def __init__(self):
|
|
self.left_margin_percentage = 0.05
|
|
self.right_margin_percentage = 0.05
|
|
self.top_margin_percentage = 0.1
|
|
self.bottom_margin_percentage = 0.1
|
|
|
|
self.margin_percentage = 0.007
|
|
self.max_depth = 3
|
|
self.initial_recursion_probability = 1
|
|
self.recursion_probability_decay = 0.1
|
|
|
|
def __call__(self, page: Image.Image) -> List[Rectangle]:
|
|
left_margin = int(page.width * self.left_margin_percentage)
|
|
right_margin = int(page.width * self.right_margin_percentage)
|
|
top_margin = int(page.height * self.top_margin_percentage)
|
|
bottom_margin = int(page.height * self.bottom_margin_percentage)
|
|
|
|
box = Rectangle(left_margin, top_margin, page.width - right_margin, page.height - bottom_margin)
|
|
boxes = lflatten(self.generate_content_boxes(box))
|
|
# boxes = self.drop_small_boxes(boxes, *page.size)
|
|
# boxes = merge_related_rectangles(boxes)
|
|
boxes = list(boxes)
|
|
return boxes
|
|
|
|
def draw_boxes(self, page: Image, boxes: Iterable[Rectangle]):
|
|
image = draw_rectangles(page, boxes, filled=False, annotate=True)
|
|
show_image(image)
|
|
|
|
def generate_content_boxes(self, box: Rectangle, depth=0):
|
|
if depth >= self.max_depth:
|
|
yield box
|
|
else:
|
|
child_boxes = self.generate_random_child_boxes(box)
|
|
if self.recurse(depth):
|
|
yield from (self.generate_content_boxes(b, depth + 1) for b in child_boxes)
|
|
else:
|
|
yield child_boxes
|
|
|
|
def generate_random_child_boxes(self, box: Rectangle) -> Tuple[Rectangle, Rectangle]:
|
|
axis = random.choice(["x", "y"])
|
|
|
|
edge_anchor_point, edge_length = (box.x1, box.width) if axis == "x" else (box.y1, box.height)
|
|
split_coordinate = random.uniform(0.3, 0.7) * edge_length + edge_anchor_point
|
|
child_boxes = self.get_child_boxes(box, split_coordinate, axis)
|
|
return child_boxes
|
|
|
|
def get_child_boxes(self, box: Rectangle, split_coordinate, axis) -> Tuple[Rectangle, Rectangle]:
|
|
def low(p):
|
|
return p * (1 + self.margin_percentage)
|
|
|
|
def high(p):
|
|
return p * (1 - self.margin_percentage)
|
|
|
|
if axis == "x":
|
|
return (
|
|
Rectangle(low(box.x1), low(box.y1), high(split_coordinate), high(box.y2)),
|
|
Rectangle(low(split_coordinate), low(box.y1), high(box.x2), high(box.y2)),
|
|
)
|
|
else:
|
|
return (
|
|
Rectangle(low(box.x1), low(box.y1), high(box.x2), high(split_coordinate)),
|
|
Rectangle(low(box.x1), low(split_coordinate), high(box.x2), high(box.y2)),
|
|
)
|
|
|
|
def recurse(self, depth):
|
|
return random.random() <= self.recursion_probability(depth)
|
|
|
|
def recursion_probability(self, depth):
|
|
return self.initial_recursion_probability * (1 - self.recursion_probability_decay) ** depth
|
|
|
|
def drop_small_boxes(
|
|
self,
|
|
boxes: Iterable[Rectangle],
|
|
page_width,
|
|
page_height,
|
|
min_percentage=0.13,
|
|
) -> List[Rectangle]:
|
|
min_width = page_width * min_percentage
|
|
min_height = page_height * min_percentage
|
|
return [b for b in boxes if b.width > min_width and b.height > min_height]
|