fix: correct column ordering for multi-column documents (#112)

* first pass on ordering algorithm

* doc strings and type hints

* added test for multicolumn layouts

* version and changelog

* handling for empty layout

* update version

Author: MthwRobinson

CHANGELOG.md

@@ -3,6 +3,7 @@
 * Preserve image format in PIL.Image.Image when loading
 * Added ONNX version of Detectron2 and make default model
 * Remove API code, we don't serve this as a standalone API any more
+* Update ordering logic to account for multicolumn documents.
 
 ## 0.4.4
 

sample-docs/design-thinking.pdf

@@ -458,3 +458,18 @@ def test_load_pdf_image_placement():
     # Image is in top half of the page, so that should be reflected in the pixel coordinates
     assert image_region.y1 < images[5].height / 2
     assert image_region.y2 < images[5].height / 2
+
+
+def test_load_pdf_with_multicolumn_layout_and_ocr(filename="sample-docs/design-thinking.pdf"):
+    layouts, images = layout.load_pdf(filename)
+    doc = layout.process_file_with_model(filename=filename, model_name=None)
+    test_snippets = ["Key to design thinking", "Design thinking also", "But in recent years"]
+
+    test_elements = []
+    for element in doc.pages[0].elements:
+        for snippet in test_snippets:
+            if element.text.startswith(snippet):
+                test_elements.append(element)
+
+    for i, element in enumerate(test_elements):
+        assert element.text.startswith(test_snippets[i])

test_unstructured_inference/inference/test_layout.py

@@ -415,7 +415,7 @@ def test_table_prediction(model_path, sample_table_transcript, platform_type):
         table_model.initialize(model=model_path)
         img = Image.open("./sample-docs/example_table.jpg").convert("RGB")
         prediction = table_model.predict(img)
-        assert prediction == sample_table_transcript
+        assert prediction.strip() == sample_table_transcript.strip()
 
 
 def test_intersect():

test_unstructured_inference/models/test_tables.py

@@ -60,6 +60,16 @@ def height(self) -> Union[int, float]:
         """Height of rectangle"""
         return self.y2 - self.y1
 
+    @property
+    def x_midpoint(self) -> Union[int, float]:
+        """Finds the horizontal midpoint of the object."""
+        return (self.x2 + self.x1) / 2
+
+    @property
+    def y_midpoint(self) -> Union[int, float]:
+        """Finds the vertical midpoint of the object."""
+        return (self.y2 + self.y1) / 2
+
     def is_disjoint(self, other: Rectangle) -> bool:
         """Checks whether this rectangle is disjoint from another rectangle."""
         return not self.intersects(other)

unstructured_inference/inference/elements.py

@@ -14,6 +14,7 @@
     ImageTextRegion,
 )
 from unstructured_inference.inference.layoutelement import LayoutElement
+from unstructured_inference.inference.ordering import order_layout
 from unstructured_inference.logger import logger
 from unstructured_inference.models.base import get_model
 from unstructured_inference.models.unstructuredmodel import UnstructuredModel
@@ -170,9 +171,7 @@ def get_elements_with_model(self, inplace=True) -> Optional[List[LayoutElement]]
     def get_elements_from_layout(self, layout: List[TextRegion]) -> List[LayoutElement]:
         """Uses the given Layout to separate the page text into elements, either extracting the
         text from the discovered layout blocks or from the image using OCR."""
-        # NOTE(robinson) - This orders the page from top to bottom. We'll need more
-        # sophisticated ordering logic for more complicated layouts.
-        layout.sort(key=lambda element: element.y1)
+        layout = order_layout(layout)
         elements = [
             get_element_from_block(
                 block=e,

unstructured_inference/inference/layout.py

@@ -0,0 +1,99 @@
+from typing import List, Union
+
+from unstructured_inference.inference.elements import TextRegion
+
+
+class Column:
+    """Class to capture a column of text in the layout. Will update the midpoint of the
+    column as layout elements are added to help with new element comparisons."""
+
+    def __init__(self, layout_elements: List[TextRegion] = []):
+        self.layout_elements = layout_elements
+
+        num_elements = len(layout_elements)
+        if num_elements > 0:
+            self.x_midpoint = sum([el.x_midpoint for el in layout_elements]) / num_elements
+        else:
+            self.x_midpoint = 0
+
+    def add_element(self, layout_element: TextRegion):
+        """Adds an elements to the column and updates the midpoint."""
+        self.layout_elements.append(layout_element)
+        num_elements = len(self.layout_elements)
+        self.x_midpoint = sum([el.x_midpoint for el in self.layout_elements]) / num_elements
+
+
+def order_layout(
+    layout: List[TextRegion],
+    column_tol_factor: float = 0.2,
+    full_page_threshold_factor: float = 0.9,
+) -> List[TextRegion]:
+    """Orders the layout elements detected on a page. For groups of elements that are not
+    the width of the page, the algorithm attempts to group elements into column based on
+    the coordinates of the bounding box. Columns are ordered left to right, and elements
+    within columns are ordered top to bottom.
+
+    Parameters
+    ----------
+    layout
+        the layout elements to order.
+    column_tol_factor
+        multiplied by the page width to find the tolerance for considering two elements as
+        part of the same column.
+    full_page_threshold_factor
+        multiplied by the page width to find the minimum width an elements need to be
+        for it to be considered a full page width element.
+    """
+    if len(layout) == 0:
+        return []
+
+    width = calculate_width(layout)
+    column_tolerance = column_tol_factor * width
+    full_page_min_width = full_page_threshold_factor * width
+
+    layout.sort(key=lambda element: element.y1)
+
+    sorted_layout = []
+    columns: List[Column] = []
+    for layout_element in layout:
+        if layout_element.width > full_page_min_width:
+            sorted_layout.extend(sorted_layout_from_columns(columns))
+            columns = []
+            sorted_layout.append(layout_element)
+
+        else:
+            added_to_column = False
+            for column in columns:
+                difference = abs(layout_element.x_midpoint - column.x_midpoint)
+                if difference < column_tolerance:
+                    column.add_element(layout_element)
+                    added_to_column = True
+                    break
+
+            if not added_to_column:
+                columns.append(Column(layout_elements=[layout_element]))
+
+    sorted_layout.extend(sorted_layout_from_columns(columns))
+    return sorted_layout
+
+
+def sorted_layout_from_columns(columns: List[Column]) -> List[TextRegion]:
+    """Creates a sorted list of elements from a list of columns. Columns will be sorted
+    left to right and elements within columns are sorted top to bottom."""
+    sorted_layout = []
+    if len(columns) > 0:
+        columns.sort(key=lambda column: column.x_midpoint)
+        for column in columns:
+            column.layout_elements.sort(key=lambda element: element.y1)
+            for layout_element in column.layout_elements:
+                sorted_layout.append(layout_element)
+    return sorted_layout
+
+
+def calculate_width(layout) -> Union[float, int]:
+    """Calculates total width of the elements in the layout. Used for computing the full
+    page threshold and column tolerance."""
+    min_x1 = min([element.x1 for element in layout])
+    max_x2 = max([element.x2 for element in layout])
+
+    return max_x2 - min_x1