diff --git a/synthesis/content_generator.py b/synthesis/content_generator.py index 7ebb0f7..021f07e 100644 --- a/synthesis/content_generator.py +++ b/synthesis/content_generator.py @@ -2,7 +2,7 @@ import itertools from typing import List, Iterable, Union from PIL import Image -from funcy import lsplit, lfilter, mapcat +from funcy import lsplit, lfilter from cv_analysis.logging import logger from cv_analysis.utils import every_nth, zipmap @@ -18,6 +18,7 @@ from synthesis.segment.segments import ( generate_recursive_random_table_with_caption, generate_random_plot_with_caption, ) +from synthesis.segment.table.table import RecursiveRandomTable class ContentGenerator: @@ -49,17 +50,24 @@ class ContentGenerator: boxes = list(unpack_boxes(boxes)) for b in boxes: - logger.info(f"Generated {b}") + logger.trace(f"Generated {b}") return boxes -def unpack_boxes(boxes: Iterable[ContentRectangle]) -> Iterable[ContentRectangle]: - yield from itertools.chain.from_iterable(map(unpack_box, boxes)) +def unpack_boxes(boxes: Iterable[ContentRectangle], depth=0) -> Iterable[ContentRectangle]: + yield from itertools.chain.from_iterable(map(lambda b: unpack_box(b, depth), boxes)) -def unpack_box(box: Union[ContentRectangle, RecursiveContentRectangle]): - print(box) - yield box +def unpack_box(box: Union[ContentRectangle, RecursiveContentRectangle], depth=0) -> Iterable[ContentRectangle]: + # Boxes for recursive tables should be cells or the root table. Interior tables should not be recognized as such, + # but rather as cells. + if isinstance(box, RecursiveRandomTable): + if not depth: + yield box + else: + pass + else: + yield box if box.has_child_boxes(): - yield from unpack_boxes(box.child_boxes) + yield from unpack_boxes(box.child_boxes, depth + 1) diff --git a/synthesis/segment/table/table.py b/synthesis/segment/table/table.py index c1a3a8f..ac47284 100644 --- a/synthesis/segment/table/table.py +++ b/synthesis/segment/table/table.py @@ -12,7 +12,7 @@ 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, shift +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 pick_colormap @@ -66,7 +66,6 @@ class RecursiveRandomTable(RandomContentRectangle, RecursiveContentRectangle): for cell in self.cells: # TODO: this is not very clean cell = deepcopy(cell) - # cell = Cell(*shift(cell, self.x1, self.y1)) cell.shift(self.x1, self.y1) yield cell @@ -86,13 +85,24 @@ class RecursiveRandomTable(RandomContentRectangle, RecursiveContentRectangle): return layout - def generate_random_table(self): + def generate_random_table(self, draw_cell_content=False): + """Generate a random table. The table is generated by first generating a random layout, and then filling the + cells with content recursively. + + Args: + draw_cell_content: Whether to draw the content of each cell. If False, only the table border is drawn. Cells + can be accessed and drawn later. + + Returns: + None + """ 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 = ... cells = list(self.draw_cell_borders(cells)) - self.content = paste_contents(self.content, cells) + # TODO: This is not very clean. + if draw_cell_content: + self.content = paste_contents(self.content, cells) assert self.content.mode == "RGBA" @@ -173,44 +183,55 @@ class RecursiveRandomTable(RandomContentRectangle, RecursiveContentRectangle): ) 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) + self.draw_cell(cell, 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 + self.draw_table_borders() - # TODO: Refactor if self.double_rule: - c1 = Cell(*self.coords) - c1.draw_top_border(width=1) - c1.draw_bottom_border(width=1) + self.draw_table_rule() - 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) + def draw_cell(self, cell, col_idx, row_index): + # TODO: Refactor + c = Cell(*cell.coords, self.background_color) + c.content = cell.content + self.draw_edges_based_on_position(c, col_idx, row_index) - c = superimpose(c1.content, c2.content) + def draw_edges_based_on_position(self, cell: Cell, col_idx, row_index): + """Draw the edges of a cell based on its position in the table.""" + if col_idx < self.n_columns - 1: + cell.draw_right_border() - self.content = superimpose(c, self.content) + if row_index < self.n_rows - 1: + cell.draw_bottom_border() + + def draw_table_rule(self): + # TODO: Refactor + 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) + + def draw_table_borders(self): + # TODO: Refactor + c = Cell(*self.coords, self.background_color) + c.content = self.content + c.draw() + yield self def generate_table(self) -> Iterable[ContentRectangle]: yield from mapcat(self.generate_column, range(self.n_columns))