Select cell content conditioned on cell size class

test/fixtures/page_generation/page.py

@@ -3,7 +3,9 @@ import itertools
 import random
 import string
 import textwrap
+from enum import Enum
 from functools import lru_cache, partial
+from math import sqrt
 from pathlib import Path
 from typing import Tuple, Union, Iterable, List
 
@@ -24,6 +26,7 @@ from cv_analysis.utils import star, rconj
 from cv_analysis.utils.common import normalize_to_gray_scale
 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
 
 Image_t = Union[Image.Image, np.ndarray]
 #
@@ -438,8 +441,8 @@ def generate_random_table(rectangle: Rectangle) -> ContentRectangle:
     return block
 
 
-def generate_recursive_random_table(rectangle: Rectangle) -> ContentRectangle:
-    block = RecursiveRandomTable(*rectangle.coords)
+def generate_recursive_random_table(rectangle: Rectangle, **kwargs) -> ContentRectangle:
+    block = RecursiveRandomTable(*rectangle.coords, **kwargs)
     block.generate_random_table()
     return block
 
@@ -456,19 +459,35 @@ class RandomContentRectangle(ContentRectangle):
         self.random = random.Random(self.seed)
 
 
+class Size(Enum):
+    SMALL = 100
+    MEDIUM = 500
+    LARGE = 2500
+
+
+def get_size_class(rectangle: Rectangle):
+    size = sqrt(area(rectangle))
+    if size < Size.SMALL.value:
+        return Size.SMALL
+    elif size < Size.LARGE.value:
+        return Size.MEDIUM
+    else:
+        return Size.LARGE
+
+
 class RecursiveRandomTable(RandomContentRectangle):
-    def __init__(self, x1, y1, x2, y2, seed=None):
+    def __init__(self, x1, y1, x2, y2, seed=None, border_width=3):
         super().__init__(x1, y1, x2, y2, seed=seed)
-        self.n_columns = max(self.width // 100, 1)
-        self.n_rows = max(self.height // random.randint(17, 100), 1)
+        self.n_columns = random.randint(1, max(self.width // 100, 1))
+        self.n_rows = random.randint(1, max(self.height // random.randint(17, 100), 1))
         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, 255))
-        self.background_color = tuple([random.randint(100, 200) for _ in range(4)])
+        self.background_color = tuple([random.randint(0, 100) for _ in range(4)])
         self.cell_border_color = (*map(lambda x: int(x * 0.8), self.background_color[:3]), 255)
+        self.draw_single_cell_borders(self, border_width)
 
     def generate_random_table(self):
-        self.draw_single_cell_borders(self, 3)
         cells = list(self.generate_cells_with_content())
         self.content = paste_contents(self.content, cells)
         assert self.content.mode == "RGBA"
@@ -476,18 +495,21 @@ class RecursiveRandomTable(RandomContentRectangle):
     def generate_cells_with_content(self):
         for cell in self.generate_table():
             self.draw_single_cell_borders(cell)
-            inner_region = shrink_rectangle(cell, 0.1)
+            inner_region = cell  # shrink_rectangle(cell, 0.1)
 
             choice = random.choice(["text", "plot", "recurse", "plain_table"])
+            size_class = get_size_class(inner_region)
 
-            if choice == "text":
+            if choice == "text" and size_class is Size.SMALL:
                 cell.content = generate_random_text_block(inner_region, n_sentences=1).content
-            elif choice == "plot":
+            elif choice == "plot" and size_class in (Size.MEDIUM, Size.LARGE):
                 cell.content = generate_random_plot(inner_region).content
-            elif choice == "recurse":
-                cell.content = generate_recursive_random_table(inner_region).content
-            elif choice == "plain_table":
+            elif choice == "plain_table" and size_class in (Size.MEDIUM, Size.LARGE):
                 cell.content = generate_random_table(inner_region).content
+            elif choice == "recurse" and size_class is Size.LARGE:
+                cell.content = generate_recursive_random_table(inner_region, border_width=0).content
+            else:  # Leave blank
+                pass
 
             assert cell.content.mode == "RGBA"