diff --git a/test/fixtures/page_generation/page.py b/test/fixtures/page_generation/page.py index 32d0464..aefd1bf 100644 --- a/test/fixtures/page_generation/page.py +++ b/test/fixtures/page_generation/page.py @@ -30,6 +30,7 @@ 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.conversion import normalize_image_format_to_array, normalize_image_format_to_pil +from cv_analysis.utils.drawing import draw_rectangles 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 @@ -455,6 +456,10 @@ def page_with_opaque_content( boxes = page_partitioner(texture) content_generator = ContentGenerator() boxes = content_generator(boxes) + + boxes.append(generate_random_page_number(texture)) # TODO: Refactor + print(boxes[-1]) + page = paste_contents(texture, boxes) return page, boxes @@ -477,6 +482,8 @@ def page_with_translucent_content( boxes = content_generator(boxes) page_content = paste_contents(blank_page, boxes) + boxes.append(generate_random_page_number(texture)) + texture = random_flip(texture) texture = texture_fn(texture) @@ -746,6 +753,17 @@ def generate_recursive_random_table(rectangle: Rectangle, **kwargs) -> ContentRe return block +# TODO: Refactor +def generate_random_page_number(page: Image) -> ContentRectangle: + + rectangle = Rectangle(0, 0, page.width, page.height) + + block = RandomPageNumber(*rectangle.coords) + block.content = page # TODO: Refactor + block.generate_random_page_number(rectangle) + return block + + @lru_cache(maxsize=None) def get_random_seed(): return rnd.randint(0, 2**32 - 1) @@ -1771,5 +1789,55 @@ 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) + page = draw_rectangles(page, boxes, filled=False, annotate=True) show_image(page, backend="pil") + + +class RandomPageNumber(ContentRectangle): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.page_number = random.randint(1, 1000) + self.margin_distance_percentage = 0.05 + self.margin_distance_x = int(self.width * self.margin_distance_percentage) + self.margin_distance_y = int(self.height * self.margin_distance_percentage) + + self.location_coordinates = self.location_to_coordinates(self.pick_location()) + + def __repr__(self): + return f"PageNumber({self.page_number}, {super().__repr__()})" + + def generate_random_page_number(self, rectangle: Rectangle): + logger.info(f"Generating random page number for {rectangle}") + self.add_page_number() + + def add_page_number(self): + page_number_image = self.get_page_number_image() + self.content.paste(page_number_image, self.location_coordinates, page_number_image) + + def pick_location(self): + return rnd.choice(["top_left", "top_right", "bottom_left", "bottom_right", "center_top", "center_bottom"]) + + def location_to_coordinates(self, locations: str): + if locations == "top_left": + return self.x1 + self.margin_distance_x, self.y1 + self.margin_distance_y + elif locations == "top_right": + return self.x2 - self.margin_distance_x, self.y1 + self.margin_distance_y + elif locations == "bottom_left": + return self.x1 + self.margin_distance_x, self.y2 - self.margin_distance_y + elif locations == "bottom_right": + return self.x2 - self.margin_distance_x, self.y2 - self.margin_distance_y + elif locations == "center_top": + return self.x1 + self.width // 2, self.y1 + self.margin_distance_y + elif locations == "center_bottom": + return self.x1 + self.width // 2, self.y2 - self.margin_distance_y + else: + raise ValueError(f"Unknown location: {locations}") + + def get_page_number_image(self): + font = pick_random_mono_space_font_available_on_system(excludes=("italic", "oblique")) + page_number_image = Image.new("RGBA", (100, 100), (255, 255, 255, 0)) + + draw = ImageDraw.Draw(page_number_image) + draw.text((0, 0), str(self.page_number), font=font, fill=(0, 0, 0, 255)) + + return page_number_image