Refactoring

- Add visitor support to content rectangle
- Refactor box unpacking to happen by applying visitor pattern

synthesis/content_generator.py

@@ -1,5 +1,5 @@
 import itertools
-from typing import List, Iterable, Union
+from typing import List, Iterable
 
 from PIL import Image
 from funcy import lsplit, lfilter
@@ -18,7 +18,6 @@ from synthesis.segment.segments import (
     generate_recursive_random_table_with_caption,
     generate_random_plot_with_caption,
 )
-from synthesis.segment.table.table import RecursiveRandomTable
 
 
 class ContentGenerator:
@@ -55,19 +54,39 @@ class ContentGenerator:
         return boxes
 
 
-def unpack_boxes(boxes: Iterable[ContentRectangle], depth=0) -> Iterable[ContentRectangle]:
-    yield from itertools.chain.from_iterable(map(lambda b: unpack_box(b, depth), boxes))
+class BoxChildrenVisitor:
+    def visit_content_rectangle(self, _box: ContentRectangle):
+        return []
+
+    def visit_recursive_content_rectangle(self, box: RecursiveContentRectangle):
+        return box.children
 
 
-def unpack_box(box: Union[ContentRectangle, RecursiveContentRectangle], depth=0) -> Iterable[ContentRectangle]:
-    # Boxes for recursive tables should be cells or the root table. Interior tables should not be recognized as such,
-    # but rather as cells.
-    if isinstance(box, RecursiveRandomTable):
-        if not depth:
-            yield box
-        else:
-            pass
-    else:
+def unpack_boxes(boxes: Iterable[ContentRectangle], recursed=False) -> Iterable[ContentRectangle]:
+    yield from itertools.chain.from_iterable(map(lambda b: unpack_box(b, recursed), boxes))
+
+
+def unpack_box(box: ContentRectangle, recursed=False) -> Iterable[ContentRectangle]:
+    children = box.accept(BoxChildrenVisitor())
+
+    def is_a_leaf():
+        return not children
+
+    def is_an_internal_node():
+        return recursed and children
+
+    def is_a_root_node():
+        return not recursed and children
+
+    if is_a_root_node():
         yield box
-    if box.has_child_boxes():
-        yield from unpack_boxes(box.child_boxes, depth + 1)
+        yield from unpack_boxes(children, True)
+
+    elif is_an_internal_node():
+        yield from unpack_boxes(children, True)
+
+    elif is_a_leaf():
+        yield box
+
+    else:
+        raise ValueError("This should not happen")

synthesis/segment/content_rectangle.py

@@ -1,5 +1,3 @@
-import abc
-
 from cv_analysis.utils.rectangle import Rectangle
 
 
@@ -11,6 +9,5 @@ class ContentRectangle(Rectangle):
     def __repr__(self):
         return f"{self.__class__.__name__}({self.x1}, {self.y1}, {self.x2}, {self.y2}, content={self.content})"
 
-    @abc.abstractmethod
-    def has_child_boxes(self):
-        raise NotImplementedError()
+    def accept(self, visitor):
+        return visitor.visit_content_rectangle(self)

synthesis/segment/plot.py

@@ -26,9 +26,6 @@ class RandomPlot(RandomContentRectangle):
     def __call__(self, *args, **kwargs):
         pass
 
-    def has_child_boxes(self):
-        return False
-
     def generate_random_plot(self, rectangle: Rectangle):
 
         if is_square_like(rectangle):

synthesis/segment/recursive_content_rectangle.py

@@ -1,5 +1,3 @@
-import abc
-
 from synthesis.segment.content_rectangle import ContentRectangle
 
 
@@ -9,5 +7,8 @@ class RecursiveContentRectangle(ContentRectangle):
         self.content = content
 
     @property
-    def child_boxes(self):
+    def children(self):
         raise NotImplementedError()
+
+    def accept(self, visitor):
+        return visitor.visit_recursive_content_rectangle(self)

synthesis/segment/table/table.py

@@ -62,16 +62,13 @@ class RecursiveRandomTable(RandomContentRectangle, RecursiveContentRectangle):
         return self.__cells
 
     @property
-    def child_boxes(self):
+    def children(self):
         for cell in self.cells:
             # TODO: this is not very clean
             cell = deepcopy(cell)
             cell.shift(self.x1, self.y1)
             yield cell
 
-    def has_child_boxes(self):
-        return True
-
     def pick_random_layout(self):
 
         if self.n_columns == 1 and self.n_rows == 1:

synthesis/segment/text_block.py

@@ -23,9 +23,6 @@ class TextBlock(ContentRectangle):
     def __call__(self, *args, **kwargs):
         pass
 
-    def has_child_boxes(self):
-        return False
-
     @debug_log()
     def generate_random_text(self, rectangle: Rectangle, n_sentences=None):
         lines = self.text_generator(rectangle, n_sentences)