diff --git a/test/fixtures/page_generation/page.py b/test/fixtures/page_generation/page.py index c8407f9..3a15d3d 100644 --- a/test/fixtures/page_generation/page.py +++ b/test/fixtures/page_generation/page.py @@ -380,7 +380,7 @@ def blank_page(texture, texture_fn) -> np.ndarray: """Creates a blank page with a given orientation and dpi.""" page = random_flip(texture) page = texture_fn(page) - page_partitioner = PagePartitioner() + page_partitioner = RandomPagePartitioner() boxes = page_partitioner(page) content_generator = ContentGenerator() boxes = content_generator(boxes) @@ -1170,7 +1170,7 @@ def paste_contents(page, contents: Iterable[ContentRectangle]): return page -class PagePartitioner: +class PagePartitioner(abc.ABC): def __init__(self): self.left_margin_percentage = 0.05 self.right_margin_percentage = 0.05 @@ -1195,6 +1195,21 @@ class PagePartitioner: # boxes = list(boxes) return boxes + @abc.abstractmethod + def generate_content_boxes(self, box: Rectangle, depth=0): + raise NotImplementedError + + def recurse(self, depth): + return random.random() <= self.recursion_probability(depth) + + def recursion_probability(self, depth): + return self.initial_recursion_probability * (1 - self.recursion_probability_decay) ** depth + + +class RandomPagePartitioner(PagePartitioner): + def __init__(self): + super().__init__() + def generate_content_boxes(self, box: Rectangle, depth=0): if depth >= self.max_recursion_depth: yield box @@ -1205,12 +1220,6 @@ class PagePartitioner: else: yield child_boxes - def recurse(self, depth): - return random.random() <= self.recursion_probability(depth) - - def recursion_probability(self, depth): - return self.initial_recursion_probability * (1 - self.recursion_probability_decay) ** depth - def generate_random_child_boxes(box: Rectangle, margin_percentage) -> Tuple[Rectangle, Rectangle]: axis = random.choice(["x", "y"])