From 10ea58414397c971df68eb25e97f8be93ac561f0 Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Tue, 24 Jan 2023 15:44:24 +0100 Subject: [PATCH] [WIP] More table / cell edge fiddling and issue fixing --- cv_analysis/utils/conversion.py | 16 +- cv_analysis/utils/display.py | 3 + test/fixtures/page_generation/page.py | 202 ++++++++++++++++++-------- 3 files changed, 159 insertions(+), 62 deletions(-) diff --git a/cv_analysis/utils/conversion.py b/cv_analysis/utils/conversion.py index 8a25abc..c43f2c7 100644 --- a/cv_analysis/utils/conversion.py +++ b/cv_analysis/utils/conversion.py @@ -1,10 +1,14 @@ import json -from typing import Sequence +from typing import Sequence, Union import cv2 +import numpy as np +from PIL import Image from cv_analysis.utils.rectangle import Rectangle +Image_t = Union[Image.Image, np.ndarray] + def contour_to_rectangle(contour): return box_to_rectangle(cv2.boundingRect(contour)) @@ -33,3 +37,13 @@ class RectangleJSONEncoder(json.JSONEncoder): def encode(self, o): result = json.JSONEncoder.encode(self, o) return result + + +def normalize_image_format_to_array(image: Image_t): + return np.array(image) if isinstance(image, Image.Image) else image + + +def normalize_image_format_to_pil(image: Image_t): + if isinstance(image, np.ndarray): + return Image.fromarray(image) + return image diff --git a/cv_analysis/utils/display.py b/cv_analysis/utils/display.py index ea286e3..8489687 100644 --- a/cv_analysis/utils/display.py +++ b/cv_analysis/utils/display.py @@ -4,8 +4,11 @@ from PIL import Image from PIL.Image import Image as Image_t from matplotlib import pyplot as plt +from cv_analysis.utils.conversion import normalize_image_format_to_array + def show_image(image, backend="mpl", **kwargs): + image = normalize_image_format_to_array(image) if backend == "mpl": show_image_mpl(image, **kwargs) elif backend == "cv2": diff --git a/test/fixtures/page_generation/page.py b/test/fixtures/page_generation/page.py index a336f17..102a77d 100644 --- a/test/fixtures/page_generation/page.py +++ b/test/fixtures/page_generation/page.py @@ -25,7 +25,7 @@ from tabulate import tabulate from cv_analysis.table_parsing import isolate_vertical_and_horizontal_components from cv_analysis.utils import star, rconj, conj from cv_analysis.utils.common import normalize_to_gray_scale -from cv_analysis.utils.drawing import draw_rectangles +from cv_analysis.utils.conversion import normalize_image_format_to_array, normalize_image_format_to_pil 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 @@ -35,7 +35,6 @@ random_seed = 3896311122 rnd = random.Random(random_seed) logger.info(f"Random seed: {random_seed}") -Image_t = Union[Image.Image, np.ndarray] # # transform = A.Compose( # [ @@ -114,6 +113,7 @@ from funcy import ( project, complement, lremove, + chunks, ) from cv_analysis.locations import TEST_PAGE_TEXTURES_DIR @@ -164,8 +164,8 @@ Color = Tuple[int, int, int] @pytest.fixture( params=[ - # "rough_grain", - "plain", + "rough_grain", + # "plain", # "digital", # "crumpled", ] @@ -268,16 +268,6 @@ def blur(image: np.ndarray): return cv.blur(image, (3, 3)) -def normalize_image_format_to_array(image: Image_t): - return np.array(image) if isinstance(image, Image.Image) else 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) @@ -526,7 +516,7 @@ class Size(Enum): # MEDIUM = sqrt((100 * 3) ** 2) # LARGE = sqrt((100 * 10) ** 2) - SMALL = 60 + SMALL = 100 MEDIUM = 180 LARGE = 300 @@ -554,9 +544,8 @@ class RecursiveRandomTable(RandomContentRectangle): 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(0, 100) for _ in range(4)]) - self.background_color = tuple([random.randint(0, 100) for _ in range(4)]) - + self.background_color = tuple([rnd.randint(100, 200) for _ in range(4)]) + # self.background_color = tuple([random.randint(0, 100) for _ in range(4)]) self.cell_border_color = (0, 0, 0, 255) # (*map(lambda x: int(x * 0.8), self.background_color[:3]), 255) self.layout = rnd.choice(["closed", "horizontal", "vertical", "open"]) @@ -573,18 +562,23 @@ class RecursiveRandomTable(RandomContentRectangle): logger.debug(f"Layout: {self.layout}") # self.draw_single_cell_borders(self, border_width, fill=(0, 0, 0, 0)) + self.cells = None + def generate_random_table(self): - cells = list(self.generate_cells_with_content()) + cells = self.generate_table() + cells = list(self.fill_cells_with_content(cells)) + self.cells = list(self.draw_cell_borders(cells)) + self.content = paste_contents(self.content, cells) assert self.content.mode == "RGBA" - def generate_cells_with_content(self): - for cell in self.generate_table(): - self.draw_single_cell_borders(cell, width=1) + 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.4) + inner_region = shrink_rectangle(cell, 0.01) choice = rnd.choice(["text", "plot", "recurse", "plain_table", "blank"]) size = get_size(inner_region) @@ -597,14 +591,12 @@ class RecursiveRandomTable(RandomContentRectangle): choice = rnd.choice(["plot", "recurse"]) - # if choice == "plain_table": - # return generate_random_table(cell) - # # cell.content = generate_random_table(inner_region).content - # # return cell if choice == "plot": # and is_square_like(cell): return generate_random_plot(cell) + elif choice == "recurse": return generate_recursive_random_table(cell) + else: return generate_text_block(cell, f"{choice} {size:.0f} {get_size_class(cell).name}") @@ -613,11 +605,13 @@ class RecursiveRandomTable(RandomContentRectangle): choice = rnd.choice(["plot", "recurse"]) logger.debug(f"Generating {choice} {size:.0f} {get_size_class(cell).name}") + if choice == "plot" and is_square_like(cell): return generate_random_plot(cell) + else: logger.debug(f"recurse {size:.0f} {get_size_class(cell).name}") - return generate_recursive_random_table(inner_region, border_width=0) + return generate_recursive_random_table(cell, border_width=0) else: return generate_text_block(cell, f"{choice} {size:.0f} {get_size_class(cell).name}") @@ -629,46 +623,67 @@ class RecursiveRandomTable(RandomContentRectangle): yield cell def draw_cell_borders(self, cells: List[ContentRectangle]): - for cell in cells: - self.draw_single_cell_borders(cell, fill=self.background_color) + # for cell in cells: + # self.draw_single_cell_borders(cell, fill=self.background_color) - 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) + 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 == 0: + cell.draw_left_border() - # 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 + cell.draw_right_border() + + if row_index == 0: + cell.draw_top_border() + + cell.draw_bottom_border() + + columns = chunks(self.n_rows, cells) + for col_idx, columns in enumerate(columns): + for row_index, cell in enumerate(columns): + # 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 + + # 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)) def generate_column(self, column_index) -> Iterable[ContentRectangle]: logger.trace(f"Generating column {column_index}.") - generate_cells_for_row = partial(self.generate_cell, column_index) - yield from map(generate_cells_for_row, range(self.n_rows)) + 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) x2, y2 = x1 + w, y1 + h logger.trace(f"Generating cell ({row_index}, {column_index}) at ({x1}, {y1}, {x2}, {y2}).") - return BlankImageRectangle(x1, y1, x2, y2, self.background_color) + return Cell(x1, y1, x2, y2, self.background_color) def generate_column_names(self): column_names = repeatedly(self.generate_column_name, self.n_columns) @@ -679,9 +694,16 @@ class RecursiveRandomTable(RandomContentRectangle): return column_name -class BlankImageRectangle(ContentRectangle): +class Cell(ContentRectangle): def __init__(self, x1, y1, x2, y2, color): super().__init__(x1, y1, x2, y2) + self.cell_border_color = (0, 0, 0, 255) + + # self.background_color = tuple([random.randint(100, 200) for _ in range(4)]) + # self.cell_border_color = (*map(lambda x: int(x * 0.8), self.background_color[:3]), 255) + + self.border_width = 1 + self.inset = 1 # image = Image.fromarray(np.random.uniform(0, 255, size=(self.height, self.width, 4)).astype(np.uint8)) # self.content = image.convert("RGBA") @@ -695,7 +717,65 @@ class BlankImageRectangle(ContentRectangle): print(color) - self.content = Image.new("RGBA", (self.width, self.height), color=color[:3]) + self.content = Image.new("RGBA", (self.width, self.height), color=color[:4]) + # self.content = Image.new("RGBA", (self.width, self.height), color=color[:3]) + + 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 - self.inset), width=width) + return self + + def draw_right_border(self, width=None): + self.draw_line( + (self.width - self.inset, +self.inset, self.width - self.inset, self.height - self.inset), 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 generate_random_words(n_min, n_max): @@ -708,10 +788,10 @@ def shrink_rectangle(rectangle: Rectangle, factor: float) -> Rectangle: logger.trace(f"Shrinking {rectangle} by {factor} to ({x1}, {y1}, {x2}, {y2}).") - assert x1 > rectangle.x1 - assert y1 > rectangle.y1 - assert x2 < rectangle.x2 - assert y2 < rectangle.y2 + assert x1 >= rectangle.x1 + assert y1 >= rectangle.y1 + assert x2 <= rectangle.x2 + assert y2 <= rectangle.y2 shrunk_rectangle = Rectangle(x1, y1, x2, y2) @@ -1330,4 +1410,4 @@ def drop_small_boxes(boxes: Iterable[Rectangle], page_width, page_height, min_pe def draw_boxes(page: Image, boxes: Iterable[Rectangle]): # page = draw_rectangles(page, boxes, filled=False, annotate=True) - show_image(page) + show_image(page, backend="pil")