enhancement: combine layouts (#123)

Add logic to combine inferred elements with extracted elements.

Author: qued

CHANGELOG.md

@@ -1,5 +1,6 @@
-## 0.5.2-dev0
+## 0.5.2-dev1
 
+* Combine inferred elements with extracted elements
 * Add ruff to keep code consistent with unstructured
 
 ## 0.5.1

setup.py

@@ -17,14 +17,15 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-from typing import List
+from typing import List, Optional, Union
 
 from setuptools import find_packages, setup
 
 from unstructured_inference.__version__ import __version__
 
 
-def load_requirements(file_list=None):
+def load_requirements(file_list: Optional[Union[str, List[str]]] = None):
+    """Loads the requirements from a .in file or list of .in files."""
     if file_list is None:
         file_list = ["requirements/base.in"]
     if isinstance(file_list, str):
@@ -33,7 +34,9 @@ def load_requirements(file_list=None):
     for file in file_list:
         with open(file, encoding="utf-8") as f:
             requirements.extend(f.readlines())
-    requirements = [req for req in requirements if not req.startswith("#")]
+    requirements = [
+        req for req in requirements if not req.startswith("#") and not req.startswith("-")
+    ]
     return requirements
 
 

test_unstructured_inference/inference/test_layout.py

@@ -8,7 +8,7 @@
 from PIL import Image
 
 import unstructured_inference.models.base as models
-from unstructured_inference.inference import elements, layout
+from unstructured_inference.inference import elements, layout, layoutelement
 from unstructured_inference.models import detectron2, tesseract
 
 
@@ -18,14 +18,30 @@ def mock_image():
 
 
 @pytest.fixture()
-def mock_page_layout():
+def mock_initial_layout():
     text_block = layout.EmbeddedTextRegion(2, 4, 6, 8, text="A very repetitive narrative. " * 10)
 
     title_block = layout.EmbeddedTextRegion(1, 2, 3, 4, text="A Catchy Title")
 
     return [text_block, title_block]
 
 
+@pytest.fixture()
+def mock_final_layout():
+    text_block = layoutelement.LayoutElement(
+        2,
+        4,
+        6,
+        8,
+        text="A very repetitive narrative. " * 10,
+        type="NarrativeText",
+    )
+
+    title_block = layoutelement.LayoutElement(1, 2, 3, 4, text="A Catchy Title", type="Title")
+
+    return [text_block, title_block]
+
+
 def test_pdf_page_converts_images_to_array(mock_image):
     page = layout.PageLayout(number=0, image=mock_image, layout=[])
     assert page.image_array is None
@@ -62,13 +78,13 @@ def initialize(self, *args, **kwargs):
         pass
 
 
-def test_get_page_elements(monkeypatch, mock_page_layout):
+def test_get_page_elements(monkeypatch, mock_final_layout):
     image = np.random.randint(12, 24, (40, 40))
     page = layout.PageLayout(
         number=0,
         image=image,
-        layout=mock_page_layout,
-        model=MockLayoutModel(mock_page_layout),
+        layout=mock_final_layout,
+        model=MockLayoutModel(mock_final_layout),
     )
 
     elements = page.get_elements_with_model(inplace=False)
@@ -94,7 +110,17 @@ def join(self):
 def test_get_page_elements_with_ocr(monkeypatch):
     text_block = layout.TextRegion(2, 4, 6, 8, text=None)
     image_block = layout.ImageTextRegion(8, 14, 16, 18)
-    doc_layout = [text_block, image_block]
+    doc_initial_layout = [text_block, image_block]
+    text_layoutelement = layoutelement.LayoutElement(
+        2,
+        4,
+        6,
+        8,
+        text=None,
+        type="UncategorizedText",
+    )
+    image_layoutelement = layoutelement.LayoutElement(8, 14, 16, 18, text=None, type="Image")
+    doc_final_layout = [text_layoutelement, image_layoutelement]
 
     monkeypatch.setattr(detectron2, "is_detectron2_available", lambda *args: True)
     monkeypatch.setattr(elements, "ocr", lambda *args, **kwargs: "An Even Catchier Title")
@@ -103,24 +129,24 @@ def test_get_page_elements_with_ocr(monkeypatch):
     page = layout.PageLayout(
         number=0,
         image=image,
-        layout=doc_layout,
-        model=MockLayoutModel(doc_layout),
+        layout=doc_initial_layout,
+        model=MockLayoutModel(doc_final_layout),
     )
     page.get_elements_with_model()
 
     assert str(page) == "\n\nAn Even Catchier Title"
 
 
-def test_read_pdf(monkeypatch, mock_page_layout):
+def test_read_pdf(monkeypatch, mock_initial_layout, mock_final_layout):
     image = np.random.randint(12, 24, (40, 40))
     images = [image, image]
 
-    layouts = [mock_page_layout, mock_page_layout]
+    layouts = [mock_initial_layout, mock_initial_layout]
 
     monkeypatch.setattr(
         models,
         "UnstructuredDetectronModel",
-        partial(MockLayoutModel, layout=mock_page_layout),
+        partial(MockLayoutModel, layout=mock_final_layout),
     )
     monkeypatch.setattr(detectron2, "is_detectron2_available", lambda *args: True)
 
@@ -139,8 +165,8 @@ def test_read_pdf(monkeypatch, mock_page_layout):
 
 
 @pytest.mark.parametrize("model_name", [None, "checkbox", "fake"])
-def test_process_data_with_model(monkeypatch, mock_page_layout, model_name):
-    monkeypatch.setattr(layout, "get_model", lambda x: MockLayoutModel(mock_page_layout))
+def test_process_data_with_model(monkeypatch, mock_final_layout, model_name):
+    monkeypatch.setattr(layout, "get_model", lambda x: MockLayoutModel(mock_final_layout))
     monkeypatch.setattr(
         layout.DocumentLayout,
         "from_file",
@@ -158,11 +184,10 @@ def test_process_data_with_model_raises_on_invalid_model_name():
 
 
 @pytest.mark.parametrize("model_name", [None, "checkbox"])
-def test_process_file_with_model(monkeypatch, mock_page_layout, model_name):
+def test_process_file_with_model(monkeypatch, mock_final_layout, model_name):
     def mock_initialize(self, *args, **kwargs):
-        self.model = MockLayoutModel(mock_page_layout)
+        self.model = MockLayoutModel(mock_final_layout)
 
-    monkeypatch.setattr(models, "get_model", lambda x: MockLayoutModel(mock_page_layout))
     monkeypatch.setattr(
         layout.DocumentLayout,
         "from_file",
@@ -276,17 +301,17 @@ def test_get_elements_from_block_raises():
 
 
 @pytest.mark.parametrize("filetype", ["png", "jpg"])
-def test_from_image_file(monkeypatch, mock_page_layout, filetype):
+def test_from_image_file(monkeypatch, mock_final_layout, filetype):
     def mock_get_elements(self, *args, **kwargs):
-        self.elements = [mock_page_layout]
+        self.elements = [mock_final_layout]
 
     monkeypatch.setattr(layout.PageLayout, "get_elements_with_model", mock_get_elements)
     elements = (
         layout.DocumentLayout.from_image_file(f"sample-docs/loremipsum.{filetype}")
         .pages[0]
         .elements
     )
-    assert elements[0] == mock_page_layout
+    assert elements[0] == mock_final_layout
 
 
 def test_from_image_file_raises_with_empty_fn():
@@ -307,9 +332,9 @@ def test_from_file_raises_on_length_mismatch(monkeypatch):
 
 
 @pytest.mark.parametrize("idx", range(2))
-def test_get_elements_from_layout(mock_page_layout, idx):
-    page = MockPageLayout(layout=mock_page_layout)
-    block = mock_page_layout[idx].pad(3)
+def test_get_elements_from_layout(mock_initial_layout, idx):
+    page = MockPageLayout(layout=mock_initial_layout)
+    block = mock_initial_layout[idx].pad(3)
     fixed_layout = [block]
     elements = page.get_elements_from_layout(fixed_layout)
     assert elements[0].text == block.text
@@ -483,6 +508,7 @@ def test_load_pdf_image_placement():
     assert image_region.y2 < images[5].height / 2
 
 
+@pytest.mark.skip("Temporarily removed multicolumn to fix ordering")
 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)
@@ -522,6 +548,34 @@ def test_annotate():
         assert ((annotated_array[:, :, 2] == 1).mean()) > 0.992
 
 
+def test_textregion_returns_empty_ocr_never(mock_image):
+    tr = elements.TextRegion(0, 0, 24, 24)
+    assert tr.extract_text(objects=None, image=mock_image, ocr_strategy="never") == ""
+
+
+@pytest.mark.parametrize(("text", "expected"), [("asdf", "asdf"), (None, "")])
+def test_embedded_text_region(text, expected):
+    etr = elements.EmbeddedTextRegion(0, 0, 24, 24, text=text)
+    assert etr.extract_text(objects=None) == expected
+
+
+@pytest.mark.parametrize(
+    ("text", "ocr_strategy", "expected"),
+    [
+        (None, "never", ""),
+        (None, "always", "asdf"),
+        ("i have text", "never", "i have text"),
+        ("i have text", "always", "i have text"),
+    ],
+)
+def test_image_text_region(text, ocr_strategy, expected, mock_image):
+    itr = elements.ImageTextRegion(0, 0, 24, 24, text=text)
+    with patch.object(elements, "ocr", return_value="asdf"):
+        assert (
+            itr.extract_text(objects=None, image=mock_image, ocr_strategy=ocr_strategy) == expected
+        )
+
+
 @pytest.fixture()
 def ordering_layout():
     elements = [
@@ -537,7 +591,11 @@ def ordering_layout():
 
 
 def test_layout_order(ordering_layout):
-    with patch.object(layout, "get_model", lambda: lambda x: ordering_layout):
+    with patch.object(layout, "get_model", lambda: lambda x: ordering_layout), patch.object(
+        layout,
+        "load_pdf",
+        lambda *args, **kwargs: ([[]], [mock_image]),
+    ):
         doc = layout.DocumentLayout.from_file("sample-docs/layout-parser-paper.pdf")
         page = doc.pages[0]
     for n, element in enumerate(page.elements):

test_unstructured_inference/models/test_yolox.py

@@ -68,4 +68,5 @@ def test_layout_yolox_local_parsing_empty_pdf_soft():
     document_layout = process_file_with_model(filename, model_name="yolox_tiny")
     assert len(document_layout.pages) == 1
     # NOTE(benjamin) The example sent to the test contains 0 detections
-    assert len(document_layout.pages[0].elements) == 0
+    text_elements_page_1 = [el for el in document_layout.pages[0].elements if el.type != "Image"]
+    assert len(text_elements_page_1) == 0