Refactoring: Move
Move table generation related code into new table module
This commit is contained in:
parent
9fd87aff8e
commit
f0072b0852
0
synthesis/segment/table/__init__.py
Normal file
0
synthesis/segment/table/__init__.py
Normal file
0
synthesis/segment/table/cell.py
Normal file
0
synthesis/segment/table/cell.py
Normal file
386
synthesis/segment/table/table.py
Normal file
386
synthesis/segment/table/table.py
Normal file
@ -0,0 +1,386 @@
|
||||
import random
|
||||
from copy import deepcopy
|
||||
from enum import Enum
|
||||
from functools import lru_cache, partial
|
||||
from math import sqrt
|
||||
from typing import List, Iterable
|
||||
|
||||
from PIL import Image, ImageDraw
|
||||
from funcy import chunks, mapcat, repeatedly
|
||||
from loguru import logger
|
||||
|
||||
from cv_analysis.utils.geometric import is_square_like
|
||||
from cv_analysis.utils.image_operations import superimpose
|
||||
|
||||
from cv_analysis.utils.rectangle import Rectangle
|
||||
from cv_analysis.utils.spacial import area
|
||||
from synthesis.random import rnd, possibly
|
||||
from synthesis.segment.content_rectangle import ContentRectangle
|
||||
from synthesis.segment.plot import RandomPlot, pick_colormap
|
||||
from synthesis.segment.random_content_rectangle import RandomContentRectangle
|
||||
from synthesis.segment.text_block import TextBlock
|
||||
from synthesis.text.font import pick_random_font_available_on_system
|
||||
from synthesis.text.text import generate_random_words, generate_random_number
|
||||
|
||||
|
||||
def generate_random_plot(rectangle: Rectangle) -> ContentRectangle:
|
||||
block = RandomPlot(*rectangle.coords)
|
||||
block.content = rectangle.content if isinstance(rectangle, ContentRectangle) else None # TODO: Refactor
|
||||
block.generate_random_plot(rectangle)
|
||||
return block
|
||||
|
||||
|
||||
def generate_recursive_random_table(rectangle: Rectangle, **kwargs) -> ContentRectangle:
|
||||
block = RecursiveRandomTable(*rectangle.coords, **kwargs)
|
||||
if isinstance(rectangle, RecursiveRandomTable):
|
||||
block.content = rectangle.content if rectangle.content else None # TODO: Refactor
|
||||
block.generate_random_table()
|
||||
return block
|
||||
|
||||
|
||||
class Size(Enum):
|
||||
SMALL = 120
|
||||
MEDIUM = 180
|
||||
LARGE = 300
|
||||
|
||||
|
||||
def get_size_class(rectangle: Rectangle):
|
||||
size = get_size(rectangle)
|
||||
if size < Size.SMALL.value:
|
||||
return Size.SMALL
|
||||
elif size < Size.LARGE.value:
|
||||
return Size.MEDIUM
|
||||
else:
|
||||
return Size.LARGE
|
||||
|
||||
|
||||
def get_size(rectangle: Rectangle):
|
||||
size = sqrt(area(rectangle))
|
||||
return size
|
||||
|
||||
|
||||
def get_random_color_complementing_color_map(colormap):
|
||||
def color_complement(r, g, b):
|
||||
"""Reference: https://stackoverflow.com/a/40234924"""
|
||||
|
||||
def hilo(a, b, c):
|
||||
if c < b:
|
||||
b, c = c, b
|
||||
if b < a:
|
||||
a, b = b, a
|
||||
if c < b:
|
||||
b, c = c, b
|
||||
return a + c
|
||||
|
||||
k = hilo(r, g, b)
|
||||
return tuple(k - u for u in (r, g, b))
|
||||
|
||||
color = colormap(0.2)[:3]
|
||||
color = [int(255 * v) for v in color]
|
||||
color = color_complement(*color)
|
||||
return color
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def get_random_background_color():
|
||||
return tuple([*get_random_color_complementing_color_map(pick_colormap()), rnd.randint(100, 210)])
|
||||
|
||||
|
||||
class RecursiveRandomTable(RandomContentRectangle):
|
||||
def __init__(self, x1, y1, x2, y2, border_width=1, layout: str = None, double_rule=False):
|
||||
"""A table with a random number of rows and columns, and random content in each cell.
|
||||
|
||||
Args:
|
||||
x1: x-coordinate of the top-left corner
|
||||
y1: y-coordinate of the top-left corner
|
||||
x2: x-coordinate of the bottom-right corner
|
||||
y2: y-coordinate of the bottom-right corner
|
||||
border_width: width of the table border
|
||||
layout: layout of the table, either "horizontal", "vertical", "closed", or "open"
|
||||
double_rule: whether to use double rules as the top and bottom rules
|
||||
"""
|
||||
|
||||
assert layout in [None, "horizontal", "vertical", "closed", "open"]
|
||||
|
||||
super().__init__(x1, y1, x2, y2)
|
||||
|
||||
self.double_rule = double_rule
|
||||
self.double_rule_width = (3 * border_width) if self.double_rule else 0
|
||||
|
||||
self.n_columns = rnd.randint(1, max(self.width // 100, 1))
|
||||
self.n_rows = rnd.randint(1, max((self.height - 2 * self.double_rule_width) // rnd.randint(17, 100), 1))
|
||||
self.cell_size = (self.width / self.n_columns, (self.height - 2 * self.double_rule_width) / self.n_rows)
|
||||
|
||||
self.content = Image.new("RGBA", (self.width, self.height), (255, 255, 255, 0))
|
||||
|
||||
self.background_color = get_random_background_color()
|
||||
|
||||
logger.info(f"Background color: {self.background_color}")
|
||||
|
||||
self.layout = layout or self.pick_random_layout()
|
||||
logger.debug(f"Layout: {self.layout}")
|
||||
|
||||
def pick_random_layout(self):
|
||||
|
||||
if self.n_columns == 1 and self.n_rows == 1:
|
||||
layout = "closed"
|
||||
elif self.n_columns == 1:
|
||||
layout = rnd.choice(["vertical", "closed"])
|
||||
elif self.n_rows == 1:
|
||||
layout = rnd.choice(["horizontal", "closed"])
|
||||
else:
|
||||
layout = rnd.choice(["closed", "horizontal", "vertical", "open"])
|
||||
|
||||
return layout
|
||||
|
||||
def generate_random_table(self):
|
||||
cells = self.generate_table()
|
||||
cells = list(self.fill_cells_with_content(cells))
|
||||
# FIXME: There is a bug here: Table rule is not drawn correctly, actually we want to do cells = ...
|
||||
list(self.draw_cell_borders(cells))
|
||||
|
||||
self.content = paste_contents(self.content, cells)
|
||||
assert self.content.mode == "RGBA"
|
||||
|
||||
def fill_cells_with_content(self, cells):
|
||||
yield from map(self.build_cell, cells)
|
||||
|
||||
def build_cell(self, cell):
|
||||
|
||||
if self.__is_a_small_cell(cell):
|
||||
cell = self.build_small_cell(cell)
|
||||
|
||||
elif self.__is_a_medium_sized_cell(cell):
|
||||
cell = self.build_medium_sized_cell(cell)
|
||||
|
||||
elif self.__is_a_large_cell(cell):
|
||||
cell = self.build_large_cell(cell)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid cell size: {get_size(cell)}")
|
||||
|
||||
assert cell.content.mode == "RGBA"
|
||||
|
||||
return cell
|
||||
|
||||
def __is_a_small_cell(self, cell):
|
||||
return get_size(cell) <= Size.SMALL.value
|
||||
|
||||
def __is_a_medium_sized_cell(self, cell):
|
||||
return get_size(cell) <= Size.MEDIUM.value
|
||||
|
||||
def __is_a_large_cell(self, cell):
|
||||
return get_size(cell) > Size.MEDIUM.value
|
||||
|
||||
def build_small_cell(self, cell):
|
||||
|
||||
content = (possibly() and generate_random_words(1, 3)) or (
|
||||
generate_random_number()
|
||||
+ ((possibly() and " " + rnd.choice(["$", "£", "%", "EUR", "USD", "CAD", "ADA"])) or "")
|
||||
)
|
||||
|
||||
return generate_text_block(cell, content)
|
||||
|
||||
def build_medium_sized_cell(self, cell):
|
||||
|
||||
choice = rnd.choice(["plot", "recurse"])
|
||||
|
||||
if choice == "plot":
|
||||
return generate_random_plot(cell)
|
||||
|
||||
elif choice == "recurse":
|
||||
return generate_recursive_random_table(
|
||||
cell,
|
||||
border_width=1,
|
||||
layout=random.choice(["open", "horizontal", "vertical"]),
|
||||
double_rule=False,
|
||||
)
|
||||
|
||||
else:
|
||||
return generate_text_block(cell, f"{choice} {get_size(cell):.0f} {get_size_class(cell).name}")
|
||||
|
||||
def build_large_cell(self, cell):
|
||||
choice = rnd.choice(["plot", "recurse"])
|
||||
|
||||
logger.debug(f"Generating {choice} {get_size(cell):.0f} {get_size_class(cell).name}")
|
||||
|
||||
if choice == "plot" and is_square_like(cell):
|
||||
return generate_random_plot(cell)
|
||||
|
||||
else:
|
||||
logger.debug(f"recurse {get_size(cell):.0f} {get_size_class(cell).name}")
|
||||
return generate_recursive_random_table(
|
||||
cell,
|
||||
border_width=1,
|
||||
layout=random.choice(["open", "horizontal", "vertical"]),
|
||||
double_rule=False,
|
||||
)
|
||||
|
||||
def draw_cell_borders(self, cells: List[ContentRectangle]):
|
||||
def draw_edges_based_on_position(cell: Cell, col_idx, row_index):
|
||||
# Draw the borders of the cell based on its position in the table
|
||||
if col_idx < self.n_columns - 1:
|
||||
cell.draw_right_border()
|
||||
|
||||
if row_index < self.n_rows - 1:
|
||||
cell.draw_bottom_border()
|
||||
|
||||
columns = chunks(self.n_rows, cells)
|
||||
for col_idx, column in enumerate(columns):
|
||||
for row_index, cell in enumerate(column):
|
||||
# TODO: Refactor
|
||||
c = Cell(*cell.coords, self.background_color)
|
||||
c.content = cell.content
|
||||
draw_edges_based_on_position(c, col_idx, row_index)
|
||||
yield cell
|
||||
|
||||
if self.layout == "closed":
|
||||
# TODO: Refactor
|
||||
c = Cell(*self.coords, self.background_color)
|
||||
c.content = self.content
|
||||
c.draw()
|
||||
yield self
|
||||
|
||||
# TODO: Refactor
|
||||
if self.double_rule:
|
||||
c1 = Cell(*self.coords)
|
||||
c1.draw_top_border(width=1)
|
||||
c1.draw_bottom_border(width=1)
|
||||
|
||||
x1, y1, x2, y2 = self.coords
|
||||
c2 = Cell(x1, y1 + self.double_rule_width, x2, y2 - self.double_rule_width)
|
||||
c2.draw_top_border(width=1)
|
||||
c2.draw_bottom_border(width=1)
|
||||
|
||||
c = superimpose(c1.content, c2.content)
|
||||
|
||||
self.content = superimpose(c, self.content)
|
||||
|
||||
yield self
|
||||
|
||||
def generate_table(self) -> Iterable[ContentRectangle]:
|
||||
yield from mapcat(self.generate_column, range(self.n_columns))
|
||||
|
||||
def generate_column(self, column_index) -> Iterable[ContentRectangle]:
|
||||
logger.trace(f"Generating column {column_index}.")
|
||||
generate_cell_for_row_index = partial(self.generate_cell, column_index)
|
||||
yield from map(generate_cell_for_row_index, range(self.n_rows))
|
||||
|
||||
def generate_cell(self, column_index, row_index) -> ContentRectangle:
|
||||
w, h = self.cell_size
|
||||
x1, y1 = (column_index * w), (row_index * h) + self.double_rule_width
|
||||
x2, y2 = x1 + w, y1 + h
|
||||
logger.trace(f"Generating cell ({row_index}, {column_index}) at ({x1}, {y1}, {x2}, {y2}).")
|
||||
return Cell(x1, y1, x2, y2, self.background_color)
|
||||
|
||||
def generate_column_names(self):
|
||||
column_names = repeatedly(self.generate_column_name, self.n_columns)
|
||||
return column_names
|
||||
|
||||
def generate_column_name(self):
|
||||
column_name = generate_random_words(1, 3)
|
||||
return column_name
|
||||
|
||||
|
||||
class Cell(ContentRectangle):
|
||||
def __init__(self, x1, y1, x2, y2, color=None):
|
||||
super().__init__(x1, y1, x2, y2)
|
||||
|
||||
self.background_color = color or (255, 255, 255, 0)
|
||||
|
||||
# to debug use random border color: tuple([random.randint(100, 200) for _ in range(3)] + [255])
|
||||
self.cell_border_color = (0, 0, 0, 255)
|
||||
|
||||
self.border_width = 1
|
||||
self.inset = 1
|
||||
|
||||
self.content = Image.new("RGBA", (self.width, self.height))
|
||||
self.fill()
|
||||
|
||||
def draw_top_border(self, width=None):
|
||||
self.draw_line((0, 0, self.width - self.inset, 0), width=width)
|
||||
return self
|
||||
|
||||
def draw_bottom_border(self, width=None):
|
||||
self.draw_line((0, self.height - self.inset, self.width - self.inset, self.height - self.inset), width=width)
|
||||
return self
|
||||
|
||||
def draw_left_border(self, width=None):
|
||||
self.draw_line((0, 0, 0, self.height), width=width)
|
||||
return self
|
||||
|
||||
def draw_right_border(self, width=None):
|
||||
self.draw_line((self.width - self.inset, 0, self.width - self.inset, self.height), width=width)
|
||||
return self
|
||||
|
||||
def draw_line(self, points, width=None):
|
||||
width = width or self.border_width
|
||||
draw = ImageDraw.Draw(self.content)
|
||||
draw.line(points, width=width, fill=self.cell_border_color)
|
||||
return self
|
||||
|
||||
def draw(self, width=None):
|
||||
self.draw_top_border(width=width)
|
||||
self.draw_bottom_border(width=width)
|
||||
self.draw_left_border(width=width)
|
||||
self.draw_right_border(width=width)
|
||||
return self
|
||||
|
||||
def draw_top_left_corner(self, width=None):
|
||||
self.draw_line((0, 0, 0, 0), width=width)
|
||||
self.draw_line((0, 0, 0, 0), width=width)
|
||||
return self
|
||||
|
||||
def draw_top_right_corner(self, width=None):
|
||||
self.draw_line((self.width - self.inset, 0, self.width - self.inset, 0), width=width)
|
||||
self.draw_line((self.width - self.inset, 0, self.width - self.inset, 0), width=width)
|
||||
return self
|
||||
|
||||
def draw_bottom_left_corner(self, width=None):
|
||||
self.draw_line((0, self.height - self.inset, 0, self.height - self.inset), width=width)
|
||||
self.draw_line((0, self.height - self.inset, 0, self.height - self.inset), width=width)
|
||||
return self
|
||||
|
||||
def draw_bottom_right_corner(self, width=None):
|
||||
self.draw_line(
|
||||
(self.width - self.inset, self.height - self.inset, self.width - self.inset, self.height - self.inset),
|
||||
width=width,
|
||||
)
|
||||
self.draw_line(
|
||||
(self.width - self.inset, self.height - self.inset, self.width - self.inset, self.height - self.inset),
|
||||
width=width,
|
||||
)
|
||||
return self
|
||||
|
||||
def fill(self, color=None):
|
||||
color = color or self.background_color
|
||||
image = Image.new("RGBA", (self.width, self.height), color=color)
|
||||
self.content = superimpose(image, self.content)
|
||||
return self
|
||||
|
||||
|
||||
def generate_text_block(rectangle: Rectangle, text) -> ContentRectangle:
|
||||
block = TextBlock(
|
||||
*rectangle.coords,
|
||||
font=pick_random_font_available_on_system(
|
||||
includes=("serif", "sans-serif", "bold"),
|
||||
excludes=("mono", "italic", "oblique", "cursive"),
|
||||
),
|
||||
font_size=30, # TODO: De-hardcode font size... Seems to have no effect on top of that
|
||||
)
|
||||
block.content = rectangle.content if isinstance(rectangle, ContentRectangle) else None # TODO: Refactor
|
||||
block.put_text(text, rectangle)
|
||||
return block
|
||||
|
||||
|
||||
def paste_content(page, content_box: ContentRectangle):
|
||||
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]):
|
||||
page = deepcopy(page)
|
||||
for content in contents:
|
||||
paste_content(page, content)
|
||||
return page
|
||||
379
test/fixtures/page_generation/page.py
vendored
379
test/fixtures/page_generation/page.py
vendored
@ -1,16 +1,11 @@
|
||||
import itertools
|
||||
import random
|
||||
import sys
|
||||
from copy import deepcopy
|
||||
from enum import Enum
|
||||
from functools import lru_cache, partial
|
||||
from math import sqrt
|
||||
from typing import Tuple, Iterable, List
|
||||
|
||||
import blend_modes
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image, ImageDraw, ImageEnhance
|
||||
from PIL import Image, ImageEnhance
|
||||
from PIL.Image import Transpose
|
||||
from loguru import logger
|
||||
|
||||
@ -18,17 +13,14 @@ from cv_analysis.utils.conversion import normalize_image_format_to_array, normal
|
||||
from cv_analysis.utils.image_operations import blur, sharpen, overlay, superimpose, compute_pasting_coordinates
|
||||
from cv_analysis.utils.merging import merge_related_rectangles
|
||||
from cv_analysis.utils.postprocessing import remove_overlapping, remove_included
|
||||
from cv_analysis.utils.spacial import area
|
||||
from synthesis.partitioner.two_column import TwoColumnPagePartitioner
|
||||
from synthesis.random import rnd, possibly, probably
|
||||
from synthesis.random import rnd, probably
|
||||
from synthesis.segment.content_rectangle import ContentRectangle
|
||||
from synthesis.segment.plot import pick_colormap, RandomPlot
|
||||
from cv_analysis.utils.geometric import is_square_like
|
||||
from synthesis.segment.random_content_rectangle import RandomContentRectangle
|
||||
from synthesis.segment.table.table import generate_random_plot, generate_recursive_random_table, paste_contents
|
||||
from synthesis.segment.text_block import TextBlock
|
||||
from synthesis.text.text_block_generator.caption import CaptionGenerator
|
||||
from synthesis.text.font import pick_random_font_available_on_system
|
||||
from synthesis.text.text import generate_random_words, generate_random_number
|
||||
|
||||
logger.remove()
|
||||
logger.add(sys.stderr, level="INFO")
|
||||
@ -42,8 +34,6 @@ from funcy import (
|
||||
lsplit,
|
||||
lfilter,
|
||||
repeatedly,
|
||||
mapcat,
|
||||
chunks,
|
||||
)
|
||||
|
||||
from cv_analysis.locations import TEST_PAGE_TEXTURES_DIR
|
||||
@ -262,342 +252,6 @@ def split_into_figure_and_caption(rectangle: Rectangle):
|
||||
return figure_box, caption_box
|
||||
|
||||
|
||||
def generate_random_plot(rectangle: Rectangle) -> ContentRectangle:
|
||||
block = RandomPlot(*rectangle.coords)
|
||||
block.content = rectangle.content if isinstance(rectangle, ContentRectangle) else None # TODO: Refactor
|
||||
block.generate_random_plot(rectangle)
|
||||
return block
|
||||
|
||||
|
||||
def generate_recursive_random_table(rectangle: Rectangle, **kwargs) -> ContentRectangle:
|
||||
block = RecursiveRandomTable(*rectangle.coords, **kwargs)
|
||||
if isinstance(rectangle, RecursiveRandomTable):
|
||||
block.content = rectangle.content if rectangle.content else None # TODO: Refactor
|
||||
block.generate_random_table()
|
||||
return block
|
||||
|
||||
|
||||
class Size(Enum):
|
||||
SMALL = 120
|
||||
MEDIUM = 180
|
||||
LARGE = 300
|
||||
|
||||
|
||||
def get_size_class(rectangle: Rectangle):
|
||||
size = get_size(rectangle)
|
||||
if size < Size.SMALL.value:
|
||||
return Size.SMALL
|
||||
elif size < Size.LARGE.value:
|
||||
return Size.MEDIUM
|
||||
else:
|
||||
return Size.LARGE
|
||||
|
||||
|
||||
def get_size(rectangle: Rectangle):
|
||||
size = sqrt(area(rectangle))
|
||||
return size
|
||||
|
||||
|
||||
def get_random_color_complementing_color_map(colormap):
|
||||
def color_complement(r, g, b):
|
||||
"""Reference: https://stackoverflow.com/a/40234924"""
|
||||
|
||||
def hilo(a, b, c):
|
||||
if c < b:
|
||||
b, c = c, b
|
||||
if b < a:
|
||||
a, b = b, a
|
||||
if c < b:
|
||||
b, c = c, b
|
||||
return a + c
|
||||
|
||||
k = hilo(r, g, b)
|
||||
return tuple(k - u for u in (r, g, b))
|
||||
|
||||
color = colormap(0.2)[:3]
|
||||
color = [int(255 * v) for v in color]
|
||||
color = color_complement(*color)
|
||||
return color
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def get_random_background_color():
|
||||
return tuple([*get_random_color_complementing_color_map(pick_colormap()), rnd.randint(100, 210)])
|
||||
|
||||
|
||||
class RecursiveRandomTable(RandomContentRectangle):
|
||||
def __init__(self, x1, y1, x2, y2, border_width=1, layout: str = None, double_rule=False):
|
||||
"""A table with a random number of rows and columns, and random content in each cell.
|
||||
|
||||
Args:
|
||||
x1: x-coordinate of the top-left corner
|
||||
y1: y-coordinate of the top-left corner
|
||||
x2: x-coordinate of the bottom-right corner
|
||||
y2: y-coordinate of the bottom-right corner
|
||||
border_width: width of the table border
|
||||
layout: layout of the table, either "horizontal", "vertical", "closed", or "open"
|
||||
double_rule: whether to use double rules as the top and bottom rules
|
||||
"""
|
||||
|
||||
assert layout in [None, "horizontal", "vertical", "closed", "open"]
|
||||
|
||||
super().__init__(x1, y1, x2, y2)
|
||||
|
||||
self.double_rule = double_rule
|
||||
self.double_rule_width = (3 * border_width) if self.double_rule else 0
|
||||
|
||||
self.n_columns = rnd.randint(1, max(self.width // 100, 1))
|
||||
self.n_rows = rnd.randint(1, max((self.height - 2 * self.double_rule_width) // rnd.randint(17, 100), 1))
|
||||
self.cell_size = (self.width / self.n_columns, (self.height - 2 * self.double_rule_width) / self.n_rows)
|
||||
|
||||
self.content = Image.new("RGBA", (self.width, self.height), (255, 255, 255, 0))
|
||||
|
||||
self.background_color = get_random_background_color()
|
||||
|
||||
logger.info(f"Background color: {self.background_color}")
|
||||
|
||||
self.layout = layout or self.pick_random_layout()
|
||||
logger.debug(f"Layout: {self.layout}")
|
||||
|
||||
def pick_random_layout(self):
|
||||
|
||||
if self.n_columns == 1 and self.n_rows == 1:
|
||||
layout = "closed"
|
||||
elif self.n_columns == 1:
|
||||
layout = rnd.choice(["vertical", "closed"])
|
||||
elif self.n_rows == 1:
|
||||
layout = rnd.choice(["horizontal", "closed"])
|
||||
else:
|
||||
layout = rnd.choice(["closed", "horizontal", "vertical", "open"])
|
||||
|
||||
return layout
|
||||
|
||||
def generate_random_table(self):
|
||||
cells = self.generate_table()
|
||||
cells = list(self.fill_cells_with_content(cells))
|
||||
# FIXME: There is a bug here: Table rule is not drawn correctly, actually we want to do cells = ...
|
||||
list(self.draw_cell_borders(cells))
|
||||
|
||||
self.content = paste_contents(self.content, cells)
|
||||
assert self.content.mode == "RGBA"
|
||||
|
||||
def fill_cells_with_content(self, cells):
|
||||
yield from map(self.build_cell, cells)
|
||||
|
||||
def build_cell(self, cell):
|
||||
|
||||
if self.__is_a_small_cell(cell):
|
||||
cell = self.build_small_cell(cell)
|
||||
|
||||
elif self.__is_a_medium_sized_cell(cell):
|
||||
cell = self.build_medium_sized_cell(cell)
|
||||
|
||||
elif self.__is_a_large_cell(cell):
|
||||
cell = self.build_large_cell(cell)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid cell size: {get_size(cell)}")
|
||||
|
||||
assert cell.content.mode == "RGBA"
|
||||
|
||||
return cell
|
||||
|
||||
def __is_a_small_cell(self, cell):
|
||||
return get_size(cell) <= Size.SMALL.value
|
||||
|
||||
def __is_a_medium_sized_cell(self, cell):
|
||||
return get_size(cell) <= Size.MEDIUM.value
|
||||
|
||||
def __is_a_large_cell(self, cell):
|
||||
return get_size(cell) > Size.MEDIUM.value
|
||||
|
||||
def build_small_cell(self, cell):
|
||||
|
||||
content = (possibly() and generate_random_words(1, 3)) or (
|
||||
generate_random_number()
|
||||
+ ((possibly() and " " + rnd.choice(["$", "£", "%", "EUR", "USD", "CAD", "ADA"])) or "")
|
||||
)
|
||||
|
||||
return generate_text_block(cell, content)
|
||||
|
||||
def build_medium_sized_cell(self, cell):
|
||||
|
||||
choice = rnd.choice(["plot", "recurse"])
|
||||
|
||||
if choice == "plot":
|
||||
return generate_random_plot(cell)
|
||||
|
||||
elif choice == "recurse":
|
||||
return generate_recursive_random_table(
|
||||
cell,
|
||||
border_width=1,
|
||||
layout=random.choice(["open", "horizontal", "vertical"]),
|
||||
double_rule=False,
|
||||
)
|
||||
|
||||
else:
|
||||
return generate_text_block(cell, f"{choice} {get_size(cell):.0f} {get_size_class(cell).name}")
|
||||
|
||||
def build_large_cell(self, cell):
|
||||
choice = rnd.choice(["plot", "recurse"])
|
||||
|
||||
logger.debug(f"Generating {choice} {get_size(cell):.0f} {get_size_class(cell).name}")
|
||||
|
||||
if choice == "plot" and is_square_like(cell):
|
||||
return generate_random_plot(cell)
|
||||
|
||||
else:
|
||||
logger.debug(f"recurse {get_size(cell):.0f} {get_size_class(cell).name}")
|
||||
return generate_recursive_random_table(
|
||||
cell,
|
||||
border_width=1,
|
||||
layout=random.choice(["open", "horizontal", "vertical"]),
|
||||
double_rule=False,
|
||||
)
|
||||
|
||||
def draw_cell_borders(self, cells: List[ContentRectangle]):
|
||||
def draw_edges_based_on_position(cell: Cell, col_idx, row_index):
|
||||
# Draw the borders of the cell based on its position in the table
|
||||
if col_idx < self.n_columns - 1:
|
||||
cell.draw_right_border()
|
||||
|
||||
if row_index < self.n_rows - 1:
|
||||
cell.draw_bottom_border()
|
||||
|
||||
columns = chunks(self.n_rows, cells)
|
||||
for col_idx, column in enumerate(columns):
|
||||
for row_index, cell in enumerate(column):
|
||||
# TODO: Refactor
|
||||
c = Cell(*cell.coords, self.background_color)
|
||||
c.content = cell.content
|
||||
draw_edges_based_on_position(c, col_idx, row_index)
|
||||
yield cell
|
||||
|
||||
if self.layout == "closed":
|
||||
# TODO: Refactor
|
||||
c = Cell(*self.coords, self.background_color)
|
||||
c.content = self.content
|
||||
c.draw()
|
||||
yield self
|
||||
|
||||
# TODO: Refactor
|
||||
if self.double_rule:
|
||||
c1 = Cell(*self.coords)
|
||||
c1.draw_top_border(width=1)
|
||||
c1.draw_bottom_border(width=1)
|
||||
|
||||
x1, y1, x2, y2 = self.coords
|
||||
c2 = Cell(x1, y1 + self.double_rule_width, x2, y2 - self.double_rule_width)
|
||||
c2.draw_top_border(width=1)
|
||||
c2.draw_bottom_border(width=1)
|
||||
|
||||
c = superimpose(c1.content, c2.content)
|
||||
|
||||
self.content = superimpose(c, self.content)
|
||||
|
||||
yield self
|
||||
|
||||
def generate_table(self) -> Iterable[ContentRectangle]:
|
||||
yield from mapcat(self.generate_column, range(self.n_columns))
|
||||
|
||||
def generate_column(self, column_index) -> Iterable[ContentRectangle]:
|
||||
logger.trace(f"Generating column {column_index}.")
|
||||
generate_cell_for_row_index = partial(self.generate_cell, column_index)
|
||||
yield from map(generate_cell_for_row_index, range(self.n_rows))
|
||||
|
||||
def generate_cell(self, column_index, row_index) -> ContentRectangle:
|
||||
w, h = self.cell_size
|
||||
x1, y1 = (column_index * w), (row_index * h) + self.double_rule_width
|
||||
x2, y2 = x1 + w, y1 + h
|
||||
logger.trace(f"Generating cell ({row_index}, {column_index}) at ({x1}, {y1}, {x2}, {y2}).")
|
||||
return Cell(x1, y1, x2, y2, self.background_color)
|
||||
|
||||
def generate_column_names(self):
|
||||
column_names = repeatedly(self.generate_column_name, self.n_columns)
|
||||
return column_names
|
||||
|
||||
def generate_column_name(self):
|
||||
column_name = generate_random_words(1, 3)
|
||||
return column_name
|
||||
|
||||
|
||||
class Cell(ContentRectangle):
|
||||
def __init__(self, x1, y1, x2, y2, color=None):
|
||||
super().__init__(x1, y1, x2, y2)
|
||||
|
||||
self.background_color = color or (255, 255, 255, 0)
|
||||
|
||||
# to debug use random border color: tuple([random.randint(100, 200) for _ in range(3)] + [255])
|
||||
self.cell_border_color = (0, 0, 0, 255)
|
||||
|
||||
self.border_width = 1
|
||||
self.inset = 1
|
||||
|
||||
self.content = Image.new("RGBA", (self.width, self.height))
|
||||
self.fill()
|
||||
|
||||
def draw_top_border(self, width=None):
|
||||
self.draw_line((0, 0, self.width - self.inset, 0), width=width)
|
||||
return self
|
||||
|
||||
def draw_bottom_border(self, width=None):
|
||||
self.draw_line((0, self.height - self.inset, self.width - self.inset, self.height - self.inset), width=width)
|
||||
return self
|
||||
|
||||
def draw_left_border(self, width=None):
|
||||
self.draw_line((0, 0, 0, self.height), width=width)
|
||||
return self
|
||||
|
||||
def draw_right_border(self, width=None):
|
||||
self.draw_line((self.width - self.inset, 0, self.width - self.inset, self.height), width=width)
|
||||
return self
|
||||
|
||||
def draw_line(self, points, width=None):
|
||||
width = width or self.border_width
|
||||
draw = ImageDraw.Draw(self.content)
|
||||
draw.line(points, width=width, fill=self.cell_border_color)
|
||||
return self
|
||||
|
||||
def draw(self, width=None):
|
||||
self.draw_top_border(width=width)
|
||||
self.draw_bottom_border(width=width)
|
||||
self.draw_left_border(width=width)
|
||||
self.draw_right_border(width=width)
|
||||
return self
|
||||
|
||||
def draw_top_left_corner(self, width=None):
|
||||
self.draw_line((0, 0, 0, 0), width=width)
|
||||
self.draw_line((0, 0, 0, 0), width=width)
|
||||
return self
|
||||
|
||||
def draw_top_right_corner(self, width=None):
|
||||
self.draw_line((self.width - self.inset, 0, self.width - self.inset, 0), width=width)
|
||||
self.draw_line((self.width - self.inset, 0, self.width - self.inset, 0), width=width)
|
||||
return self
|
||||
|
||||
def draw_bottom_left_corner(self, width=None):
|
||||
self.draw_line((0, self.height - self.inset, 0, self.height - self.inset), width=width)
|
||||
self.draw_line((0, self.height - self.inset, 0, self.height - self.inset), width=width)
|
||||
return self
|
||||
|
||||
def draw_bottom_right_corner(self, width=None):
|
||||
self.draw_line(
|
||||
(self.width - self.inset, self.height - self.inset, self.width - self.inset, self.height - self.inset),
|
||||
width=width,
|
||||
)
|
||||
self.draw_line(
|
||||
(self.width - self.inset, self.height - self.inset, self.width - self.inset, self.height - self.inset),
|
||||
width=width,
|
||||
)
|
||||
return self
|
||||
|
||||
def fill(self, color=None):
|
||||
color = color or self.background_color
|
||||
image = Image.new("RGBA", (self.width, self.height), color=color)
|
||||
self.content = superimpose(image, self.content)
|
||||
return self
|
||||
|
||||
|
||||
def shrink_rectangle(rectangle: Rectangle, factor: float) -> Rectangle:
|
||||
x1, y1, x2, y2 = compute_scaled_coordinates(rectangle, (1 - factor))
|
||||
|
||||
@ -665,33 +319,6 @@ def generate_random_caption(rectangle: Rectangle, caption_start, n_sentences=100
|
||||
return block
|
||||
|
||||
|
||||
def generate_text_block(rectangle: Rectangle, text) -> ContentRectangle:
|
||||
block = TextBlock(
|
||||
*rectangle.coords,
|
||||
font=pick_random_font_available_on_system(
|
||||
includes=("serif", "sans-serif", "bold"),
|
||||
excludes=("mono", "italic", "oblique", "cursive"),
|
||||
),
|
||||
font_size=30, # TODO: De-hardcode font size... Seems to have no effect on top of that
|
||||
)
|
||||
block.content = rectangle.content if isinstance(rectangle, ContentRectangle) else None # TODO: Refactor
|
||||
block.put_text(text, rectangle)
|
||||
return block
|
||||
|
||||
|
||||
def paste_content(page, content_box: ContentRectangle):
|
||||
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]):
|
||||
page = deepcopy(page)
|
||||
for content in contents:
|
||||
paste_content(page, content)
|
||||
return page
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
params=[
|
||||
TwoColumnPagePartitioner,
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user