Refactoring

Pull base class for page partioner out of page partioner and add random
page partioner derived class.
This commit is contained in:
Matthias Bisping 2023-01-24 10:13:02 +01:00
parent c4eeb956ca
commit 3772ca021a

View File

@ -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"])