fix: store page no (#118)

Fixes page numbers not being correctly stored in the PageLayout objects.

Author: qued

CHANGELOG.md

@@ -1,5 +1,6 @@
-## 0.5.1-dev0
+## 0.5.1-dev1
 
+* Store page numbers when processing PDFs
 * Hotfix to handle inference of blank pages using ONNX detectron2
 
 ## 0.5.0

test_unstructured_inference/inference/test_layout.py

@@ -308,6 +308,15 @@ def test_get_elements_from_layout(mock_page_layout, idx):
     assert elements[0].text == block.text
 
 
+def test_page_numbers_in_page_objects():
+    with patch(
+        "unstructured_inference.inference.layout.PageLayout.get_elements_with_model"
+    ) as mock_get_elements:
+        doc = layout.DocumentLayout.from_file("sample-docs/layout-parser-paper.pdf")
+        mock_get_elements.assert_called()
+        assert [page.number for page in doc.pages] == list(range(1, len(doc.pages) + 1))
+
+
 @pytest.mark.parametrize(
     "fixed_layouts, called_method, not_called_method",
     [

unstructured_inference/__version__.py

@@ -1 +1 @@
-__version__ = "0.5.1-dev0"  # pragma: no cover
+__version__ = "0.5.1-dev1"  # pragma: no cover

unstructured_inference/inference/layout.py

@@ -76,11 +76,12 @@ def from_file(
         pages: List[PageLayout] = list()
         if fixed_layouts is None:
             fixed_layouts = [None for _ in layouts]
-        for image, layout, fixed_layout in zip(images, layouts, fixed_layouts):
+        for i, (image, layout, fixed_layout) in enumerate(zip(images, layouts, fixed_layouts)):
             # NOTE(robinson) - In the future, maybe we detect the page number and default
             # to the index if it is not detected
             page = PageLayout.from_image(
                 image,
+                number=i + 1,
                 model=model,
                 layout=layout,
                 ocr_strategy=ocr_strategy,
@@ -195,6 +196,7 @@ def _get_image_array(self) -> Union[np.ndarray, None]:
     def from_image(
         cls,
         image,
+        number=1,
         model: Optional[UnstructuredModel] = None,
         layout: Optional[List[TextRegion]] = None,
         ocr_strategy: str = "auto",
@@ -204,7 +206,7 @@ def from_image(
     ):
         """Creates a PageLayout from an already-loaded PIL Image."""
         page = cls(
-            number=0,
+            number=number,
             image=image,
             layout=layout,
             model=model,