Improve Layout (#24)

* Enforce keyword arguments for page_data

* Add sort func for layout, suggested by #23

Author: lolipopshock

src/layoutparser/elements.py

@@ -274,13 +274,13 @@ def is_in(self, other, soft_margin={}, center=False):
 
     @abstractmethod
     def intersect(self, other: "BaseCoordElement", strict: bool = True):
-        """Intersect the current shape with the other object, with operations defined in 
+        """Intersect the current shape with the other object, with operations defined in
         :doc:`../notes/shape_operations`.
         """
 
     @abstractmethod
     def union(self, other: "BaseCoordElement", strict: bool = True):
-        """Union the current shape with the other object, with operations defined in 
+        """Union the current shape with the other object, with operations defined in
         :doc:`../notes/shape_operations`.
         """
 
@@ -1727,18 +1727,19 @@ class Layout(MutableSequence):
             A list of layout element blocks
         page_data (Dict, optional):
             A dictionary storing the page (canvas) related information
-            like `height`, `width`, etc.
+            like `height`, `width`, etc. It should be passed in as a
+            keyword argument to avoid any confusion.
             Defaults to None.
     """
 
-    def __init__(self, blocks: Optional[List] = None, page_data: Dict = None):
+    def __init__(self, blocks: Optional[List] = None, *, page_data: Dict = None):
         self._blocks = blocks if blocks is not None else []
         self.page_data = page_data or {}
 
     def __getitem__(self, key):
         blocks = self._blocks[key]
         if isinstance(key, slice):
-            return self.__class__(self._blocks[key], self.page_data)
+            return self.__class__(self._blocks[key], page_data=self.page_data)
         else:
             return blocks
 
@@ -1771,17 +1772,20 @@ def __eq__(self, other):
     def __add__(self, other):
         if isinstance(other, Layout):
             if self.page_data == other.page_data:
-                return self.__class__(self._blocks + other._blocks, self.page_data)
+                return self.__class__(
+                    self._blocks + other._blocks, page_data=self.page_data
+                )
             elif self.page_data == {} or other.page_data == {}:
                 return self.__class__(
-                    self._blocks + other._blocks, self.page_data or other.page_data
+                    self._blocks + other._blocks,
+                    page_data=self.page_data or other.page_data,
                 )
             else:
                 raise ValueError(
                     f"Incompatible page_data for two innputs: {self.page_data} vs {other.page_data}."
                 )
         elif isinstance(other, list):
-            return self.__class__(self._blocks + other, self.page_data)
+            return self.__class__(self._blocks + other, page_data=self.page_data)
         else:
             raise ValueError(
                 f"Invalid input type for other {other.__class__.__name__}."
@@ -1791,19 +1795,50 @@ def insert(self, key, value):
         self._blocks.insert(key, value)
 
     def copy(self):
-        return self.__class__(copy(self._blocks), self.page_data)
+        return self.__class__(copy(self._blocks), page_data=self.page_data)
 
     def relative_to(self, other):
-        return self.__class__([ele.relative_to(other) for ele in self], self.page_data)
+        return self.__class__(
+            [ele.relative_to(other) for ele in self], page_data=self.page_data
+        )
 
     def condition_on(self, other):
-        return self.__class__([ele.condition_on(other) for ele in self], self.page_data)
+        return self.__class__(
+            [ele.condition_on(other) for ele in self], page_data=self.page_data
+        )
 
     def is_in(self, other, soft_margin={}, center=False):
         return self.__class__(
-            [ele.is_in(other, soft_margin, center) for ele in self], self.page_data
+            [ele.is_in(other, soft_margin, center) for ele in self],
+            page_data=self.page_data,
         )
 
+    def sort(self, key=None, reverse=False, inplace=False) -> Optional["Layout"]:
+        """Sort the list of blocks based on the given 
+
+        Args:
+            key ([type], optional): key specifies a function of one argument that 
+            is used to extract a comparison key from each list element. 
+            Defaults to None.
+            reverse (bool, optional): reverse is a boolean value. If set to True, 
+            then the list elements are sorted as if each comparison were reversed. 
+            Defaults to False.
+            inplace (bool, optional): whether to perform the sort inplace. If set 
+            to False, it will return another object instance with _block sorted in
+            the order. Defaults to False.
+
+        Examples::
+            >>> import layoutparser as lp
+            >>> i = lp.Interval(4, 5, axis="y")
+            >>> l = lp.Layout([i, i.shift(2)])
+            >>> l.sort(key=lambda x: x.coordinates[1], reverse=True)
+
+        """
+        if not inplace:
+            return self.__class__(sorted(self._blocks, key=key, reverse=reverse), page_data=self.page_data)
+        else:
+            self._blocks.sort(key=key, reverse=reverse)
+
     def filter_by(self, other, soft_margin={}, center=False):
         """
         Return a `Layout` object containing the elements that are in the `other` object.
@@ -1818,7 +1853,7 @@ def filter_by(self, other, soft_margin={}, center=False):
         """
         return self.__class__(
             [ele for ele in self if ele.is_in(other, soft_margin, center)],
-            self.page_data,
+            page_data=self.page_data,
         )
 
     def shift(self, shift_distance):
@@ -1835,7 +1870,7 @@ def shift(self, shift_distance):
                 A new layout object with all the elements shifted in the specified values.
         """
         return self.__class__(
-            [ele.shift(shift_distance) for ele in self], self.page_data
+            [ele.shift(shift_distance) for ele in self], page_data=self.page_data
         )
 
     def pad(self, left=0, right=0, top=0, bottom=0, safe_mode=True):
@@ -1856,7 +1891,7 @@ def pad(self, left=0, right=0, top=0, bottom=0, safe_mode=True):
         """
         return self.__class__(
             [ele.pad(left, right, top, bottom, safe_mode) for ele in self],
-            self.page_data,
+            page_data=self.page_data,
         )
 
     def scale(self, scale_factor):
@@ -1871,7 +1906,9 @@ def scale(self, scale_factor):
             :obj:`Layout`:
                 A new layout object with all the elements scaled in the specified values.
         """
-        return self.__class__([ele.scale(scale_factor) for ele in self], self.page_data)
+        return self.__class__(
+            [ele.scale(scale_factor) for ele in self], page_data=self.page_data
+        )
 
     def crop_image(self, image):
         return [ele.crop_image(image) for ele in self]

tests/test_elements.py

@@ -2,7 +2,15 @@
 import numpy as np
 import pandas as pd
 
-from layoutparser.elements import Interval, Rectangle, Quadrilateral, TextBlock, Layout, InvalidShapeError, NotSupportedShapeError
+from layoutparser.elements import (
+    Interval,
+    Rectangle,
+    Quadrilateral,
+    TextBlock,
+    Layout,
+    InvalidShapeError,
+    NotSupportedShapeError,
+)
 
 
 def test_interval():
@@ -46,7 +54,7 @@ def test_quadrilateral():
 
     points = np.array([[2, 2], [6, 2], [6, 7], [2, 6]])
     q = Quadrilateral(points)
-    q.to_interval(axis='x')
+    q.to_interval(axis="x")
     q.to_rectangle()
     assert q.shift(1) == Quadrilateral(points + 1)
     assert q.shift([1, 2]) == Quadrilateral(points + np.array([1, 2]))
@@ -70,7 +78,7 @@ def test_quadrilateral():
 
     q = Quadrilateral([1, 2, 3, 4, 5, 6, 7, 8])
     assert (q.points == np.array([[1, 2], [3, 4], [5, 6], [7, 8]])).all()
-    
+
     q = Quadrilateral([[1, 2], [3, 4], [5, 6], [7, 8]])
     assert (q.points == np.array([[1, 2], [3, 4], [5, 6], [7, 8]])).all()
 
@@ -173,25 +181,48 @@ def test_textblock():
 
     t = TextBlock(q, score=0.2)
 
-    # Additional test for shape conversion 
-    assert TextBlock(i, id=1, type=2, text="12").to_interval() == TextBlock(i, id=1, type=2, text="12")
-    assert TextBlock(i, id=1, type=2, text="12").to_rectangle() == TextBlock(i.to_rectangle(), id=1, type=2, text="12")
-    assert TextBlock(i, id=1, type=2, text="12").to_quadrilateral() == TextBlock(i.to_quadrilateral(), id=1, type=2, text="12")
-    
-    assert TextBlock(r, id=1, type=2, parent="a").to_interval(axis="x") == TextBlock(r.to_interval(axis="x"), id=1, type=2, parent="a")
-    assert TextBlock(r, id=1, type=2, parent="a").to_interval(axis="y") == TextBlock(r.to_interval(axis="y"), id=1, type=2, parent="a")
-    assert TextBlock(r, id=1, type=2, parent="a").to_rectangle() == TextBlock(r, id=1, type=2, parent="a")
-    assert TextBlock(r, id=1, type=2, parent="a").to_quadrilateral() == TextBlock(r.to_quadrilateral(), id=1, type=2, parent="a")
-    
-    assert TextBlock(q, id=1, type=2, parent="a").to_interval(axis="x") == TextBlock(q.to_interval(axis="x"), id=1, type=2, parent="a")
-    assert TextBlock(q, id=1, type=2, parent="a").to_interval(axis="y") == TextBlock(q.to_interval(axis="y"), id=1, type=2, parent="a")
-    assert TextBlock(q, id=1, type=2, parent="a").to_rectangle() == TextBlock(q.to_rectangle(), id=1, type=2, parent="a")
-    assert TextBlock(q, id=1, type=2, parent="a").to_quadrilateral() == TextBlock(q, id=1, type=2, parent="a")
+    # Additional test for shape conversion
+    assert TextBlock(i, id=1, type=2, text="12").to_interval() == TextBlock(
+        i, id=1, type=2, text="12"
+    )
+    assert TextBlock(i, id=1, type=2, text="12").to_rectangle() == TextBlock(
+        i.to_rectangle(), id=1, type=2, text="12"
+    )
+    assert TextBlock(i, id=1, type=2, text="12").to_quadrilateral() == TextBlock(
+        i.to_quadrilateral(), id=1, type=2, text="12"
+    )
+
+    assert TextBlock(r, id=1, type=2, parent="a").to_interval(axis="x") == TextBlock(
+        r.to_interval(axis="x"), id=1, type=2, parent="a"
+    )
+    assert TextBlock(r, id=1, type=2, parent="a").to_interval(axis="y") == TextBlock(
+        r.to_interval(axis="y"), id=1, type=2, parent="a"
+    )
+    assert TextBlock(r, id=1, type=2, parent="a").to_rectangle() == TextBlock(
+        r, id=1, type=2, parent="a"
+    )
+    assert TextBlock(r, id=1, type=2, parent="a").to_quadrilateral() == TextBlock(
+        r.to_quadrilateral(), id=1, type=2, parent="a"
+    )
+
+    assert TextBlock(q, id=1, type=2, parent="a").to_interval(axis="x") == TextBlock(
+        q.to_interval(axis="x"), id=1, type=2, parent="a"
+    )
+    assert TextBlock(q, id=1, type=2, parent="a").to_interval(axis="y") == TextBlock(
+        q.to_interval(axis="y"), id=1, type=2, parent="a"
+    )
+    assert TextBlock(q, id=1, type=2, parent="a").to_rectangle() == TextBlock(
+        q.to_rectangle(), id=1, type=2, parent="a"
+    )
+    assert TextBlock(q, id=1, type=2, parent="a").to_quadrilateral() == TextBlock(
+        q, id=1, type=2, parent="a"
+    )
 
     with pytest.raises(ValueError):
         TextBlock(q, id=1, type=2, parent="a").to_interval()
         TextBlock(r, id=1, type=2, parent="a").to_interval()
 
+
 def test_layout():
     i = Interval(4, 5, axis="y")
     q = Quadrilateral(np.array([[2, 2], [6, 2], [6, 7], [2, 5]]))
@@ -241,6 +272,20 @@ def test_layout():
         l.page_data = {"width": 200, "height": 400}
         l + l2
 
+    # Test sort
+    l = Layout([i, i.shift(2)])
+    l.sort(key=lambda x: x.coordinates[1], reverse=True)
+    assert l == Layout([i.shift(2), i])
+
+    l = Layout([q, r, i], page_data={"width": 200, "height": 400})
+    assert l.sort(key=lambda x: x.coordinates[0], inplace=False) == Layout(
+        [i, q, r], page_data={"width": 200, "height": 400}
+    )
+
+    l = Layout([q, t])
+    assert l.sort(key=lambda x: x.coordinates[0], inplace=False) == Layout([q, t])
+
+
 def test_shape_operations():
     i_1 = Interval(1, 2, axis="y", canvas_height=30, canvas_width=400)
     i_2 = TextBlock(Interval(1, 2, axis="x"))
@@ -249,19 +294,19 @@ def test_shape_operations():
     r_1 = Rectangle(0.5, 0.5, 2.5, 1.5)
     r_2 = TextBlock(Rectangle(0.5, 0.5, 2, 2.5))
 
-    q_1 = Quadrilateral([[1,1], [2.5, 1.2], [2.5, 3], [1.5, 3]])
-    q_2 = TextBlock(Quadrilateral([[0.5, 0.5], [2,1], [1.5, 2.5], [0.5, 2]]))
+    q_1 = Quadrilateral([[1, 1], [2.5, 1.2], [2.5, 3], [1.5, 3]])
+    q_2 = TextBlock(Quadrilateral([[0.5, 0.5], [2, 1], [1.5, 2.5], [0.5, 2]]))
 
-    # I and I in different axes 
+    # I and I in different axes
     assert i_1.intersect(i_1) == i_1
-    assert i_1.intersect(i_2) == Rectangle(1,1,2,2)
-    assert i_1.intersect(i_3) == i_1 # Ensure intersect copy the canvas size 
+    assert i_1.intersect(i_2) == Rectangle(1, 1, 2, 2)
+    assert i_1.intersect(i_3) == i_1  # Ensure intersect copy the canvas size
 
     assert i_1.union(i_1) == i_1
     with pytest.raises(InvalidShapeError):
-        assert i_1.union(i_2) == Rectangle(1,1,2,2)
+        assert i_1.union(i_2) == Rectangle(1, 1, 2, 2)
 
-    # I and R in different axes 
+    # I and R in different axes
     assert i_1.intersect(r_1) == Rectangle(0.5, 1, 2.5, 1.5)
     assert i_2.intersect(r_1).block == Rectangle(1, 0.5, 2, 1.5)
     assert i_1.union(r_1) == Rectangle(0.5, 0.5, 2.5, 2)
@@ -271,12 +316,12 @@ def test_shape_operations():
     with pytest.raises(NotSupportedShapeError):
         i_1.intersect(q_1)
         i_1.union(q_1)
-        
+
     # I and Q in different axes
-    assert i_1.intersect(q_1, strict=False) == Rectangle(1,1,2.5,2)
-    assert i_1.union(q_1, strict=False)  == Rectangle(1,1,2.5,3)   
+    assert i_1.intersect(q_1, strict=False) == Rectangle(1, 1, 2.5, 2)
+    assert i_1.union(q_1, strict=False) == Rectangle(1, 1, 2.5, 3)
     assert i_2.intersect(q_1, strict=False).block == Rectangle(1, 1, 2, 3)
-    assert i_2.union(q_1, strict=False).block  == Rectangle(1,1,2.5,3)
+    assert i_2.union(q_1, strict=False).block == Rectangle(1, 1, 2.5, 3)
 
     # R and I
     assert r_1.intersect(i_1) == i_1.intersect(r_1)
@@ -289,7 +334,7 @@ def test_shape_operations():
     with pytest.raises(NotSupportedShapeError):
         r_1.intersect(q_1)
         r_1.union(q_1)
-        
+
     assert r_1.intersect(q_1, strict=False) == Rectangle(1, 1, 2.5, 1.5)
     assert r_1.union(q_1, strict=False) == Rectangle(0.5, 0.5, 2.5, 3)
     assert r_1.intersect(q_2, strict=False) == r_1.intersect(q_2.to_rectangle())
@@ -300,7 +345,7 @@ def test_shape_operations():
         q_1.intersect(i_1)
         q_1.intersect(r_1)
         q_1.intersect(q_2)
-        
+
     # Q and I
     assert q_1.intersect(i_1, strict=False) == i_1.intersect(q_1, strict=False)
     assert q_1.union(i_1, strict=False) == i_1.union(q_1, strict=False)
@@ -315,6 +360,7 @@ def test_shape_operations():
     assert q_1.union(q_2, strict=False) == q_2.union(q_1, strict=False).block
     assert q_1.union(q_2, strict=False) == Rectangle(0.5, 0.5, 2.5, 3)
 
+
 def test_dict():
 
     i = Interval(1, 2, "y", canvas_height=5)