[WIP] More table / cell edge fiddling and issue fixing
This commit is contained in:
parent
7676a8148e
commit
10ea584143
@ -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
|
||||
|
||||
@ -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":
|
||||
|
||||
202
test/fixtures/page_generation/page.py
vendored
202
test/fixtures/page_generation/page.py
vendored
@ -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")
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user