diff --git a/test/fixtures/page_generation/page.py b/test/fixtures/page_generation/page.py index 9c717e3..4dd1198 100644 --- a/test/fixtures/page_generation/page.py +++ b/test/fixtures/page_generation/page.py @@ -12,6 +12,7 @@ from typing import Tuple, Union, Iterable, List import albumentations as A import cv2 as cv +import matplotlib import numpy as np import pandas as pd import pytest @@ -20,6 +21,7 @@ from PIL.Image import Transpose from faker import Faker from loguru import logger from matplotlib import pyplot as plt +from matplotlib.colors import ListedColormap from tabulate import tabulate from cv_analysis.table_parsing import isolate_vertical_and_horizontal_components @@ -33,7 +35,9 @@ from cv_analysis.utils.spacial import area random_seed = random.randint(0, 2**32 - 1) # random_seed = 3896311122 -random_seed = 1986343479 +# random_seed = 1986343479 + +random_seed = 273244862 # empty large table rnd = random.Random(random_seed) logger.info(f"Random seed: {random_seed}") @@ -447,6 +451,14 @@ def is_square_like(box: Rectangle): return box.width / box.height > 0.5 and box.height / box.width > 0.5 +def is_wide(box: Rectangle): + return box.width / box.height > 1.5 + + +def is_tall(box: Rectangle): + return box.height / box.width > 1.5 + + def every_nth(n, iterable): return itertools.islice(iterable, 0, None, n) @@ -541,47 +553,62 @@ def get_size(rectangle: Rectangle): return size +def get_random_color_complementing_color_map(colormap): + def color_complement(r, g, b): + 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, seed=None, border_width=1, layout=None): super().__init__(x1, y1, x2, y2, seed=seed) self.n_columns = rnd.randint(1, max(self.width // 100, 1)) self.n_rows = rnd.randint(1, max(self.height // rnd.randint(17, 100), 1)) - cell_width, cell_height = (self.width / self.n_columns, self.height / self.n_rows) - - width_delta = (self.width - cell_width * self.n_columns) / self.n_columns - height_delta = (self.height - cell_height * self.n_rows) / self.n_rows - - logger.debug(f"width_delta: {width_delta}, height_delta: {height_delta}") - - self.cell_size = cell_width, cell_height - - logger.debug(f"cell size: {self.cell_size}") + self.cell_size = (self.width / self.n_columns, self.height / self.n_rows) self.content = Image.new("RGBA", (self.width, self.height), (255, 255, 255, 0)) - self.background_color = tuple([rnd.randint(100, 200) for _ in range(3)] + [rnd.randint(180, 210)]) - self.cell_border_color = (0, 0, 0, 255) # (*map(lambda x: int(x * 0.8), self.background_color[:3]), 255) + self.background_color = get_random_background_color() - # TODO: Refactor layout selection - # self.layout = rnd.choice(["closed", "horizontal", "vertical", "open"]) - - # # Overwrite the layout choice in some cases - # if self.n_columns == 1 and self.n_rows == 1: - # self.layout = "closed" - # elif self.n_columns == 1: - # self.layout = rnd.choice(["vertical", "closed"]) - # elif self.n_rows == 1: - # self.layout = rnd.choice(["horizontal", "closed"]) - - self.layout = "closed" # TODO: Remove this line - - self.layout = layout or self.layout + logger.info(f"Background color: {self.background_color}") + self.layout = layout or self.pick_random_layout() logger.debug(f"Layout: {self.layout}") - # self.draw_single_cell_borders(self, border_width, fill=(0, 0, 0, 0)) self.cells = None + 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)) @@ -592,12 +619,9 @@ class RecursiveRandomTable(RandomContentRectangle): def fill_cells_with_content(self, cells): for cell in cells: - # self.draw_single_cell_borders(cell, width=1) def inner(cell): - # inner_region = shrink_rectangle(cell, 0.11) - choice = rnd.choice(["text", "plot", "recurse", "plain_table", "blank"]) size = get_size(cell) @@ -609,11 +633,13 @@ class RecursiveRandomTable(RandomContentRectangle): choice = rnd.choice(["plot", "recurse"]) - if choice == "plot": # and is_square_like(cell): + if choice == "plot": return generate_random_plot(cell) elif choice == "recurse": - return generate_recursive_random_table(cell, border_width=1, layout="open") + return generate_recursive_random_table( + cell, border_width=1, layout=random.choice(["open", "horizontal", "vertical"]) + ) else: return generate_text_block(cell, f"{choice} {size:.0f} {get_size_class(cell).name}") @@ -624,26 +650,24 @@ class RecursiveRandomTable(RandomContentRectangle): logger.debug(f"Generating {choice} {size:.0f} {get_size_class(cell).name}") - if choice == "plot" and is_square_like(cell): + if choice == "plot": return generate_random_plot(cell) else: logger.debug(f"recurse {size:.0f} {get_size_class(cell).name}") - return generate_recursive_random_table(cell, border_width=1, layout="open") + return generate_recursive_random_table( + cell, border_width=1, layout=random.choice(["open", "horizontal", "vertical"]) + ) else: return generate_text_block(cell, f"{choice} {size:.0f} {get_size_class(cell).name}") cell = inner(cell) - # self.draw_single_cell_borders(cell, fill=None, width=2) assert cell.content.mode == "RGBA" yield cell def draw_cell_borders(self, cells: List[ContentRectangle]): - # for cell in cells: - # self.draw_single_cell_borders(cell, fill=self.background_color) - 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: @@ -652,8 +676,6 @@ class RecursiveRandomTable(RandomContentRectangle): if row_index < self.n_rows - 1: cell.draw_bottom_border() - # cell.draw() - columns = chunks(self.n_rows, cells) for col_idx, columns in enumerate(columns): for row_index, cell in enumerate(columns): @@ -664,34 +686,12 @@ class RecursiveRandomTable(RandomContentRectangle): yield cell if self.layout == "closed": + # TODO: Refactor c = Cell(*self.coords, self.background_color) c.content = self.content c.draw() yield self - # def draw_single_cell_borders(self, cell: ContentRectangle, width=1, fill=None): - # # fill = (0, 0, 0, 0) if fill is None else fill - # image = cell.content or Image.new("RGBA", (cell.width, cell.height), (255, 255, 255, 0)) - # assert image.mode == "RGBA" - # draw = ImageDraw.Draw(image) - # - # # TODO: Refactor - # if self.layout == "closed": - # draw.rectangle((0, 0, cell.width - 1, cell.height - 1), outline=self.cell_border_color, width=width) - # elif self.layout == "vertical": - # draw.line((0, 0, 0, cell.height - 1), width=width, fill=self.cell_border_color) - # draw.line((cell.width - 1, 0, cell.width - 1, cell.height - 1), width=width, fill=self.cell_border_color) - # elif self.layout == "horizontal": - # draw.line((0, 0, cell.width - 1, 0), width=width, fill=self.cell_border_color) - # draw.line((0, cell.height - 1, cell.width - 1, cell.height - 1), width=width, fill=self.cell_border_color) - # elif self.layout == "open": - # pass - # else: - # raise ValueError(f"Invalid layout '{self.layout}'") - # cell.content = image - # assert cell.content.mode == "RGBA" - # return cell - def generate_table(self) -> Iterable[ContentRectangle]: yield from mapcat(self.generate_column, range(self.n_columns)) @@ -721,9 +721,9 @@ class Cell(ContentRectangle): super().__init__(x1, y1, x2, y2) self.background_color = color or (255, 255, 255, 0) - self.cell_border_color = (*map(lambda x: int(x * 0.8), self.background_color[:3]), 255) - self.cell_border_color = tuple([random.randint(100, 200) for _ in range(3)] + [255]) # FIXME: DEBUG; REMOVE + # to debug use random border color: tuple([random.randint(100, 200) for _ in range(3)] + [255]) + self.cell_border_color = (10, 10, 10, 255) self.border_width = 1 self.inset = 1 @@ -736,7 +736,6 @@ class Cell(ContentRectangle): return self def draw_bottom_border(self, width=None): - # self.draw_line((0, self.height, self.width - self.inset, self.height), width=width) self.draw_line((0, self.height - self.inset, self.width - self.inset, self.height - self.inset), width=width) return self @@ -996,8 +995,8 @@ def pick_random_font_available_on_system(**kwargs): @lru_cache(maxsize=None) -def pick_colormap(): - return rnd.choice( +def pick_colormap() -> ListedColormap: + cmap_name = rnd.choice( [ "viridis", "plasma", @@ -1006,29 +1005,49 @@ def pick_colormap(): "cividis", ], ) + cmap = plt.get_cmap(cmap_name) + return cmap class RandomPlot(RandomContentRectangle): def __init__(self, x1, y1, x2, y2, seed=None): super().__init__(x1, y1, x2, y2, seed=seed) - cmap_name = pick_colormap() - self.cmap = plt.get_cmap(cmap_name) + self.cmap = pick_colormap() def __call__(self, *args, **kwargs): pass def generate_random_plot(self, rectangle: Rectangle): - # noinspection PyArgumentList - rnd.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) + + if is_square_like(rectangle): + plt_fn = rnd.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, + ] + ) + elif is_wide(rectangle): + plt_fn = rnd.choice( + [ + self.generate_random_line_plot, + self.generate_random_scatter_plot, + ] + ) + elif is_tall(rectangle): + plt_fn = rnd.choice( + [ + self.generate_random_bar_plot, + self.generate_random_histogram, + ] + ) + else: + plt_fn = self.generate_random_scatter_plot + + plt_fn(rectangle) def generate_random_bar_plot(self, rectangle: Rectangle): x = sorted(np.random.randint(low=1, high=11, size=5)) @@ -1054,7 +1073,17 @@ class RandomPlot(RandomContentRectangle): def generate_random_pie_chart(self, rectangle: Rectangle): x = np.random.uniform(size=10) - self.__generate_random_plot(plt.pie, rectangle, x, None, plot_kwargs=self.generate_plot_kwargs(keywords=["a"])) + pie_fn = partial( + plt.pie, + labels=[f"Label {i}" for i in range(10)], + autopct="%1.1f%%", + shadow=True, + startangle=90, + pctdistance=0.85, + labeldistance=1.1, + colors=self.cmap(np.linspace(0, 1, 10)), + ) + self.__generate_random_plot(pie_fn, rectangle, x, None, plot_kwargs=self.generate_plot_kwargs(keywords=["a"])) def generate_plot_kwargs(self, keywords=None):