feat: partition image (#144)

Adds partition_image to partition image file types, which is integrated into the partition brick. This relies on the 0.2.2 version of unstructured-inference.

Author: qued

.github/workflows/ci.yml

@@ -98,7 +98,7 @@ jobs:
         source .venv/bin/activate
         make install-nltk-models
         make install-detectron2
-        sudo apt-get install -y libmagic-dev poppler-utils
+        sudo apt-get install -y libmagic-dev poppler-utils tesseract-ocr
         make test
         make check-coverage
 

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.4.2-dev0
+* Added `partition_image` to process documents in an image format.
+
+
 ## 0.4.1
 
 * Added support for text files in the `partition` function

requirements/local-inference.in

@@ -1 +1 @@
-unstructured-inference>=0.2.1
+unstructured-inference>=0.2.2

requirements/local-inference.txt

@@ -150,7 +150,7 @@ typing-extensions==4.4.0
     #   starlette
     #   torch
     #   torchvision
-unstructured-inference==0.2.1
+unstructured-inference==0.2.2
     # via -r requirements/local-inference.in
 urllib3==1.26.13
     # via requests

test_unstructured/partition/test_auto.py

@@ -151,6 +151,19 @@ def test_auto_partition_pdf_from_file():
     assert len(elements) > 0
 
 
+def test_auto_partition_jpg():
+    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "..", "..", "example-docs", "example.jpg")
+    elements = partition(filename=filename)
+    assert len(elements) > 0
+
+
+def test_auto_partition_jpg_from_file():
+    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "..", "..", "example-docs", "example.jpg")
+    with open(filename, "rb") as f:
+        elements = partition(file=f)
+    assert len(elements) > 0
+
+
 def test_auto_partition_raises_with_bad_type(monkeypatch):
     monkeypatch.setattr(auto, "detect_filetype", lambda *args, **kwargs: None)
     with pytest.raises(ValueError):

test_unstructured/partition/test_image.py

@@ -0,0 +1,126 @@
+import pytest
+import requests
+from unittest import mock
+
+import unstructured.partition.pdf as pdf
+import unstructured.partition.image as image
+import unstructured_inference.inference.layout as layout
+
+
+class MockResponse:
+    def __init__(self, status_code, response):
+        self.status_code = status_code
+        self.response = response
+
+    def json(self):
+        return self.response
+
+
+def mock_healthy_get(url, **kwargs):
+    return MockResponse(status_code=200, response={})
+
+
+def mock_unhealthy_get(url, **kwargs):
+    return MockResponse(status_code=500, response={})
+
+
+def mock_unsuccessful_post(url, **kwargs):
+    return MockResponse(status_code=500, response={})
+
+
+def mock_successful_post(url, **kwargs):
+    response = {
+        "pages": [
+            {
+                "number": 0,
+                "elements": [{"type": "Title", "text": "Charlie Brown and the Great Pumpkin"}],
+            }
+        ]
+    }
+    return MockResponse(status_code=200, response=response)
+
+
+class MockPageLayout(layout.PageLayout):
+    def __init__(self, number: int):
+        pass
+
+    @property
+    def elements(self):
+        return [
+            layout.LayoutElement(
+                type="Title",
+                coordinates=[(0, 0), (2, 2)],
+                text="Charlie Brown and the Great Pumpkin",
+            )
+        ]
+
+
+class MockDocumentLayout(layout.DocumentLayout):
+    @property
+    def pages(self):
+        return [
+            MockPageLayout(
+                number=0,
+            )
+        ]
+
+
+def test_partition_image_api(monkeypatch, filename="example-docs/example.jpg"):
+    monkeypatch.setattr(requests, "post", mock_successful_post)
+    monkeypatch.setattr(requests, "get", mock_healthy_get)
+
+    partition_image_response = pdf._partition_via_api(filename)
+    assert partition_image_response[0]["type"] == "Title"
+    assert partition_image_response[0]["text"] == "Charlie Brown and the Great Pumpkin"
+
+
+@pytest.mark.parametrize("filename, file", [("example-docs/example.jpg", None), (None, b"0000")])
+def test_partition_image_local(monkeypatch, filename, file):
+    monkeypatch.setattr(
+        layout, "process_data_with_model", lambda *args, **kwargs: MockDocumentLayout()
+    )
+    monkeypatch.setattr(
+        layout, "process_file_with_model", lambda *args, **kwargs: MockDocumentLayout()
+    )
+
+    partition_image_response = pdf._partition_pdf_or_image_local(filename, file, is_image=True)
+    assert partition_image_response[0].type == "Title"
+    assert partition_image_response[0].text == "Charlie Brown and the Great Pumpkin"
+
+
+@pytest.mark.skip("Needs to be fixed upstream in unstructured-inference")
+def test_partition_image_local_raises_with_no_filename():
+    with pytest.raises(FileNotFoundError):
+        pdf._partition_pdf_or_image_local(filename="", file=None, is_image=True)
+
+
+def test_partition_image_api_raises_with_failed_healthcheck(
+    monkeypatch, filename="example-docs/example.jpg"
+):
+    monkeypatch.setattr(requests, "post", mock_successful_post)
+    monkeypatch.setattr(requests, "get", mock_unhealthy_get)
+
+    with pytest.raises(ValueError):
+        pdf._partition_via_api(filename=filename, url="http://ml.unstructured.io/layout/image")
+
+
+def test_partition_image_api_raises_with_failed_api_call(
+    monkeypatch, filename="example-docs/example.jpg"
+):
+    monkeypatch.setattr(requests, "post", mock_unsuccessful_post)
+    monkeypatch.setattr(requests, "get", mock_healthy_get)
+
+    with pytest.raises(ValueError):
+        pdf._partition_via_api(filename=filename, url="http://ml.unstructured.io/layout/image")
+
+
+@pytest.mark.parametrize(
+    "url, api_called, local_called", [("fakeurl", True, False), (None, False, True)]
+)
+def test_partition_image(url, api_called, local_called):
+    with mock.patch.object(
+        pdf, attribute="_partition_via_api", new=mock.MagicMock()
+    ), mock.patch.object(pdf, "_partition_pdf_or_image_local", mock.MagicMock()):
+        image.partition_image(filename="fake.pdf", url=url)
+        assert pdf._partition_via_api.called == api_called
+        assert pdf._partition_pdf_or_image_local.called == local_called

test_unstructured/partition/test_pdf.py

@@ -68,7 +68,7 @@ def test_partition_pdf_api(monkeypatch, filename="example-docs/layout-parser-pap
     monkeypatch.setattr(requests, "post", mock_successful_post)
     monkeypatch.setattr(requests, "get", mock_healthy_get)
 
-    partition_pdf_response = pdf._partition_pdf_via_api(filename)
+    partition_pdf_response = pdf._partition_via_api(filename)
     assert partition_pdf_response[0]["type"] == "Title"
     assert partition_pdf_response[0]["text"] == "Charlie Brown and the Great Pumpkin"
 
@@ -77,12 +77,14 @@ def test_partition_pdf_api(monkeypatch, filename="example-docs/layout-parser-pap
     "filename, file", [("example-docs/layout-parser-paper-fast.pdf", None), (None, b"0000")]
 )
 def test_partition_pdf_local(monkeypatch, filename, file):
-    monkeypatch.setattr(layout, "process_data_with_model", lambda *args: MockDocumentLayout())
+    monkeypatch.setattr(
+        layout, "process_data_with_model", lambda *args, **kwargs: MockDocumentLayout()
+    )
     monkeypatch.setattr(
         layout, "process_file_with_model", lambda *args, **kwargs: MockDocumentLayout()
     )
 
-    partition_pdf_response = pdf._partition_pdf_via_local(filename, file)
+    partition_pdf_response = pdf._partition_pdf_or_image_local(filename, file)
     assert partition_pdf_response[0].type == "Title"
     assert partition_pdf_response[0].text == "Charlie Brown and the Great Pumpkin"
 
@@ -92,15 +94,12 @@ def test_partition_pdf_api_raises_with_no_filename(monkeypatch):
     monkeypatch.setattr(requests, "get", mock_healthy_get)
 
     with pytest.raises(FileNotFoundError):
-        pdf._partition_pdf_via_api(filename=None, file=None)
-
+        pdf._partition_via_api(filename=None, file=None)
 
-def test_partition_pdf_local_raises_with_no_filename(monkeypatch):
-    monkeypatch.setattr(requests, "post", mock_successful_post)
-    monkeypatch.setattr(requests, "get", mock_healthy_get)
 
+def test_partition_pdf_local_raises_with_no_filename():
     with pytest.raises(FileNotFoundError):
-        pdf._partition_pdf_via_api(filename=None, file=None)
+        pdf._partition_pdf_or_image_local(filename="", file=None, is_image=False)
 
 
 def test_partition_pdf_api_raises_with_failed_healthcheck(
@@ -110,7 +109,7 @@ def test_partition_pdf_api_raises_with_failed_healthcheck(
     monkeypatch.setattr(requests, "get", mock_unhealthy_get)
 
     with pytest.raises(ValueError):
-        pdf._partition_pdf_via_api(filename=filename)
+        pdf._partition_via_api(filename=filename)
 
 
 def test_partition_pdf_api_raises_with_failed_api_call(
@@ -120,16 +119,16 @@ def test_partition_pdf_api_raises_with_failed_api_call(
     monkeypatch.setattr(requests, "get", mock_healthy_get)
 
     with pytest.raises(ValueError):
-        pdf._partition_pdf_via_api(filename=filename)
+        pdf._partition_via_api(filename=filename)
 
 
 @pytest.mark.parametrize(
     "url, api_called, local_called", [("fakeurl", True, False), (None, False, True)]
 )
 def test_partition_pdf(url, api_called, local_called):
-    with mock.patch(
-        "unstructured.partition.pdf._partition_pdf_via_api", mock.MagicMock()
-    ), mock.patch("unstructured.partition.pdf._partition_pdf_via_local", mock.MagicMock()):
+    with mock.patch.object(
+        pdf, attribute="_partition_via_api", new=mock.MagicMock()
+    ), mock.patch.object(pdf, "_partition_pdf_or_image_local", mock.MagicMock()):
         pdf.partition_pdf(filename="fake.pdf", url=url)
-        assert pdf._partition_pdf_via_api.called == api_called
-        assert pdf._partition_pdf_via_local.called == local_called
+        assert pdf._partition_via_api.called == api_called
+        assert pdf._partition_pdf_or_image_local.called == local_called

unstructured/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.1"  # pragma: no cover
+__version__ = "0.4.2-dev0"  # pragma: no cover

unstructured/partition/__init__.py

@@ -0,0 +1,42 @@
+import requests  # type: ignore
+from typing import BinaryIO, List, Optional, Union, Tuple, Mapping
+
+from unstructured.documents.elements import Element
+
+
+def _partition_via_api(
+    filename: str = "",
+    file: Optional[bytes] = None,
+    url: str = "https://ml.unstructured.io/layout/pdf",
+    token: Optional[str] = None,
+    data: Optional[dict] = None,  # NOTE(alan): Remove after different models are handled by routing
+) -> List[Element]:
+    """Use API for partitioning."""
+    if not filename and not file:
+        raise FileNotFoundError("No filename nor file were specified")
+
+    healthcheck_response = requests.models.Response()
+    if not token:
+        healthcheck_response = requests.get(url=f"{url}healthcheck")
+
+    if healthcheck_response.status_code != 200:
+        raise ValueError("endpoint api healthcheck has failed!")
+
+    file_: Mapping[str, Tuple[str, Union[BinaryIO, bytes]]] = {
+        "file": (
+            filename,
+            file if file else open(filename, "rb"),
+        )
+    }
+    response = requests.post(
+        url=url,
+        headers={"Authorization": f"Bearer {token}" if token else ""},
+        files=file_,
+        data=data,  # NOTE(alan): Remove after unstructured API is using routing
+    )
+
+    if response.status_code == 200:
+        pages = response.json()["pages"]
+        return [element for page in pages for element in page["elements"]]
+    else:
+        raise ValueError(f"response status code = {response.status_code}")

unstructured/partition/auto.py

@@ -5,6 +5,7 @@
 from unstructured.partition.email import partition_email
 from unstructured.partition.html import partition_html
 from unstructured.partition.pdf import partition_pdf
+from unstructured.partition.image import partition_image
 from unstructured.partition.text import partition_text
 
 
@@ -34,6 +35,8 @@ def partition(filename: Optional[str] = None, file: Optional[IO] = None):
         return partition_html(filename=filename, file=file)
     elif filetype == FileType.PDF:
         return partition_pdf(filename=filename, file=file, url=None)  # type: ignore
+    elif (filetype == FileType.PNG) or (filetype == FileType.JPG):
+        return partition_image(filename=filename, file=file, url=None)  # type: ignore
     elif filetype == FileType.TXT:
         return partition_text(filename=filename, file=file)
     else: